Repository navigation
feat(dynacell): modernize pix2pix3d training recipe (StyleGAN2/R3GAN-style) - #428
Merged
Merged
Conversation
Introduces a new gan subpackage with PatchGAN3D and MultiScalePatchGAN3D (pix2pixHD-style, 2-scale by default), plus lsgan_d_loss / lsgan_g_loss list-based helpers. These will pair with UNetViT3D to form a 3D pix2pix adversarial baseline alongside the existing regression and flow-matching DynaCell baselines. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
DynacellGAN composes the existing UNetViT3D generator with the new MultiScalePatchGAN3D discriminator, training with LSGAN + L1 (λ=100) via manual optimization (automatic_optimization=False). A requires_grad toggle separates the D step from the G step so each backward only populates the active side's grads; D grads are also explicitly cleared between optimizers so the test invariant (no D grads after a G step) holds for Phase 2 DDP correctness. Validation logs the existing loss/validate weighted-mean alias, so ModelCheckpoint(monitor="loss/validate") keys off the same metric used across the other DynaCell baselines. configure_optimizers builds two AdamW + WarmupCosine pairs directly to fit the two-optimizer return shape Lightning's manual mode expects. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…ke leaf The recipe declares dynacell.engine.DynacellGAN with UNetViT3D's production generator config and a 2-scale spectral-normalized PatchGAN3D. fit / predict model overlays mirror the unetvit3d ones with lr_g/lr_d split (no scalar lr) and bf16-mixed precision. The single-GPU smoke leaf for er/pix2pix3d_unetvit/ipsc_confocal sets benchmark.dataset_ref: null to disable the compose-hook resolver so it can point data_path at SEC61B_test48.zarr directly, and uses batch_size=2 to match the inherited num_samples=2 (HCSDataModule divisibility rule). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Mirrors er/unetvit3d/ipsc_confocal/train.yml with the model overlay and all model_name / experiment_id / run_root / job_name strings retargeted to pix2pix3d_unetvit. Keeps ModelCheckpoint.monitor: loss/validate unchanged — DynacellGAN logs that key via on_validation_epoch_end. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…organelles
Adds 7 production single-GPU training leaves for pix2pix3d_unetvit:
- er/pix2pix3d_unetvit/a549_mantis
- {membrane, mito, nucleus}/pix2pix3d_unetvit/{ipsc_confocal, a549_mantis}
Each leaf mirrors the matching unetvit3d cell with model_overlay,
model_name, experiment_id, logger name/save_dir, ckpt dirpath, launcher
run_root/job_name retargeted to pix2pix3d_unetvit. ModelCheckpoint
monitor=loss/validate preserved.
Joint training leaves (joint_ipsc_confocal_a549_mantis) deferred — still
gated on the open BatchedConcatDataModule DDP deadlock (handoff
2026-04-26).
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Mirrors er/unetvit3d/ipsc_confocal/{predict,eval}__ipsc_confocal.* with
the model overlay, model_name, experiment_id, ckpt_path, output_store
filename, and eval pred_path / save_dir all retargeted to
pix2pix3d_unetvit. predict_method stays full_image (the only mode
DynacellGAN.predict_step supports for v1).
ckpt_path uses an explicit REPLACE_ME_WITH_PRODUCTION_CHECKPOINT_PATH
placeholder so torch.load raises FileNotFoundError with the placeholder
visible if the user forgets to override it — running predict against an
untrained generator would silently produce garbage.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…3d_unetvit/ipsc Adds the 3 a549_mantis test splits (mock, denv, zikv) for the iPSC-trained pix2pix3d_unetvit ER cell — enables cross-dataset transfer evaluation. Mirrors the unetvit3d ipsc_confocal predict/eval triplet with the same retargeting rule already used for the iPSC pair. ckpt_path uses the REPLACE_ME_WITH_PRODUCTION_CHECKPOINT_PATH placeholder so any forgotten override fails loudly via FileNotFoundError. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
… dedupe _log_samples DynacellGAN.configure_optimizers used to rebuild AdamW + WarmupCosineSchedule inline despite the class docstring claiming it called the shared helper. Now calls configure_adamw_scheduler twice and concatenates, dropping the redundant schedule guard and the unused monai/torch.optim imports. _log_samples was a byte-identical instance method on DynacellUNet, DynacellFlowMatching, and DynacellGAN. Promoted to a module-level free function and updated the five call sites across all three classes. Updated test_dynacell_gan_validate_logs_alias to monkeypatch the module attribute instead of the instance method. Also trimmed three narrating comments inside DynacellGAN (the plan/PR context comments belong in the PR description, not the source) and five per-layer banner comments in patchgan3d.py (the conv literals state the shape and the class docstring covers the layer table). Safe under the running production training's Slurm auto-requeue: no constructor signatures change, AdamW + WarmupCosine pairs are built with the same args, optimizer state shape is unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Adds the DDP infrastructure for pix2pix3d_unetvit so the GAN baseline can train on 4 H200s in ~13 h instead of the ~50 h single-GPU pace measured on job 33132605 (now cancelled). - recipes/topology/ddp_4gpu_gan.yml — DDP variant that wraps the strategy in find_unused_parameters: true. Required because DynacellGAN runs manual_optimization with two optimizers, and the per-step requires_grad toggle leaves one side's params "unused" from the reducer's perspective on each backward. Plain ddp_4gpu.yml stays single-strategy for the other models. - model_overlays/pix2pix3d_unetvit_fit_ddp.yml — mirrors the single-GPU fit overlay but composes ddp_4gpu_gan.yml topology + 4-GPU hardware. - er/pix2pix3d_unetvit/ipsc_confocal/train_smoke_4gpu.yml — 4-GPU smoke variant. Same SEC61B_test48 store, batch_size=2 matching the inherited num_samples=2 divisibility rule. Validates DDP sharding + the find_unused_parameters DDP reducer behavior under the GAN's manual-opt pattern before committing the 4-GPU production run. - er/pix2pix3d_unetvit/ipsc_confocal/train_4gpu.yml — 4-GPU production leaf. Keeps per-rank batch_size=4 (effective 16 across 4 GPUs) and unchanged lr_g/lr_d so the DDP run is comparable to the single-GPU reference at the per-rank optimization settings. logger name and job_name suffixed with _4gpu to disambiguate in W&B/slurm; run_root and ckpt dirpath unchanged so Phase 4 predict leaves keep working. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The membrane single-set train leaves landed with placeholder paths under .../ipsc/memb_temp/ — diverging from the celldiff convention (memb, not memb_temp) and conflating the a549_mantis run dir with the iPSC tree. Both leaves now point at the canonical .../ipsc/memb/ and .../a549_mantis/memb/ trees, matching the celldiff layout so all organelles can be cross-referenced uniformly. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Mirrors the er/pix2pix3d_unetvit/ipsc_confocal/ predict+eval set across membrane, mito, and nucleus so every iPSC-trained pix2pix3d_unetvit run has an apples-to-apples inference + evaluation path against both the ipsc_confocal self-target and the three a549_mantis conditions (mock, denv, zikv). Predict leaves swap target-inherited normalizations for a source-only NormalizeSampled (Phase3D) and clear the RandWeightedCropd augmentation. Eval leaves enable compute_feature_metrics only on the ipsc_confocal self-target, matching the existing er convention. Ckpt paths are placeholders pending the corresponding production runs. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Adds the joint_ipsc_confocal_a549_mantis variant for all four organelles (er, membrane, mito, nucleus): one train.yml authoring BatchedConcatDataModule with two HCSDataModule children plus the four predict leaves (self iPSC + three a549_mantis conditions). The joint train differs from the celldiff joint by using NormalizeSampled (fov_statistics) instead of MinMaxSampled so joint and single-set pix2pix3d_unetvit runs share the same normalization contract for apples-to-apples ablations. batch_size=2 × num_samples=2 matches the single-set 4 GPU-samples/step (per the BatchedConcatDataModule rule in CLAUDE.md that joint does NOT divide by num_samples). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
PatchGAN3D and MultiScalePatchGAN3D accepted a num_layers constructor parameter that only worked with the value 5 — any other value raised ValueError. The argument was justified as future-proofing for ablations that don't exist yet, which is exactly the speculative flexibility CLAUDE.md prohibits: "No 'flexibility' or 'configurability' that wasn't requested. ... No error handling for impossible scenarios." Re-add it only if and when a non-5 variant lands. Recipe + test fixture updated in lockstep; uv run pytest packages/viscy-models/tests/test_gan/ and the DynacellGAN init/forward tests pass (11/11). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
The 1:1 LR schedule collapsed to D-dominance on DDP run 33176091
(wandb 20260520-002629): loss/d_train fell from 1.3e-2 (ep1) to
~1e-4 (ep12+), loss/g_adv pinned at ~1.0 (saturated MSE(d_fake, 1)),
loss/g_l1 flat at ~0.45 from epoch ~5 onward — the adversarial
gradient was effectively zero, generator was training as pure L1.
Slowing the discriminator 10x (lr_d = lr_g / 10 = 3e-5) is the
canonical pix2pix-style fix for D-dominance. Applies to both the
single-GPU and DDP fit overlays so the entire Phase 3 baseline
(all 4 organelles × {iPSC, A549, joint}) inherits the same
adversarial balance regime.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…into gan-baseline-pix2pix3d
Adds three loss families and two gradient penalties to viscy_models.gan, extending the existing multi-scale 3D PatchGAN convention to the modern StyleGAN2 / R3GAN recipe: - nonsat_d_loss / nonsat_g_loss: softplus non-saturating loss, verbatim from stylegan2-ada-pytorch/training/loss.py. - rpgan_d_loss / rpgan_g_loss: relativistic pairing loss, verbatim from brownvc/R3GAN/Trainer.py. - r1_penalty / r2_penalty: Mescheder zero-centered gradient penalty on real / fake input. Multi-scale aggregation = mean of per-scale ||grad||^2 (NOT ||grad of sum||^2, which would mix cross-scale derivatives through MultiScalePatchGAN3D's pool-derived second scale). All losses preserve the existing list-of-per-scale-Tensors convention. LSGAN losses are unchanged. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
… recipe)
Adds opt-in modernization to DynacellGAN with legacy-safe defaults
(loss_type=lsgan, r1_gamma=0, ema_kimg=None, lecam_gamma=0). The 14
existing leaves composing pix2pix3d_unetvit_fit.yml are unaffected
because none of them set the new knobs and the new defaults preserve
pre-modernization behavior.
New constructor knobs:
- loss_type {lsgan, nonsat, rpgan} + lambda_adv
- r1_gamma / r2_gamma / r1_every (lazy schedule with StyleGAN2-style
* r1_every unbiased rescaling). RpGAN raises ValueError unless
r1_gamma > 0 (R3GAN Theorem 3.1 convergence requirement).
- ema_kimg + use_ema_at_predict for the generator EMA shadow. Decay
derived per-step as 0.5 ** (global_batch / (ema_kimg * 1000)).
- lecam_gamma + lecam_decay (decay=0.9 matches google/lecam-gan reference).
Training-step structural changes:
- D-step counter (_d_step_count buffer) drives the lazy R1 schedule;
self.global_step would advance per opt.step() (2 per training_step
with manual_optimization) so it can't be used directly.
- R1 / R2 path is outer-guarded on (r1_gamma > 0 or r2_gamma > 0) so
legacy leaves skip the grad-of-grad compute and the spectral-norm
power-iteration advance.
- R1 / R2 compute is wrapped in autocast(enabled=False); Lightning's
bf16-mixed plugin would otherwise put bf16 into D forwards which
is fragile under create_graph=True grad-of-grad.
- RpGAN G step recomputes fresh d_real against the post-update
discriminator (matches R3GAN Trainer.py convention).
- LeCam EMA update uses sync_dist'd batch means via self.all_gather so
all ranks update the buffers identically without explicit all_reduce.
- Generator EMA update after opt_g.step() (params are DDP-synced at
that point, so per-rank EMA stays in lockstep).
Inference paths:
- _inference_generator helper returns EMA when available AND
use_ema_at_predict; used by forward and predict_step.
- validation_step runs both raw and EMA generator forwards (when EMA
exists); emits both loss/validate (back-compat) and loss/validate_ema
aliases.
Checkpoint loading at engine.py:802 changed to strict=False with
filtered warning logs (expected_missing for generator_ema.* /
_d_step_count / _lecam_ema_*). When loading a pre-modernization
checkpoint with EMA enabled, the EMA shadow is seeded from the loaded
generator weights so inference paths return the loaded weights instead
of the random-init deepcopy.
16 DynacellGAN tests cover legacy-default safety, modernized smoke, the
R1/R2 gamma guard, lazy-R1 step counter, RpGAN constructor validation,
EMA seeding on legacy load, LeCam buffer registration gating, two-pass
validation logging, and inference-helper EMA selection.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
New opt-in configs for the StyleGAN2/R3GAN-style modernized pix2pix3d recipe: - pix2pix3d_unetvit_fit_modernized.yml: single-GPU overlay with loss_type=nonsat, r1_gamma=10, r1_every=16, ema_kimg=10, lr_g=lr_d=2e-4. Existing pix2pix3d_unetvit_fit.yml stays untouched so the 14 legacy leaves continue running LSGAN unchanged. - pix2pix3d_unetvit_fit_ddp_modernized.yml: 4-GPU DDP sibling, same model knobs, composes ddp_4gpu_gan topology (find_unused_parameters=True). - er/pix2pix3d_unetvit/ipsc_confocal/train_4gpu_modernized.yml: Run A of the plan. SEC61B headline experiment. ModelCheckpoint monitors loss/validate_ema (the EMA generator's val loss = inference path). Checkpoints land in a distinct subtree from the LSGAN baseline so there is no collision risk with in-flight runs. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…view) Addresses findings from the PR review subagents: - ema_kimg=0 raises ValueError (silent EMA freeze prevention). - RpGAN now requires both r1_gamma > 0 AND r2_gamma > 0 (R3GAN Theorem 3.1 requires both penalties for sharp-distribution convergence; previously only r1_gamma was validated). - ckpt_path strict=False load: missing-but-not-expected keys now RAISE RuntimeError instead of warnings.warn (silent half-loaded models burn hours of training before users notice). Unexpected-keys path switched from warnings.warn to _logger.warning. - Partial-EMA checkpoints (some but not all generator_ema.* keys) now RAISE — silently half-seeding the EMA shadow would produce nonsense outputs in the missing layers. - on_predict_start logs which generator (EMA vs raw) will be used at inference, so users notice silent fallbacks (predict overlay forgot ema_kimg, use_ema_at_predict=False, etc.). New tests: - test_dynacell_gan_rpgan_requires_r2 - test_dynacell_gan_ema_kimg_zero_raises - test_dynacell_gan_rpgan_smoke (end-to-end trainer.fit with rpgan + R1+R2; exercises the fresh-d_real RpGAN G-step path) - test_dynacell_gan_ema_update_math (closed-form EMA decay verification) - test_dynacell_gan_use_ema_at_predict_false (predict_step returns raw output when EMA is gated off) - test_dynacell_gan_ckpt_unexpected_missing_raises (corrupted ckpt surfaces as RuntimeError, not a swallowed warning) Existing test_dynacell_gan_ema_seed_on_legacy_load updated to use use_spectral_norm=False on both sides — with the stricter unexpected-missing check, SN parametrization name mismatches would otherwise trip the raise. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Contributor
There was a problem hiding this comment.
Pull request overview
Opt-in modernization of the pix2pix3d adversarial training recipe in Dynacell, adding StyleGAN2/R3GAN-style loss families (non-saturating softplus and relativistic pairing) plus optional R1/R2 regularization and generator EMA, while keeping existing LSGAN defaults legacy-safe.
Changes:
- Added modern GAN loss functions and multi-scale R1/R2 gradient penalties to
viscy_models.gan.losses, and exported them in the GAN package API. - Extended
DynacellGANto supportloss_typedispatch, lazy R1/R2 scheduling, generator EMA, LeCam regularization, and dual raw/EMA validation metrics. - Added unit + integration tests and introduced modernized config overlays/leaves to enable the new recipe without affecting existing runs.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 5 comments.
Show a summary per file
| File | Description |
|---|---|
| packages/viscy-models/src/viscy_models/gan/losses.py | Adds nonsat/RpGAN losses and multi-scale R1/R2 penalties with scale validation. |
| packages/viscy-models/src/viscy_models/gan/init.py | Re-exports new loss/penalty functions via viscy_models.gan. |
| packages/viscy-models/tests/test_gan/test_modern_losses.py | Unit tests for new loss families and gradient-penalty behavior (including AMP wrapper expectations). |
| applications/dynacell/src/dynacell/engine.py | Wires new loss/regularization/EMA knobs into DynacellGAN, adds EMA-aware validation and inference selection. |
| applications/dynacell/tests/test_engine.py | Adds integration tests covering legacy-default safety, modernization smoke, lazy-R1 counter, EMA seeding, and validation aliases. |
| applications/dynacell/configs/benchmarks/virtual_staining/_internal/shared/model/model_overlays/pix2pix3d_unetvit_fit_modernized.yml | New single-GPU model overlay enabling the modernized recipe. |
| applications/dynacell/configs/benchmarks/virtual_staining/_internal/shared/model/model_overlays/pix2pix3d_unetvit_fit_ddp_modernized.yml | DDP variant overlay enabling the same modernized knobs for GAN-aware DDP topology. |
| applications/dynacell/configs/benchmarks/virtual_staining/er/pix2pix3d_unetvit/ipsc_confocal/train_4gpu_modernized.yml | New 4-GPU training leaf composing the modernized overlay and monitoring EMA validation loss. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
r1_every is the modulo period for the lazy R1/R2 schedule and is also used as the unbiased rescaling factor on the penalty term. Without an init-time guard, r1_every=0 would ZeroDivisionError on the first D step when R1/R2 is enabled, and a negative value would silently fire the gate on the wrong cadence. Raise at construction so the bad config fails before any training time is burned. Surfaced by Copilot review on PR #428. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
DynacellGAN.validation_step assigned pred_for_samples=pred_ema whenever the EMA submodule existed, regardless of use_ema_at_predict. This contradicted predict_step / forward (both routed through _inference_generator, which respects the flag) — the dashboard could show EMA-generator samples while real predictions were emitted by the raw generator. Now matches the inference path. Regression test stubs out both generators with distinguishable constant modules and reads back validation_step_outputs to confirm which one fed the dashboard for each value of use_ema_at_predict. Surfaced by Copilot review on PR #428. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Lazy-reg gate previously read `int(self._d_step_count)` each training step where the counter was a CUDA buffer. Each .item() call forced a device→host sync. Switch to a plain Python int (incremented in training_step, modulo-gated against r1_every) and persist it via on_save_checkpoint / on_load_checkpoint hooks. All DDP ranks start at 0 and advance deterministically, so the buffer-was-broadcast invariant is preserved without per-step sync overhead. Side benefit: `_d_step_count` is no longer a state_dict key, so the strict-load filter drops one expected-missing prefix. Legacy checkpoints without `_d_step_count` resume cleanly at 0 via the dict.get fallback. Surfaced by Copilot review on PR #428. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
… path The header comment referenced /home/alex.kalinin/.claude/plans/..., which won't exist for other contributors or in CI. Swap for the PR URL so the recipe rationale is reachable from the file. Surfaced by Copilot review on PR #428. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
…path Same fix as the modernized overlay header — replace /home/alex.kalinin/.claude/plans/... reference with the PR URL so it resolves for everyone. Surfaced by Copilot review on PR #428. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
edyoshikun
added a commit
that referenced
this pull request
May 22, 2026
Co-authored-by: Claude Opus 4.7 (1M context) <noreply@anthropic.com> Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com> Co-authored-by: Eduardo Hirata-Miyasaki <edhiratam@gmail.com> Co-authored-by: dihan.zheng <dihan.zheng@ucsf.edu> fix(viscy-data): drop use_thread_workers to fix DDP deadlock (#413) fix(viscy-data): port to iohub 0.3.2 ImageArray API (#407) fix shipped via PR #412 fixes triggered by Copilot's review on PR #427 plus an independent fix two CLAUDE.md doc nits flagged in PR #427 self-review fix(dynacell): address PR #426 code-review findings fix(dynacell): harden DynacellGAN against silent failures (PR #428 review)
alxndrkalinin
added a commit
that referenced
this pull request
May 22, 2026
…g/dynacell-dynadtw PR #432 base modular-viscy-staging diverged from staging/dynacell-dynadtw when #404 squashed all of dynacell-models history into a single commit (393abdf). dynacell-models has 30 newer commits since dynadtw's last sync at f1df36b — re-eval campaign generator, controlled FOV-parallel eval (#426), 3D pix2pix GAN baseline (#425), DeepFeatureBatcher / precompute deep features, cache identity refactors, pix2pix3d_unetvit modernization (#428), backbone bookkeeping refactor onto _BackboneLists. Conflict triage (92 conflicts): - 89 files where staging/dynacell-dynadtw carried stale copies of dynacell-models content (no post-branch edits on dynadtw side): 84 AA conflicts (eval YAML configs, dynacell evaluation sources/tests, CLAUDE.md, tools) + 5 UU conflicts (engine.py, __init__.py, __main__.py, test_engine.py, .github/workflows/test.yml). Resolved by taking the modular-viscy-staging post-#404 version. - 3 files with genuine concurrent edits, resolved per direction: * .gitignore — dynadtw side is a strict superset (added local evaluation_matrix.md exclusion). Take dynadtw. * packages/viscy-models/.../foundation/cell_dino.py — dynadtw added CellDinoModel wrapper + per-image normalize flag; modular-viscy-staging didn't touch post-branch. Take dynadtw. * pyproject.toml — modular-viscy-staging added pytest "slow" marker + addopts "-m not slow" (no-overlap, kept); both sides identically added aics* git sources, py312 bump, pythonpath. iohub: modular-viscy-staging kept a git override pin at afe11f6a (~0.3.3); dynadtw dropped it for iohub>=0.3.5 from PyPI. Took dynadtw's strategy (PyPI 0.3.5 is released; viscy-data dep pin already at >=0.3.5 on this side). - uv.lock regenerated from the merged pyproject.toml (iohub 0.3.5 from PyPI registry). Sanity check: uv sync --all-packages --all-extras + uv run pytest -m "not slow" applications/dynacell/tests/ shows 51 failed / 459 passed / 78 skipped. All 51 failures are pre-existing on dynacell-models HEAD (test_benchmark_config_composition strategy/leaf composition issues + test_evaluation_metrics gain-and-offset) — verified by running the same suite on the clean backbone-refactor worktree. The merge introduces no new regressions. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Opt-in modernization of the pix2pix3d adversarial training recipe, replacing the LSGAN-era defaults with the verified StyleGAN2 / R3GAN-style recipe. Architectures (
UNetViT3Dgenerator,MultiScalePatchGAN3Ddiscriminator) are unchanged for benchmark fairness.The 14 existing pix2pix3d leaves continue to run LSGAN unchanged — Python
__init__defaults are legacy-safe, so opt-in is only via the newpix2pix3d_unetvit_fit_modernized.ymloverlay.What changed
New losses (
packages/viscy-models/src/viscy_models/gan/losses.py) — verbatim from reference codebases:nonsat_d_loss/nonsat_g_loss— softplus loss fromstylegan2-ada-pytorch/training/loss.pyrpgan_d_loss/rpgan_g_loss— relativistic pairing frombrownvc/R3GAN/Trainer.pyr1_penalty/r2_penalty— multi-scale Mescheder zero-centered gradient penalty (per-scale grad, not grad-of-sum, to matchMultiScalePatchGAN3D's pool-derived second scale)DynacellGANwiring (applications/dynacell/src/dynacell/engine.py) — all knobs default to legacy-safe values:loss_type ∈ {lsgan, nonsat, rpgan},lambda_adv,r1_gamma,r2_gamma,r1_every,ema_kimg,lecam_gamma,lecam_decay,use_ema_at_predict_d_step_countbuffer drives lazy R1 (Lightning'sglobal_stepadvances per opt.step → can't be used; would fire R1 every 8 batches instead of 16)(γ > 0) and (counter % r1_every == 0)outer guard so legacy leaves skip the grad-of-grad path entirely (no spectral-norm power-iteration drift)autocast(enabled=False)— Lightning bf16-mixed gotcha; reference codebases don't AMPd_realagainst post-update Dema_kimg is not None; decay derived per-step from global batch size via the StyleGAN2 formulaall_gather-synced batch means (no implicit DDP all-reduce on buffers)validation_stepruns both raw and EMA forwards; emits bothloss/validate(back-compat) andloss/validate_ema(modernized monitor target)ckpt_pathload switched tostrict=Falsewith filtered warnings + post-load EMA seed for legacy ckptsNew configs:
pix2pix3d_unetvit_fit_modernized.yml+_ddpsiblinger/pix2pix3d_unetvit/ipsc_confocal/train_4gpu_modernized.yml(Run A)Run A launched
SLURM job
33207301— SEC61B headline run. Distinct save subtree (.../pix2pix3d_unetvit_modernized/) so there's no collision with the in-flight LSGAN run (job 33190465).Monitor armed for state transitions, R1 firing, and terminal failures.
Plan
/home/alex.kalinin/.claude/plans/pix2pix3d-gan-modernization.md(local) — went through 4 rounds of/review-plancovering correctness, regressions, merge compatibility, and external-source cross-reference. All 11 blocking findings verified addressed via/verify-planpost-implementation.Test plan
loss/d_train_epochstabilizes in [0.2, 0.6] band (vs. LSGAN baseline's collapse to 0.017 in job 33190465)loss/g_adv_train_epochnear log(2) ≈ 0.69 (vs. LSGAN's 0.97)loss/validate_ema≤ LSGAN baseline at matched epoch count🤖 Generated with Claude Code