diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc_confocal/predict__a549_mantis_denv.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc_confocal/predict__a549_mantis_denv.yml new file mode 100644 index 000000000..8b7c61640 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc_confocal/predict__a549_mantis_denv.yml @@ -0,0 +1,44 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons` arm (predict) of the baseline leaf nucleus/fcmae_vscyto3d_scratch/ipsc_confocal/predict__a549_mantis_denv.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/a549_mantis_h2b_denv.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/fcmae_vscyto3d_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: a549_mantis_h2b_denv + model_name: fcmae_vscyto3d_scratch_segcons + experiment_id: nucleus__ipsc_confocal__fcmae_vscyto3d_scratch_segcons__a549_mantis_h2b_denv + dataset_ref: + target: h2b +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fcmae_vscyto3d_scratch_segcons/checkpoints/last.ckpt +data: + init_args: + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc/a549__denv/prediction.zarr +launcher: + job_name: FCMAE_VSCyto3D_Scratch_PRED_NUCL_ON_A549_DENV_segcons + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc/a549__denv diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc_confocal/predict__a549_mantis_mock.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc_confocal/predict__a549_mantis_mock.yml new file mode 100644 index 000000000..f2934660d --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc_confocal/predict__a549_mantis_mock.yml @@ -0,0 +1,44 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons` arm (predict) of the baseline leaf nucleus/fcmae_vscyto3d_scratch/ipsc_confocal/predict__a549_mantis_mock.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/a549_mantis_h2b_mock.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/fcmae_vscyto3d_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: a549_mantis_h2b_mock + model_name: fcmae_vscyto3d_scratch_segcons + experiment_id: nucleus__ipsc_confocal__fcmae_vscyto3d_scratch_segcons__a549_mantis_h2b_mock + dataset_ref: + target: h2b +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fcmae_vscyto3d_scratch_segcons/checkpoints/last.ckpt +data: + init_args: + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc/a549__mock/prediction.zarr +launcher: + job_name: FCMAE_VSCyto3D_Scratch_PRED_NUCL_ON_A549_MOCK_segcons + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc/a549__mock diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc_confocal/predict__a549_mantis_zikv.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc_confocal/predict__a549_mantis_zikv.yml new file mode 100644 index 000000000..616b620c5 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc_confocal/predict__a549_mantis_zikv.yml @@ -0,0 +1,44 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons` arm (predict) of the baseline leaf nucleus/fcmae_vscyto3d_scratch/ipsc_confocal/predict__a549_mantis_zikv.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/a549_mantis_h2b_zikv.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/fcmae_vscyto3d_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: a549_mantis_h2b_zikv + model_name: fcmae_vscyto3d_scratch_segcons + experiment_id: nucleus__ipsc_confocal__fcmae_vscyto3d_scratch_segcons__a549_mantis_h2b_zikv + dataset_ref: + target: h2b +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fcmae_vscyto3d_scratch_segcons/checkpoints/last.ckpt +data: + init_args: + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc/a549__zikv/prediction.zarr +launcher: + job_name: FCMAE_VSCyto3D_Scratch_PRED_NUCL_ON_A549_ZIKV_segcons + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc/a549__zikv diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc_confocal/predict__ipsc_confocal.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc_confocal/predict__ipsc_confocal.yml new file mode 100644 index 000000000..8ca841ec8 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc_confocal/predict__ipsc_confocal.yml @@ -0,0 +1,42 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons` arm (predict) of the baseline leaf nucleus/fcmae_vscyto3d_scratch/ipsc_confocal/predict__ipsc_confocal.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/ipsc_confocal.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/fcmae_vscyto3d_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: ipsc_confocal + model_name: fcmae_vscyto3d_scratch_segcons + experiment_id: nucleus__ipsc_confocal__fcmae_vscyto3d_scratch_segcons__ipsc_confocal +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fcmae_vscyto3d_scratch_segcons/checkpoints/last.ckpt +data: + init_args: + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc/ipsc/prediction.zarr +launcher: + job_name: FCMAE_VSCyto3D_Scratch_PRED_NUCL_segcons + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc/ipsc diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc_confocal/train.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc_confocal/train.yml new file mode 100644 index 000000000..c00b1ee30 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc_confocal/train.yml @@ -0,0 +1,71 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons` arm (fit) of the baseline leaf nucleus/fcmae_vscyto3d_scratch/ipsc_confocal/train.yml. +# Recipe delta vs that leaf: the segcons_ctrl fork + model seg_consistency (SegConsistencyLoss, Cellpose-SAM cpsam, slab halfwidth 1, 240-px crop (one cellpose tile), per-tile activation checkpointing) + seg_consistency_weight (gradient-norm calibrated at the fork checkpoint). +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/train_sets/ipsc_confocal.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/data_overlays/fcmae_vscyto3d_fit.yml +- ../../../_internal/shared/model/model_overlays/fcmae_vscyto3d_fit.yml +- ../../../_internal/shared/model/launcher_profiles/mode_fit.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_4gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + train_set: ipsc_confocal + model_name: fcmae_vscyto3d_scratch_segcons + experiment_id: nucleus__ipsc_confocal__fcmae_vscyto3d_scratch_segcons +data: + init_args: + augmentations: + - class_path: viscy_transforms.RandWeightedCropd + init_args: + keys: + - Phase3D + - Nuclei + w_key: Nuclei + spatial_size: + - 20 + - 600 + - 600 + num_samples: 4 +trainer: + logger: + init_args: + name: FCMAE_VSCyto3D_Scratch_iPSC_Nucleus_segcons + save_dir: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fcmae_vscyto3d_scratch_segcons + callbacks: + - class_path: lightning.pytorch.callbacks.LearningRateMonitor + init_args: + logging_interval: step + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: loss/validate + filename: epoch={epoch}-step={step}-loss={loss/validate:.3f} + auto_insert_metric_name: false + every_n_epochs: 1 + save_top_k: 5 + save_last: true + dirpath: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fcmae_vscyto3d_scratch_segcons/checkpoints + max_epochs: 20 +launcher: + job_name: FCMAE_VSCyto3D_Scratch_Nucleus_segcons + run_root: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fcmae_vscyto3d_scratch_segcons +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fcmae_vscyto3d_scratch_v2/checkpoints/latest-epoch=199-step=62400.ckpt + lr: 4.0e-05 + warmup_steps: 312 + seg_consistency: + class_path: viscy_utils.losses.SegConsistencyLoss + init_args: + segmenter: + class_path: dynacell.cellpose_dense.CellposeDense + init_args: + model_name: cpsam + checkpoint_tiles: 1 + slab_halfwidth: 1 + crop_size: 240 + seg_consistency_weight: 0.032 diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc_confocal/predict__a549_mantis_denv.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc_confocal/predict__a549_mantis_denv.yml new file mode 100644 index 000000000..3d576b846 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc_confocal/predict__a549_mantis_denv.yml @@ -0,0 +1,44 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons_ctrl` arm (predict) of the baseline leaf nucleus/fcmae_vscyto3d_scratch/ipsc_confocal/predict__a549_mantis_denv.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/a549_mantis_h2b_denv.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/fcmae_vscyto3d_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: a549_mantis_h2b_denv + model_name: fcmae_vscyto3d_scratch_segcons_ctrl + experiment_id: nucleus__ipsc_confocal__fcmae_vscyto3d_scratch_segcons_ctrl__a549_mantis_h2b_denv + dataset_ref: + target: h2b +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/checkpoints/last.ckpt +data: + init_args: + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc/a549__denv/prediction.zarr +launcher: + job_name: FCMAE_VSCyto3D_Scratch_PRED_NUCL_ON_A549_DENV_segcons_ctrl + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc/a549__denv diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc_confocal/predict__a549_mantis_mock.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc_confocal/predict__a549_mantis_mock.yml new file mode 100644 index 000000000..4123c2041 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc_confocal/predict__a549_mantis_mock.yml @@ -0,0 +1,44 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons_ctrl` arm (predict) of the baseline leaf nucleus/fcmae_vscyto3d_scratch/ipsc_confocal/predict__a549_mantis_mock.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/a549_mantis_h2b_mock.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/fcmae_vscyto3d_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: a549_mantis_h2b_mock + model_name: fcmae_vscyto3d_scratch_segcons_ctrl + experiment_id: nucleus__ipsc_confocal__fcmae_vscyto3d_scratch_segcons_ctrl__a549_mantis_h2b_mock + dataset_ref: + target: h2b +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/checkpoints/last.ckpt +data: + init_args: + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc/a549__mock/prediction.zarr +launcher: + job_name: FCMAE_VSCyto3D_Scratch_PRED_NUCL_ON_A549_MOCK_segcons_ctrl + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc/a549__mock diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc_confocal/predict__a549_mantis_zikv.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc_confocal/predict__a549_mantis_zikv.yml new file mode 100644 index 000000000..cbcd8072a --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc_confocal/predict__a549_mantis_zikv.yml @@ -0,0 +1,44 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons_ctrl` arm (predict) of the baseline leaf nucleus/fcmae_vscyto3d_scratch/ipsc_confocal/predict__a549_mantis_zikv.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/a549_mantis_h2b_zikv.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/fcmae_vscyto3d_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: a549_mantis_h2b_zikv + model_name: fcmae_vscyto3d_scratch_segcons_ctrl + experiment_id: nucleus__ipsc_confocal__fcmae_vscyto3d_scratch_segcons_ctrl__a549_mantis_h2b_zikv + dataset_ref: + target: h2b +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/checkpoints/last.ckpt +data: + init_args: + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc/a549__zikv/prediction.zarr +launcher: + job_name: FCMAE_VSCyto3D_Scratch_PRED_NUCL_ON_A549_ZIKV_segcons_ctrl + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc/a549__zikv diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc_confocal/predict__ipsc_confocal.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc_confocal/predict__ipsc_confocal.yml new file mode 100644 index 000000000..ed2966df8 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc_confocal/predict__ipsc_confocal.yml @@ -0,0 +1,42 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons_ctrl` arm (predict) of the baseline leaf nucleus/fcmae_vscyto3d_scratch/ipsc_confocal/predict__ipsc_confocal.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/ipsc_confocal.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/fcmae_vscyto3d_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: ipsc_confocal + model_name: fcmae_vscyto3d_scratch_segcons_ctrl + experiment_id: nucleus__ipsc_confocal__fcmae_vscyto3d_scratch_segcons_ctrl__ipsc_confocal +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/checkpoints/last.ckpt +data: + init_args: + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc/ipsc/prediction.zarr +launcher: + job_name: FCMAE_VSCyto3D_Scratch_PRED_NUCL_segcons_ctrl + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc/ipsc diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc_confocal/train.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc_confocal/train.yml new file mode 100644 index 000000000..fc7159dc5 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc_confocal/train.yml @@ -0,0 +1,60 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons_ctrl` arm (fit) of the baseline leaf nucleus/fcmae_vscyto3d_scratch/ipsc_confocal/train.yml. +# Recipe delta vs that leaf: fork: model ckpt_path (the converged baseline checkpoint, weights only) + the fine-tune budget and LR of FORK_SCHEDULES; no other change (the no-term control of segcons). +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/train_sets/ipsc_confocal.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/data_overlays/fcmae_vscyto3d_fit.yml +- ../../../_internal/shared/model/model_overlays/fcmae_vscyto3d_fit.yml +- ../../../_internal/shared/model/launcher_profiles/mode_fit.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_4gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + train_set: ipsc_confocal + model_name: fcmae_vscyto3d_scratch_segcons_ctrl + experiment_id: nucleus__ipsc_confocal__fcmae_vscyto3d_scratch_segcons_ctrl +data: + init_args: + augmentations: + - class_path: viscy_transforms.RandWeightedCropd + init_args: + keys: + - Phase3D + - Nuclei + w_key: Nuclei + spatial_size: + - 20 + - 600 + - 600 + num_samples: 4 +trainer: + logger: + init_args: + name: FCMAE_VSCyto3D_Scratch_iPSC_Nucleus_segcons_ctrl + save_dir: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl + callbacks: + - class_path: lightning.pytorch.callbacks.LearningRateMonitor + init_args: + logging_interval: step + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: loss/validate + filename: epoch={epoch}-step={step}-loss={loss/validate:.3f} + auto_insert_metric_name: false + every_n_epochs: 1 + save_top_k: 5 + save_last: true + dirpath: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/checkpoints + max_epochs: 20 +launcher: + job_name: FCMAE_VSCyto3D_Scratch_Nucleus_segcons_ctrl + run_root: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fcmae_vscyto3d_scratch_v2/checkpoints/latest-epoch=199-step=62400.ckpt + lr: 4.0e-05 + warmup_steps: 312 diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/predict__a549_mantis_h2b_denv.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/predict__a549_mantis_h2b_denv.yml new file mode 100644 index 000000000..d399ca4d5 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/predict__a549_mantis_h2b_denv.yml @@ -0,0 +1,44 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons` arm (predict) of the baseline leaf nucleus/fnet3d_vscyto3daug/ipsc_confocal/predict__a549_mantis_h2b_denv.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/a549_mantis_h2b_denv.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/fnet3d_paper_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: a549_mantis_h2b_denv + model_name: fnet3d_vscyto3daug_segcons + experiment_id: nucleus__ipsc_confocal__fnet3d_vscyto3daug_segcons__a549_mantis_h2b_denv + dataset_ref: + target: h2b +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fnet3d_vscyto3daug_segcons/checkpoints/last.ckpt +data: + init_args: + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/fnet3d_vscyto3daug_segcons/ipsc/a549__denv/prediction.zarr +launcher: + job_name: FNet3DVSAug_PRED_NUCL_ON_A549_DENV_segcons + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/fnet3d_vscyto3daug_segcons/ipsc/a549__denv diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/predict__a549_mantis_h2b_mock.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/predict__a549_mantis_h2b_mock.yml new file mode 100644 index 000000000..086b84190 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/predict__a549_mantis_h2b_mock.yml @@ -0,0 +1,44 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons` arm (predict) of the baseline leaf nucleus/fnet3d_vscyto3daug/ipsc_confocal/predict__a549_mantis_h2b_mock.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/a549_mantis_h2b_mock.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/fnet3d_paper_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: a549_mantis_h2b_mock + model_name: fnet3d_vscyto3daug_segcons + experiment_id: nucleus__ipsc_confocal__fnet3d_vscyto3daug_segcons__a549_mantis_h2b_mock + dataset_ref: + target: h2b +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fnet3d_vscyto3daug_segcons/checkpoints/last.ckpt +data: + init_args: + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/fnet3d_vscyto3daug_segcons/ipsc/a549__mock/prediction.zarr +launcher: + job_name: FNet3DVSAug_PRED_NUCL_ON_A549_MOCK_segcons + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/fnet3d_vscyto3daug_segcons/ipsc/a549__mock diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/predict__a549_mantis_h2b_zikv.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/predict__a549_mantis_h2b_zikv.yml new file mode 100644 index 000000000..4e0b0365e --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/predict__a549_mantis_h2b_zikv.yml @@ -0,0 +1,44 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons` arm (predict) of the baseline leaf nucleus/fnet3d_vscyto3daug/ipsc_confocal/predict__a549_mantis_h2b_zikv.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/a549_mantis_h2b_zikv.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/fnet3d_paper_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: a549_mantis_h2b_zikv + model_name: fnet3d_vscyto3daug_segcons + experiment_id: nucleus__ipsc_confocal__fnet3d_vscyto3daug_segcons__a549_mantis_h2b_zikv + dataset_ref: + target: h2b +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fnet3d_vscyto3daug_segcons/checkpoints/last.ckpt +data: + init_args: + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/fnet3d_vscyto3daug_segcons/ipsc/a549__zikv/prediction.zarr +launcher: + job_name: FNet3DVSAug_PRED_NUCL_ON_A549_ZIKV_segcons + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/fnet3d_vscyto3daug_segcons/ipsc/a549__zikv diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/predict__ipsc_confocal.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/predict__ipsc_confocal.yml new file mode 100644 index 000000000..7bf84abaa --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/predict__ipsc_confocal.yml @@ -0,0 +1,42 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons` arm (predict) of the baseline leaf nucleus/fnet3d_vscyto3daug/ipsc_confocal/predict__ipsc_confocal.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/ipsc_confocal.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/fnet3d_paper_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: ipsc_confocal + model_name: fnet3d_vscyto3daug_segcons + experiment_id: nucleus__ipsc_confocal__fnet3d_vscyto3daug_segcons__ipsc_confocal +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fnet3d_vscyto3daug_segcons/checkpoints/last.ckpt +data: + init_args: + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/fnet3d_vscyto3daug_segcons/ipsc/ipsc/prediction.zarr +launcher: + job_name: FNet3DVSAug_PRED_NUCL_segcons + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/fnet3d_vscyto3daug_segcons/ipsc/ipsc diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/train.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/train.yml new file mode 100644 index 000000000..d043e890b --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/train.yml @@ -0,0 +1,99 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons` arm (fit) of the baseline leaf nucleus/fnet3d_vscyto3daug/ipsc_confocal/train.yml. +# Recipe delta vs that leaf: the segcons_ctrl fork + model seg_consistency (SegConsistencyLoss, Cellpose-SAM cpsam, slab halfwidth 1, 240-px crop (one cellpose tile), per-tile activation checkpointing) + seg_consistency_weight (gradient-norm calibrated at the fork checkpoint). +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/train_sets/ipsc_confocal.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/data_overlays/fnet3d_vscyto3daug_fit.yml +- ../../../_internal/shared/model/model_overlays/fnet3d_paper_fit.yml +- ../../../_internal/shared/model/launcher_profiles/mode_fit.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_gpu_any_long.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + train_set: ipsc_confocal + model_name: fnet3d_vscyto3daug_segcons + experiment_id: nucleus__ipsc_confocal__fnet3d_vscyto3daug_segcons +data: + init_args: + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Nuclei + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: + - class_path: viscy_transforms.RandWeightedCropd + init_args: + keys: + - Phase3D + - Nuclei + w_key: Nuclei + spatial_size: + - 32 + - 600 + - 600 + num_samples: 4 + val_augmentations: + - class_path: viscy_transforms.CenterSpatialCropd + init_args: + keys: + - Phase3D + - Nuclei + roi_size: + - 32 + - 384 + - 384 +trainer: + precision: bf16-mixed + logger: + init_args: + name: FNet3D_VSAug_iPSC_NUCL_segcons + save_dir: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fnet3d_vscyto3daug_segcons + callbacks: + - class_path: lightning.pytorch.callbacks.LearningRateMonitor + init_args: + logging_interval: step + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: loss/validate + filename: epoch={epoch}-step={step}-loss={loss/validate:.3f} + auto_insert_metric_name: false + every_n_epochs: 1 + save_top_k: 4 + save_last: true + dirpath: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fnet3d_vscyto3daug_segcons/checkpoints + max_steps: 20000 +launcher: + job_name: FNet3DVSAug_iPSC_NUCL_segcons + run_root: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fnet3d_vscyto3daug_segcons + sbatch: + constraint: h100|h200 + mem: 512G +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fnet3d_vscyto3daug_v2/checkpoints/latest-epoch=76-step=200000.ckpt + lr: 0.0001 + seg_consistency: + class_path: viscy_utils.losses.SegConsistencyLoss + init_args: + segmenter: + class_path: dynacell.cellpose_dense.CellposeDense + init_args: + model_name: cpsam + checkpoint_tiles: 1 + slab_halfwidth: 1 + crop_size: 240 + seg_consistency_weight: 0.37 diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc_confocal/predict__a549_mantis_h2b_denv.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc_confocal/predict__a549_mantis_h2b_denv.yml new file mode 100644 index 000000000..4b5777af5 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc_confocal/predict__a549_mantis_h2b_denv.yml @@ -0,0 +1,44 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons_ctrl` arm (predict) of the baseline leaf nucleus/fnet3d_vscyto3daug/ipsc_confocal/predict__a549_mantis_h2b_denv.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/a549_mantis_h2b_denv.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/fnet3d_paper_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: a549_mantis_h2b_denv + model_name: fnet3d_vscyto3daug_segcons_ctrl + experiment_id: nucleus__ipsc_confocal__fnet3d_vscyto3daug_segcons_ctrl__a549_mantis_h2b_denv + dataset_ref: + target: h2b +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fnet3d_vscyto3daug_segcons_ctrl/checkpoints/last.ckpt +data: + init_args: + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc/a549__denv/prediction.zarr +launcher: + job_name: FNet3DVSAug_PRED_NUCL_ON_A549_DENV_segcons_ctrl + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc/a549__denv diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc_confocal/predict__a549_mantis_h2b_mock.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc_confocal/predict__a549_mantis_h2b_mock.yml new file mode 100644 index 000000000..0393320d5 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc_confocal/predict__a549_mantis_h2b_mock.yml @@ -0,0 +1,44 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons_ctrl` arm (predict) of the baseline leaf nucleus/fnet3d_vscyto3daug/ipsc_confocal/predict__a549_mantis_h2b_mock.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/a549_mantis_h2b_mock.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/fnet3d_paper_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: a549_mantis_h2b_mock + model_name: fnet3d_vscyto3daug_segcons_ctrl + experiment_id: nucleus__ipsc_confocal__fnet3d_vscyto3daug_segcons_ctrl__a549_mantis_h2b_mock + dataset_ref: + target: h2b +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fnet3d_vscyto3daug_segcons_ctrl/checkpoints/last.ckpt +data: + init_args: + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc/a549__mock/prediction.zarr +launcher: + job_name: FNet3DVSAug_PRED_NUCL_ON_A549_MOCK_segcons_ctrl + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc/a549__mock diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc_confocal/predict__a549_mantis_h2b_zikv.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc_confocal/predict__a549_mantis_h2b_zikv.yml new file mode 100644 index 000000000..077c033d2 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc_confocal/predict__a549_mantis_h2b_zikv.yml @@ -0,0 +1,44 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons_ctrl` arm (predict) of the baseline leaf nucleus/fnet3d_vscyto3daug/ipsc_confocal/predict__a549_mantis_h2b_zikv.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/a549_mantis_h2b_zikv.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/fnet3d_paper_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: a549_mantis_h2b_zikv + model_name: fnet3d_vscyto3daug_segcons_ctrl + experiment_id: nucleus__ipsc_confocal__fnet3d_vscyto3daug_segcons_ctrl__a549_mantis_h2b_zikv + dataset_ref: + target: h2b +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fnet3d_vscyto3daug_segcons_ctrl/checkpoints/last.ckpt +data: + init_args: + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc/a549__zikv/prediction.zarr +launcher: + job_name: FNet3DVSAug_PRED_NUCL_ON_A549_ZIKV_segcons_ctrl + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc/a549__zikv diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc_confocal/predict__ipsc_confocal.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc_confocal/predict__ipsc_confocal.yml new file mode 100644 index 000000000..c7013cdf1 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc_confocal/predict__ipsc_confocal.yml @@ -0,0 +1,42 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons_ctrl` arm (predict) of the baseline leaf nucleus/fnet3d_vscyto3daug/ipsc_confocal/predict__ipsc_confocal.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/ipsc_confocal.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/fnet3d_paper_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: ipsc_confocal + model_name: fnet3d_vscyto3daug_segcons_ctrl + experiment_id: nucleus__ipsc_confocal__fnet3d_vscyto3daug_segcons_ctrl__ipsc_confocal +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fnet3d_vscyto3daug_segcons_ctrl/checkpoints/last.ckpt +data: + init_args: + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc/ipsc/prediction.zarr +launcher: + job_name: FNet3DVSAug_PRED_NUCL_segcons_ctrl + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc/ipsc diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc_confocal/train.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc_confocal/train.yml new file mode 100644 index 000000000..98ef16741 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc_confocal/train.yml @@ -0,0 +1,88 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons_ctrl` arm (fit) of the baseline leaf nucleus/fnet3d_vscyto3daug/ipsc_confocal/train.yml. +# Recipe delta vs that leaf: fork: model ckpt_path (the converged baseline checkpoint, weights only) + the fine-tune budget and LR of FORK_SCHEDULES; no other change (the no-term control of segcons). +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/train_sets/ipsc_confocal.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/data_overlays/fnet3d_vscyto3daug_fit.yml +- ../../../_internal/shared/model/model_overlays/fnet3d_paper_fit.yml +- ../../../_internal/shared/model/launcher_profiles/mode_fit.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_gpu_any_long.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + train_set: ipsc_confocal + model_name: fnet3d_vscyto3daug_segcons_ctrl + experiment_id: nucleus__ipsc_confocal__fnet3d_vscyto3daug_segcons_ctrl +data: + init_args: + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Nuclei + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: + - class_path: viscy_transforms.RandWeightedCropd + init_args: + keys: + - Phase3D + - Nuclei + w_key: Nuclei + spatial_size: + - 32 + - 600 + - 600 + num_samples: 4 + val_augmentations: + - class_path: viscy_transforms.CenterSpatialCropd + init_args: + keys: + - Phase3D + - Nuclei + roi_size: + - 32 + - 384 + - 384 +trainer: + precision: bf16-mixed + logger: + init_args: + name: FNet3D_VSAug_iPSC_NUCL_segcons_ctrl + save_dir: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fnet3d_vscyto3daug_segcons_ctrl + callbacks: + - class_path: lightning.pytorch.callbacks.LearningRateMonitor + init_args: + logging_interval: step + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: loss/validate + filename: epoch={epoch}-step={step}-loss={loss/validate:.3f} + auto_insert_metric_name: false + every_n_epochs: 1 + save_top_k: 4 + save_last: true + dirpath: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fnet3d_vscyto3daug_segcons_ctrl/checkpoints + max_steps: 20000 +launcher: + job_name: FNet3DVSAug_iPSC_NUCL_segcons_ctrl + run_root: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fnet3d_vscyto3daug_segcons_ctrl + sbatch: + constraint: h100|h200 + mem: 512G +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/fnet3d_vscyto3daug_v2/checkpoints/latest-epoch=76-step=200000.ckpt + lr: 0.0001 diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons/ipsc_confocal/predict__a549_mantis_denv.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons/ipsc_confocal/predict__a549_mantis_denv.yml new file mode 100644 index 000000000..963531e56 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons/ipsc_confocal/predict__a549_mantis_denv.yml @@ -0,0 +1,45 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons` arm (predict) of the baseline leaf nucleus/pix2pix2d_unetvit/ipsc_confocal/predict__a549_mantis_denv.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/a549_mantis_h2b_denv.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/pix2pix2d_unetvit_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: a549_mantis_h2b_denv + model_name: pix2pix2d_unetvit_segcons + experiment_id: nucleus__ipsc_confocal__pix2pix2d_unetvit_segcons__a549_mantis_h2b_denv + dataset_ref: + target: h2b +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/pix2pix2d_unetvit_segcons/checkpoints/last.ckpt +data: + init_args: + predict_z_planes: /hpc/projects/virtual_staining/training/dynacell/focus2d_planes/nucleus/a549__denv__h1.json + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/pix2pix2d_unetvit_segcons/ipsc/a549__denv/prediction.zarr +launcher: + job_name: pix2pix2D_PRED_NUCL_ON_A549_DENV_segcons + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/pix2pix2d_unetvit_segcons/ipsc/a549__denv diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons/ipsc_confocal/predict__a549_mantis_mock.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons/ipsc_confocal/predict__a549_mantis_mock.yml new file mode 100644 index 000000000..da1d53470 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons/ipsc_confocal/predict__a549_mantis_mock.yml @@ -0,0 +1,45 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons` arm (predict) of the baseline leaf nucleus/pix2pix2d_unetvit/ipsc_confocal/predict__a549_mantis_mock.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/a549_mantis_h2b_mock.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/pix2pix2d_unetvit_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: a549_mantis_h2b_mock + model_name: pix2pix2d_unetvit_segcons + experiment_id: nucleus__ipsc_confocal__pix2pix2d_unetvit_segcons__a549_mantis_h2b_mock + dataset_ref: + target: h2b +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/pix2pix2d_unetvit_segcons/checkpoints/last.ckpt +data: + init_args: + predict_z_planes: /hpc/projects/virtual_staining/training/dynacell/focus2d_planes/nucleus/a549__mock__h1.json + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/pix2pix2d_unetvit_segcons/ipsc/a549__mock/prediction.zarr +launcher: + job_name: pix2pix2D_PRED_NUCL_ON_A549_MOCK_segcons + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/pix2pix2d_unetvit_segcons/ipsc/a549__mock diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons/ipsc_confocal/predict__a549_mantis_zikv.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons/ipsc_confocal/predict__a549_mantis_zikv.yml new file mode 100644 index 000000000..17b7d7c16 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons/ipsc_confocal/predict__a549_mantis_zikv.yml @@ -0,0 +1,45 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons` arm (predict) of the baseline leaf nucleus/pix2pix2d_unetvit/ipsc_confocal/predict__a549_mantis_zikv.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/a549_mantis_h2b_zikv.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/pix2pix2d_unetvit_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: a549_mantis_h2b_zikv + model_name: pix2pix2d_unetvit_segcons + experiment_id: nucleus__ipsc_confocal__pix2pix2d_unetvit_segcons__a549_mantis_h2b_zikv + dataset_ref: + target: h2b +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/pix2pix2d_unetvit_segcons/checkpoints/last.ckpt +data: + init_args: + predict_z_planes: /hpc/projects/virtual_staining/training/dynacell/focus2d_planes/nucleus/a549__zikv__h1.json + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/pix2pix2d_unetvit_segcons/ipsc/a549__zikv/prediction.zarr +launcher: + job_name: pix2pix2D_PRED_NUCL_ON_A549_ZIKV_segcons + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/pix2pix2d_unetvit_segcons/ipsc/a549__zikv diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons/ipsc_confocal/predict__ipsc_confocal.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons/ipsc_confocal/predict__ipsc_confocal.yml new file mode 100644 index 000000000..8f4200afe --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons/ipsc_confocal/predict__ipsc_confocal.yml @@ -0,0 +1,43 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons` arm (predict) of the baseline leaf nucleus/pix2pix2d_unetvit/ipsc_confocal/predict__ipsc_confocal.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/ipsc_confocal.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/pix2pix2d_unetvit_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: ipsc_confocal + model_name: pix2pix2d_unetvit_segcons + experiment_id: nucleus__ipsc_confocal__pix2pix2d_unetvit_segcons__ipsc_confocal +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/pix2pix2d_unetvit_segcons/checkpoints/last.ckpt +data: + init_args: + predict_z_planes: /hpc/projects/virtual_staining/training/dynacell/focus2d_planes/nucleus/ipsc__h1.json + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/pix2pix2d_unetvit_segcons/ipsc/ipsc/prediction.zarr +launcher: + job_name: pix2pix2D_PRED_NUCL_ON_IPSC_segcons + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/pix2pix2d_unetvit_segcons/ipsc/ipsc diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons/ipsc_confocal/train.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons/ipsc_confocal/train.yml new file mode 100644 index 000000000..9b84b4839 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons/ipsc_confocal/train.yml @@ -0,0 +1,83 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons` arm (fit) of the baseline leaf nucleus/pix2pix2d_unetvit/ipsc_confocal/train.yml. +# Recipe delta vs that leaf: the segcons_ctrl fork + model seg_consistency (SegConsistencyLoss, Cellpose-SAM cpsam, slab halfwidth 1, 240-px crop (one cellpose tile), per-tile activation checkpointing) + seg_consistency_weight (gradient-norm calibrated at the fork checkpoint). +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/train_sets/ipsc_confocal.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/data_overlays/pix2pix2d_unetvit_fit.yml +- ../../../_internal/shared/model/model_overlays/pix2pix2d_unetvit_fit.yml +- ../../../_internal/shared/model/launcher_profiles/mode_fit.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_4gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + gene: Nuclei + target: nucleus + target_id: nucleus + train_set: ipsc_confocal + model_name: pix2pix2d_unetvit_segcons + experiment_id: nucleus__ipsc_confocal__pix2pix2d_unetvit_segcons + dataset_ref: null +data: + init_args: + source_channel: Phase3D + target_channel: Nuclei + data_path: /hpc/projects/virtual_staining/training/dynacell/ipsc/dataset_v4/train/cell_focus.zarr + augmentations: + - class_path: viscy_transforms.RandWeightedCropd + init_args: + keys: + - Phase3D + - Nuclei + w_key: Nuclei + spatial_size: + - 1 + - 624 + - 624 + num_samples: 2 +trainer: + logger: + init_args: + name: pix2pix2d_unetvit_iPSC_NUCL_segcons + save_dir: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/pix2pix2d_unetvit_segcons + callbacks: + - class_path: lightning.pytorch.callbacks.LearningRateMonitor + init_args: + logging_interval: step + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: loss/validate_ema + filename: epoch={epoch}-step={step}-loss={loss/validate_ema:.3f} + auto_insert_metric_name: false + every_n_epochs: 1 + save_top_k: 4 + save_last: true + dirpath: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/pix2pix2d_unetvit_segcons/checkpoints + max_epochs: 4 +launcher: + job_name: pix2pix2d_unetvit_iPSC_NUCL_segcons + run_root: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/pix2pix2d_unetvit_segcons + sbatch: + constraint: a40|a6000|l40s|h100|h200 + mem: 128G +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/pix2pix2d_unetvit/checkpoints/latest-epoch=39-step=20000.ckpt + lr_g: 2.0e-05 + lr_d: 2.0e-05 + warmup_steps: 50 + ema_kimg: 2.0 + seg_consistency: + class_path: viscy_utils.losses.SegConsistencyLoss + init_args: + segmenter: + class_path: dynacell.cellpose_dense.CellposeDense + init_args: + model_name: cpsam + checkpoint_tiles: 1 + slab_halfwidth: 1 + crop_size: 240 + seg_consistency_weight: 0.19 diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/predict__a549_mantis_denv.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/predict__a549_mantis_denv.yml new file mode 100644 index 000000000..4007ddf28 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/predict__a549_mantis_denv.yml @@ -0,0 +1,45 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons_ctrl` arm (predict) of the baseline leaf nucleus/pix2pix2d_unetvit/ipsc_confocal/predict__a549_mantis_denv.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/a549_mantis_h2b_denv.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/pix2pix2d_unetvit_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: a549_mantis_h2b_denv + model_name: pix2pix2d_unetvit_segcons_ctrl + experiment_id: nucleus__ipsc_confocal__pix2pix2d_unetvit_segcons_ctrl__a549_mantis_h2b_denv + dataset_ref: + target: h2b +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/pix2pix2d_unetvit_segcons_ctrl/checkpoints/last.ckpt +data: + init_args: + predict_z_planes: /hpc/projects/virtual_staining/training/dynacell/focus2d_planes/nucleus/a549__denv__h1.json + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc/a549__denv/prediction.zarr +launcher: + job_name: pix2pix2D_PRED_NUCL_ON_A549_DENV_segcons_ctrl + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc/a549__denv diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/predict__a549_mantis_mock.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/predict__a549_mantis_mock.yml new file mode 100644 index 000000000..372c4b234 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/predict__a549_mantis_mock.yml @@ -0,0 +1,45 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons_ctrl` arm (predict) of the baseline leaf nucleus/pix2pix2d_unetvit/ipsc_confocal/predict__a549_mantis_mock.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/a549_mantis_h2b_mock.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/pix2pix2d_unetvit_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: a549_mantis_h2b_mock + model_name: pix2pix2d_unetvit_segcons_ctrl + experiment_id: nucleus__ipsc_confocal__pix2pix2d_unetvit_segcons_ctrl__a549_mantis_h2b_mock + dataset_ref: + target: h2b +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/pix2pix2d_unetvit_segcons_ctrl/checkpoints/last.ckpt +data: + init_args: + predict_z_planes: /hpc/projects/virtual_staining/training/dynacell/focus2d_planes/nucleus/a549__mock__h1.json + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc/a549__mock/prediction.zarr +launcher: + job_name: pix2pix2D_PRED_NUCL_ON_A549_MOCK_segcons_ctrl + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc/a549__mock diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/predict__a549_mantis_zikv.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/predict__a549_mantis_zikv.yml new file mode 100644 index 000000000..60f17874e --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/predict__a549_mantis_zikv.yml @@ -0,0 +1,45 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons_ctrl` arm (predict) of the baseline leaf nucleus/pix2pix2d_unetvit/ipsc_confocal/predict__a549_mantis_zikv.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/a549_mantis_h2b_zikv.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/pix2pix2d_unetvit_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: a549_mantis_h2b_zikv + model_name: pix2pix2d_unetvit_segcons_ctrl + experiment_id: nucleus__ipsc_confocal__pix2pix2d_unetvit_segcons_ctrl__a549_mantis_h2b_zikv + dataset_ref: + target: h2b +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/pix2pix2d_unetvit_segcons_ctrl/checkpoints/last.ckpt +data: + init_args: + predict_z_planes: /hpc/projects/virtual_staining/training/dynacell/focus2d_planes/nucleus/a549__zikv__h1.json + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc/a549__zikv/prediction.zarr +launcher: + job_name: pix2pix2D_PRED_NUCL_ON_A549_ZIKV_segcons_ctrl + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc/a549__zikv diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/predict__ipsc_confocal.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/predict__ipsc_confocal.yml new file mode 100644 index 000000000..4b645ad0d --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/predict__ipsc_confocal.yml @@ -0,0 +1,43 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons_ctrl` arm (predict) of the baseline leaf nucleus/pix2pix2d_unetvit/ipsc_confocal/predict__ipsc_confocal.yml. +# Inference is identical to the baseline's; only the checkpoint, store and names differ. +# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`. +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/predict_sets/ipsc_confocal.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/model_overlays/pix2pix2d_unetvit_predict.yml +- ../../../_internal/shared/model/launcher_profiles/mode_predict.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_predict_any_gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + trained_on: ipsc_confocal + predict_set: ipsc_confocal + model_name: pix2pix2d_unetvit_segcons_ctrl + experiment_id: nucleus__ipsc_confocal__pix2pix2d_unetvit_segcons_ctrl__ipsc_confocal +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/pix2pix2d_unetvit_segcons_ctrl/checkpoints/last.ckpt +data: + init_args: + predict_z_planes: /hpc/projects/virtual_staining/training/dynacell/focus2d_planes/nucleus/ipsc__h1.json + normalizations: + - class_path: viscy_transforms.NormalizeSampled + init_args: + keys: + - Phase3D + level: fov_statistics + subtrahend: mean + divisor: std + augmentations: [] +trainer: + callbacks: + - class_path: viscy_utils.callbacks.prediction_writer.HCSPredictionWriter + init_args: + output_store: /hpc/projects/virtual_staining/training/dynacell/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc/ipsc/prediction.zarr +launcher: + job_name: pix2pix2D_PRED_NUCL_ON_IPSC_segcons_ctrl + run_root: /hpc/projects/virtual_staining/training/dynacell/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc/ipsc diff --git a/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/train.yml b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/train.yml new file mode 100644 index 000000000..3bde11e21 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/train.yml @@ -0,0 +1,72 @@ +# GENERATED by applications/dynacell/tools/generate_spotlight_v2_leaves.py -- edit the +# generator, not this file (plan: ~/.claude/plans/spotlight-v2.md). +# Spotlight-v2 `segcons_ctrl` arm (fit) of the baseline leaf nucleus/pix2pix2d_unetvit/ipsc_confocal/train.yml. +# Recipe delta vs that leaf: fork: model ckpt_path (the converged baseline checkpoint, weights only) + the fine-tune budget and LR of FORK_SCHEDULES; no other change (the no-term control of segcons). +# Everything else, and the rationale for it, is the baseline leaf's. +base: +- ../../../_internal/shared/model/train_sets/ipsc_confocal.yml +- ../../../_internal/shared/model/targets/nucleus.yml +- ../../../_internal/shared/model/data_overlays/pix2pix2d_unetvit_fit.yml +- ../../../_internal/shared/model/model_overlays/pix2pix2d_unetvit_fit.yml +- ../../../_internal/shared/model/launcher_profiles/mode_fit.yml +- ../../../_internal/shared/model/launcher_profiles/hardware_4gpu.yml +- ../../../_internal/shared/model/launcher_profiles/runtime_shared.yml +benchmark: + task: virtual_staining + organelle: nucleus + gene: Nuclei + target: nucleus + target_id: nucleus + train_set: ipsc_confocal + model_name: pix2pix2d_unetvit_segcons_ctrl + experiment_id: nucleus__ipsc_confocal__pix2pix2d_unetvit_segcons_ctrl + dataset_ref: null +data: + init_args: + source_channel: Phase3D + target_channel: Nuclei + data_path: /hpc/projects/virtual_staining/training/dynacell/ipsc/dataset_v4/train/cell_focus.zarr + augmentations: + - class_path: viscy_transforms.RandWeightedCropd + init_args: + keys: + - Phase3D + - Nuclei + w_key: Nuclei + spatial_size: + - 1 + - 624 + - 624 + num_samples: 2 +trainer: + logger: + init_args: + name: pix2pix2d_unetvit_iPSC_NUCL_segcons_ctrl + save_dir: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/pix2pix2d_unetvit_segcons_ctrl + callbacks: + - class_path: lightning.pytorch.callbacks.LearningRateMonitor + init_args: + logging_interval: step + - class_path: lightning.pytorch.callbacks.ModelCheckpoint + init_args: + monitor: loss/validate_ema + filename: epoch={epoch}-step={step}-loss={loss/validate_ema:.3f} + auto_insert_metric_name: false + every_n_epochs: 1 + save_top_k: 4 + save_last: true + dirpath: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/pix2pix2d_unetvit_segcons_ctrl/checkpoints + max_epochs: 4 +launcher: + job_name: pix2pix2d_unetvit_iPSC_NUCL_segcons_ctrl + run_root: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/pix2pix2d_unetvit_segcons_ctrl + sbatch: + constraint: a40|a6000|l40s|h100|h200 + mem: 128G +model: + init_args: + ckpt_path: /hpc/projects/comp.micro/virtual_staining/models/dynacell/ipsc/nucleus/pix2pix2d_unetvit/checkpoints/latest-epoch=39-step=20000.ckpt + lr_g: 2.0e-05 + lr_d: 2.0e-05 + warmup_steps: 50 + ema_kimg: 2.0 diff --git a/applications/dynacell/src/dynacell/cellpose_dense.py b/applications/dynacell/src/dynacell/cellpose_dense.py new file mode 100644 index 000000000..260ed9750 --- /dev/null +++ b/applications/dynacell/src/dynacell/cellpose_dense.py @@ -0,0 +1,206 @@ +"""Differentiable dense outputs of a frozen Cellpose-SAM model (flows + cell-probability logit). + +Ported from the G4 segmenter gate (``experiments/2026-10-15_conditional-generation-review/ +gates/G4/segmenters.py``), which verified it against stock ``CellposeModel.eval``: flows and +cell-probability logits within 2.9e-6 max abs on real iPSC frames, and central +finite-difference gradients within 2.5e-3 relative error (cpsam). + +The input may be in any positive-affine intensity units (raw fluorescence, or a +normalized target or prediction): cellpose's ``normalize99`` (1st/99th percentile -> +[0, 1], no clip, flat image -> 0) is applied per image with its percentiles computed +under ``no_grad``, so gradients flow through the pixels and not through the statistics. +This mirrors how the eval feeds cpdino (``normalize=True`` on each side independently). + +The network is the stock checkpoint in cellpose's model directory +(``CELLPOSE_LOCAL_MODELS_PATH`` or ``~/.cellpose/models``), loaded by absolute path so a +missing file raises instead of falling back to another model or a download. +""" + +import numpy as np +import torch +import torch.nn.functional as F +from cellpose import models as cp_models +from cellpose import transforms as cp_transforms +from torch import Tensor, nn +from torch.utils.checkpoint import checkpoint + +__all__ = ["CellposeDense", "normalize99_diff"] + + +def normalize99_diff(x: Tensor) -> Tensor: + """Apply cellpose ``normalize99`` per image, with the percentiles under ``no_grad``. + + Parameters + ---------- + x : Tensor + Images, shape ``(B, ...)``. + + Returns + ------- + Tensor + ``(x - p1) / (p99 - p1)`` per image, unclipped; 0 where the percentile + range is below 1e-3 (cellpose's flat-image rule). + """ + with torch.no_grad(): + flat = x.reshape(x.shape[0], -1).float() + lo, hi = torch.quantile(flat, 0.01, dim=1), torch.quantile(flat, 0.99, dim=1) + ok = (hi - lo) > 1e-3 + scale = torch.where(ok, 1.0 / (hi - lo).clamp_min(1e-12), torch.zeros_like(hi)) + shape = (-1, *([1] * (x.ndim - 1))) + return (x - lo.view(shape)) * scale.view(shape) + + +class CellposeDense(nn.Module): + """Frozen Cellpose-SAM network returning differentiable dense outputs. + + Reproduces ``core.run_net`` (non-augmented): the image in channel 0 of 3 (the + others zero, as ``convert_image`` does), padded by ``get_pad_yx``, cut into + 256-px tiles at ``tile_overlap=0.1`` and blended with cellpose's taper mask. No + diameter rescale, matching the eval's ``diameter=None``. + + It is loaded on the CPU, in eval mode, and follows the module's ``.to(device)``. + The mode matters: cellpose builds the network in train mode, whose flows differ + from stock inference by up to 1.7 (measured), so ``train()`` keeps it in eval even + if this wrapper is ever a registered submodule. Parameter gradients are the caller's + to switch off; ``SegConsistencyLoss`` does it. + + Parameters + ---------- + model_name : str + File name of a Cellpose-SAM checkpoint in cellpose's model directory, e.g. + ``"cpsam"``. The cpdino family is the evaluation segmenter and is refused. + use_bfloat16 : bool + Load the weights in bfloat16, cellpose's default. + checkpoint_tiles : int or None + When the input needs a gradient, run the network on this many tiles at a time + under activation checkpointing: the backward pass then holds one such group's + activations (~5.7 GB per 256-px tile in bfloat16, measured on an A40) instead of + every tile's, for one extra forward pass. ``None`` keeps all activations. + """ + + BSIZE = 256 + TILE_OVERLAP = 0.1 + + def __init__(self, model_name: str, use_bfloat16: bool = True, checkpoint_tiles: int | None = 1) -> None: + super().__init__() + if checkpoint_tiles is not None and checkpoint_tiles < 1: + raise ValueError(f"checkpoint_tiles must be >= 1 or None, got {checkpoint_tiles}") + if model_name.startswith("cpdino"): + raise ValueError("cpdino is the evaluation segmenter and must not be used as a training segmenter") + path = cp_models.MODEL_DIR / model_name + if not path.is_file(): + raise FileNotFoundError(f"cellpose model {model_name!r} not found at {path}") + # Building the net random-initializes it before loading the checkpoint, which would + # advance the global RNG that HCSDataModule's train/val split draws from later; fork it + # so a segcons arm and its same-seed control split identically. + with torch.random.fork_rng(devices=[]): + cpm = cp_models.CellposeModel( + device=torch.device("cpu"), pretrained_model=str(path), use_bfloat16=use_bfloat16 + ) + if cpm.backbone != "sam_vitl": + raise ValueError(f"{model_name}: expected a SAM backbone, got {cpm.backbone}") + self.net = cpm.net.eval() + self.model_name = model_name + self.checkpoint_tiles = checkpoint_tiles + + def train(self, mode: bool = True) -> "CellposeDense": + """Keep the network in eval mode whatever ``mode`` is.""" + return super().train(False) + + @classmethod + def tiling(cls, ly: int, lx: int) -> tuple[tuple[int, int, int, int], list[tuple[int, int]]]: + """Return cellpose's padding and tile corners for a ``(ly, lx)`` image. + + ``get_pad_yx`` then ``make_tiles`` (non-augmented): the image is padded to a + multiple of 16 of at least ``BSIZE`` plus a margin, so a 240-px image pads to + exactly one 256-px tile while a 256-px one pads to 272 and needs 2 x 2 tiles. + + Returns + ------- + tuple + ``(yp1, yp2, xp1, xp2)`` padding and the ``(y, x)`` corner of every tile. + """ + bs = cls.BSIZE + yp1, yp2, xp1, xp2 = cp_transforms.get_pad_yx(ly, lx, min_size=(bs, bs)) + py, px = ly + yp1 + yp2, lx + xp1 + xp2 + ny = 1 if py <= bs else int(np.ceil((1.0 + 2 * cls.TILE_OVERLAP) * py / bs)) + nx = 1 if px <= bs else int(np.ceil((1.0 + 2 * cls.TILE_OVERLAP) * px / bs)) + ys = np.linspace(0, py - bs, ny).astype(int) + xs = np.linspace(0, px - bs, nx).astype(int) + return (yp1, yp2, xp1, xp2), [(int(a), int(c)) for a in ys for c in xs] + + def _net_flows(self, tiles: Tensor) -> Tensor: + """Run the network on ``(N, 3, 256, 256)`` tiles; return its ``(N, 3, 256, 256)`` output.""" + return self.net(tiles)[0] + + def _taper(self, device: torch.device) -> Tensor: + """Return cellpose's ``transforms._taper_mask`` for one tile.""" + bsize = self.BSIZE + xm = torch.arange(bsize, device=device, dtype=torch.float32) + xm = (xm - xm.mean()).abs() + mask = 1 / (1 + torch.exp((xm - (bsize / 2 - 20)) / 7.5)) + return mask[:, None] * mask[None, :] + + def forward(self, x: Tensor, batch_size: int = 8) -> dict[str, Tensor]: + """Compute the dense outputs of ``(B, 1, Y, X)`` images in any positive-affine units. + + Parameters + ---------- + x : Tensor + Images, shape ``(B, 1, Y, X)``. + batch_size : int + Tiles per network call without a gradient (see ``checkpoint_tiles`` with one). + + Returns + ------- + dict of str to Tensor + ``"flow"`` ``(B, 2, Y, X)`` (dY, dX) and ``"cellprob_logit"`` ``(B, 1, Y, X)``, + float32. + """ + if x.ndim != 4 or x.shape[1] != 1: + raise ValueError(f"expected (B, 1, Y, X), got {tuple(x.shape)}") + return self.dense_normalized(normalize99_diff(x.float()), batch_size=batch_size) + + def dense_normalized(self, xn: Tensor, batch_size: int = 8) -> dict[str, Tensor]: + """Compute the dense outputs of already ``normalize99``-ed images. + + Exposed for finite-difference checks: the normalization statistics are + stop-gradient by design, so only this path has an exact finite difference. + + Parameters + ---------- + xn : Tensor + Normalized images, shape ``(B, 1, Y, X)``. + batch_size : int + Tiles per network call without a gradient. + + Returns + ------- + dict of str to Tensor + As :meth:`forward`. + """ + b, _, ly0, lx0 = xn.shape + net_dtype = next(self.net.parameters()).dtype + xn = torch.cat([xn, torch.zeros_like(xn), torch.zeros_like(xn)], dim=1) + bs = self.BSIZE + pads, coords = self.tiling(ly0, lx0) + yp1, yp2, xp1, xp2 = pads + xn = F.pad(xn, (xp1, xp2, yp1, yp2)) + ly, lx = xn.shape[-2:] + tiles = torch.stack([xn[:, :, a : a + bs, c : c + bs] for a, c in coords], dim=1).reshape(-1, 3, bs, bs) + checkpointed = self.checkpoint_tiles is not None and tiles.requires_grad + step = self.checkpoint_tiles if checkpointed else batch_size + outs = [] + for j in range(0, tiles.shape[0], step): + chunk = tiles[j : j + step].to(net_dtype) + y = checkpoint(self._net_flows, chunk, use_reentrant=False) if checkpointed else self._net_flows(chunk) + outs.append(y.float()) + y = torch.cat(outs).reshape(b, len(coords), 3, bs, bs) + taper = self._taper(xn.device) + acc = xn.new_zeros((b, 3, ly, lx), dtype=torch.float32) + wsum = xn.new_zeros((ly, lx), dtype=torch.float32) + for i, (a, c) in enumerate(coords): + acc = acc + F.pad(y[:, i] * taper, (c, lx - c - bs, a, ly - a - bs)) + wsum = wsum + F.pad(taper, (c, lx - c - bs, a, ly - a - bs)) + yf = (acc / wsum)[:, :, yp1 : ly - yp2, xp1 : lx - xp2] + return {"flow": yf[:, :2], "cellprob_logit": yf[:, 2:3]} diff --git a/applications/dynacell/src/dynacell/engine.py b/applications/dynacell/src/dynacell/engine.py index 8b8425299..1eb75db8c 100644 --- a/applications/dynacell/src/dynacell/engine.py +++ b/applications/dynacell/src/dynacell/engine.py @@ -39,7 +39,7 @@ ) from viscy_models.unet.fcmae import FullyConvolutionalMAE from viscy_utils.log_images import detach_sample, log_image_grid -from viscy_utils.losses import BackgroundLowPass, SegAuxDice +from viscy_utils.losses import BackgroundLowPass, SegAuxDice, SegConsistencyLoss from viscy_utils.optimizers import configure_adamw_scheduler _logger = logging.getLogger("lightning.pytorch") @@ -122,6 +122,31 @@ def _check_seg_aux_args( ) +def _check_seg_consistency_args( + seg_consistency: SegConsistencyLoss | None, weight: float | None, seg_aux: SegAuxDice | None +) -> None: + """Reject a ``seg_consistency`` weight without the term, the term without a weight, and it with ``seg_aux``. + + The weight has no default: it is calibrated per cell (median over validation + batches of ``||grad base|| / ||grad seg_consistency||`` at the starting + checkpoint), so a config that sets the term must also set its weight. + With ``seg_aux`` too, the logged ``base`` would stop meaning the loss before the term. + """ + if seg_consistency is None: + if weight is not None: + raise ValueError(f"seg_consistency_weight={weight} has no effect without seg_consistency.") + return + if weight is None: + raise ValueError( + "seg_consistency is set without seg_consistency_weight, which has no default: calibrate it as the " + "median gradient-norm ratio ||grad base|| / ||grad seg_consistency|| at the starting checkpoint." + ) + if weight <= 0: + raise ValueError(f"seg_consistency_weight must be > 0, got {weight}; drop seg_consistency to turn it off") + if seg_aux is not None: + raise ValueError("seg_consistency cannot be combined with seg_aux; use one auxiliary term per arm.") + + def _require_fg_mask(batch: Sample) -> Tensor: """Return the batch's ``fg_mask`` or raise naming the missing config.""" if "fg_mask" not in batch: @@ -593,6 +618,13 @@ class DynacellUNet(LightningModule): must then carry ``fg_mask`` (set ``fg_mask_key`` on the datamodule). seg_aux_weight : float Weight of the ``seg_aux`` term. Must be ``0.0`` when ``seg_aux`` is None. + seg_consistency : SegConsistencyLoss or None + Segmenter-consistency term on ``(pred, target)``, added as + ``seg_consistency_weight * seg_consistency(pred, target)`` on top of the + rest of the training loss. ``loss/validate`` stays without it; its value + is logged as ``loss/validate_segcons``. Needs no ``fg_mask``. + seg_consistency_weight : float or None + Weight of the ``seg_consistency`` term; required with it, ``None`` without it. channels_last : bool, default False Store the backbone's convolution weights in channels-last layout (``torch.channels_last_3d`` for 5-D weights, ``torch.channels_last`` @@ -633,6 +665,8 @@ def __init__( encoder_only: bool = False, seg_aux: SegAuxDice | None = None, seg_aux_weight: float = 0.0, + seg_consistency: SegConsistencyLoss | None = None, + seg_consistency_weight: float | None = None, channels_last: bool = False, compile_model: bool = False, ) -> None: @@ -641,7 +675,15 @@ def __init__( # resume takes them from its own config and older checkouts can still # resume checkpoints written with them on. self.save_hyperparameters( - ignore=["loss_function", "ckpt_path", "encoder_only", "seg_aux", "channels_last", "compile_model"] + ignore=[ + "loss_function", + "ckpt_path", + "encoder_only", + "seg_aux", + "seg_consistency", + "channels_last", + "compile_model", + ] ) if model_config is None: model_config = {} @@ -662,6 +704,9 @@ def __init__( self.predict_phase_shifts = predict_phase_shifts self.seg_aux = seg_aux self.seg_aux_weight = seg_aux_weight + _check_seg_consistency_args(seg_consistency, seg_consistency_weight, seg_aux) + self.seg_consistency = seg_consistency + self.seg_consistency_weight = seg_consistency_weight self.channels_last = channels_last self.compile_model = compile_model # Fit-scoped compiled backbone; a plain dict keeps it out of the module tree. @@ -671,6 +716,7 @@ def __init__( # Each entry is a list of (loss, batch_size) tuples for weighted aggregation. self.validation_losses: list[list[tuple[Tensor, int]]] = [] self.validation_dice_losses: list[list[tuple[Tensor, int]]] = [] + self.validation_segcons_losses: list[list[tuple[Tensor, int]]] = [] self.validation_step_outputs: list = [] # Cache fg_mask compatibility to avoid per-batch inspect.signature(). @@ -818,6 +864,10 @@ def training_step(self, batch: Sample, batch_idx: int) -> Tensor: base, dice, n_valid = self._compute_seg_aux_loss(pred, target, batch) loss = base + self.seg_aux_weight * dice components = {"base": base, "dice": dice, "dice_n_valid": n_valid} + if self.seg_consistency is not None: + segcons, segcons_components = self.seg_consistency(pred, target, return_components=True) + components = {"base": loss, **components, **segcons_components} + loss = loss + self.seg_consistency_weight * segcons if batch_idx < self.log_batches_per_epoch: self.training_step_outputs.extend(detach_sample((source, target, pred), self.log_samples_per_batch)) self.log( @@ -834,6 +884,7 @@ def training_step(self, batch: Sample, batch_idx: int) -> Tensor: # path. Without these the fused scalar cannot say whether the Dice term # is contributing anything, which is what the lambda_mse choice rests on. # With seg_aux: the base term, the unweighted Dice, and its valid-patch count. + # With seg_consistency: the loss without it, and its unweighted flow/probability terms. for name, value in components.items(): self.log( f"loss/{name}_train", @@ -868,6 +919,9 @@ def validation_step(self, batch: Sample, batch_idx: int, dataloader_idx: int = 0 # the same as the baseline's; the Dice goes to loss/validate_dice. loss, dice, _ = self._compute_seg_aux_loss(pred, target, batch) _record_val_loss(self.validation_dice_losses, dataloader_idx, dice.detach(), source.shape[0]) + if self.seg_consistency is not None: + segcons = self.seg_consistency(pred, target) + _record_val_loss(self.validation_segcons_losses, dataloader_idx, segcons.detach(), source.shape[0]) _record_val_loss(self.validation_losses, dataloader_idx, loss.detach(), source.shape[0]) self.log( f"loss/val/{dataloader_idx}", @@ -928,9 +982,14 @@ def on_validation_epoch_end(self): self.log("loss/validate", _aggregate_validation_losses(self.validation_losses), sync_dist=True) if self.validation_dice_losses: self.log("loss/validate_dice", _aggregate_validation_losses(self.validation_dice_losses), sync_dist=True) + if self.validation_segcons_losses: + self.log( + "loss/validate_segcons", _aggregate_validation_losses(self.validation_segcons_losses), sync_dist=True + ) self.validation_step_outputs.clear() self.validation_losses.clear() self.validation_dice_losses.clear() + self.validation_segcons_losses.clear() def configure_optimizers(self): """Configure AdamW optimizer with LR scheduler.""" @@ -1628,6 +1687,14 @@ class DynacellGAN(LightningModule): Every batch must carry ``fg_mask``. seg_aux_weight : float Weight of the ``seg_aux`` term. Must be ``0.0`` when ``seg_aux`` is None. + seg_consistency : SegConsistencyLoss or None + Segmenter-consistency term on the generator output, added to the + generator objective as ``seg_consistency_weight * seg_consistency(pred, + target)`` (not scaled by ``lambda_l1``). ``loss/validate`` / + ``loss/validate_ema`` stay reconstruction-only; its value is logged as + ``loss/validate_segcons`` / ``loss/validate_ema_segcons``. + seg_consistency_weight : float or None + Weight of the ``seg_consistency`` term; required with it, ``None`` without it. """ def __init__( @@ -1661,6 +1728,8 @@ def __init__( ckpt_path: str | None = None, seg_aux: SegAuxDice | None = None, seg_aux_weight: float = 0.0, + seg_consistency: SegConsistencyLoss | None = None, + seg_consistency_weight: float | None = None, ) -> None: super().__init__() # Lightning's manual-optimization API: required because the GAN @@ -1672,7 +1741,7 @@ def __init__( # subcommands for any key the predict leaf does not declare. A single # pickled class also makes torch.load(weights_only=True) — which # _ckpt_state_dict uses — raise for the whole file. - self.save_hyperparameters(ignore=["ckpt_path", "recon_loss", "seg_aux"]) + self.save_hyperparameters(ignore=["ckpt_path", "recon_loss", "seg_aux", "seg_consistency"]) net_class = _ARCHITECTURE.get(architecture) if net_class is None: @@ -1710,6 +1779,9 @@ def __init__( _check_seg_aux_args(seg_aux, seg_aux_weight, self.recon_loss, self._recon_accepts_fg_mask) self.seg_aux = seg_aux self.seg_aux_weight = seg_aux_weight + _check_seg_consistency_args(seg_consistency, seg_consistency_weight, seg_aux) + self.seg_consistency = seg_consistency + self.seg_consistency_weight = seg_consistency_weight self.lambda_l1 = lambda_l1 self.loss_type = loss_type @@ -1762,6 +1834,8 @@ def __init__( self.validation_losses_ema: list[list[tuple[Tensor, int]]] = [] self.validation_dice_raw: list[list[tuple[Tensor, int]]] = [] self.validation_dice_ema: list[list[tuple[Tensor, int]]] = [] + self.validation_segcons_raw: list[list[tuple[Tensor, int]]] = [] + self.validation_segcons_ema: list[list[tuple[Tensor, int]]] = [] self.validation_step_outputs: list = [] # Build example_input_array for graph logging (TensorBoard/W&B). @@ -2087,6 +2161,9 @@ def training_step(self, batch: Sample, batch_idx: int) -> None: if self.seg_aux is not None: dice_train, dice_n_valid = self._compute_seg_aux(pred, target, batch) g_loss = g_loss + self.seg_aux_weight * dice_train + if self.seg_consistency is not None: + segcons_train, segcons_components = self.seg_consistency(pred, target, return_components=True) + g_loss = g_loss + self.seg_consistency_weight * segcons_train opt_g.zero_grad(set_to_none=True) self.manual_backward(g_loss) opt_g.step() @@ -2137,6 +2214,8 @@ def training_step(self, batch: Sample, batch_idx: int) -> None: if self.seg_aux is not None: log_payload["loss/g_dice_train"] = dice_train log_payload["loss/g_dice_n_valid_train"] = dice_n_valid + if self.seg_consistency is not None: + log_payload.update({f"loss/g_{name}_train": value for name, value in segcons_components.items()}) self.log_dict( log_payload, on_step=True, @@ -2192,6 +2271,9 @@ def validation_step(self, batch: Sample, batch_idx: int, dataloader_idx: int = 0 if self.seg_aux is not None: dice_raw, _ = self._compute_seg_aux(pred_raw, target, batch) _record_val_loss(self.validation_dice_raw, dataloader_idx, dice_raw.detach(), source.shape[0]) + if self.seg_consistency is not None: + segcons_raw = self.seg_consistency(pred_raw, target) + _record_val_loss(self.validation_segcons_raw, dataloader_idx, segcons_raw.detach(), source.shape[0]) self.log( f"loss/val/{dataloader_idx}", recon_raw, @@ -2208,6 +2290,9 @@ def validation_step(self, batch: Sample, batch_idx: int, dataloader_idx: int = 0 if self.seg_aux is not None: dice_ema, _ = self._compute_seg_aux(pred_ema, target, batch) _record_val_loss(self.validation_dice_ema, dataloader_idx, dice_ema.detach(), source.shape[0]) + if self.seg_consistency is not None: + segcons_ema = self.seg_consistency(pred_ema, target) + _record_val_loss(self.validation_segcons_ema, dataloader_idx, segcons_ema.detach(), source.shape[0]) self.log( f"loss/val_ema/{dataloader_idx}", recon_ema, @@ -2256,11 +2341,19 @@ def on_validation_epoch_end(self) -> None: self.log("loss/validate_dice", _aggregate_validation_losses(self.validation_dice_raw), sync_dist=True) if self.validation_dice_ema: self.log("loss/validate_ema_dice", _aggregate_validation_losses(self.validation_dice_ema), sync_dist=True) + if self.validation_segcons_raw: + self.log("loss/validate_segcons", _aggregate_validation_losses(self.validation_segcons_raw), sync_dist=True) + if self.validation_segcons_ema: + self.log( + "loss/validate_ema_segcons", _aggregate_validation_losses(self.validation_segcons_ema), sync_dist=True + ) self.validation_step_outputs.clear() self.validation_losses_raw.clear() self.validation_losses_ema.clear() self.validation_dice_raw.clear() self.validation_dice_ema.clear() + self.validation_segcons_raw.clear() + self.validation_segcons_ema.clear() def configure_optimizers(self): """Build two AdamW optimizers + WarmupCosine schedulers via the shared helper.""" diff --git a/applications/dynacell/src/dynacell/evaluation/paths.py b/applications/dynacell/src/dynacell/evaluation/paths.py index 68bee4dd4..5be73a3a0 100644 --- a/applications/dynacell/src/dynacell/evaluation/paths.py +++ b/applications/dynacell/src/dynacell/evaluation/paths.py @@ -350,6 +350,15 @@ # Checkpoint-stage control: the nucleus vscyto3daug segaux_halfw arm re-predicted from # latest-epoch=76-step=200000 instead of its best-val epoch=51, at the e74 stage-control floor. "fnet3d_vscyto3daug_segaux_halfw_e76": "fnet3d_vscyto3daug_segaux_halfw_e76", + # P1 segmenter-consistency nucleus fine-tunes (condgen review, 2026-10-15): `_segcons` forks + # the converged baseline and fine-tunes it with SegConsistencyLoss (Cellpose-SAM); + # `_segcons_ctrl` is the same fork, steps and LR without the term. + "fnet3d_vscyto3daug_segcons": "fnet3d_vscyto3daug_segcons", + "fnet3d_vscyto3daug_segcons_ctrl": "fnet3d_vscyto3daug_segcons_ctrl", + "fcmae_vscyto3d_scratch_segcons": "unext2_segcons", + "fcmae_vscyto3d_scratch_segcons_ctrl": "unext2_segcons_ctrl", + "pix2pix2d_unetvit_segcons": "pix2pix2d_segcons", + "pix2pix2d_unetvit_segcons_ctrl": "pix2pix2d_segcons_ctrl", } # Organelle code token -> paper-script organelle key. Mito uses the long form diff --git a/applications/dynacell/tests/test_benchmark_config_composition.py b/applications/dynacell/tests/test_benchmark_config_composition.py index 1b58fe2a8..6afb760af 100644 --- a/applications/dynacell/tests/test_benchmark_config_composition.py +++ b/applications/dynacell/tests/test_benchmark_config_composition.py @@ -1126,6 +1126,13 @@ def test_joint_train_smoke_leaf_composes() -> None: for organelle in ("nucleus", "membrane") for suffix in ("segaux", "segauxself", "seed1") ) + | frozenset( + # P1 fine-tunes of the same baseline (nucleus). The control changes only the budget, LR and + # starting weights. The segcons arm adds cpsam: ~1.2 GB of bf16 weights plus ~0.07 GB per + # sample with per-tile checkpointing (A40 bench, 2026-10-15), so ~15 GiB/rank at bs=4. + f"nucleus/pix2pix2d_unetvit_{suffix}/ipsc_confocal/train.yml" + for suffix in ("segcons", "segcons_ctrl") + ) ) _WIDE_GPU_EXTRA = frozenset({"a40", "a6000", "l40s"}) diff --git a/applications/dynacell/tests/test_cellpose_dense.py b/applications/dynacell/tests/test_cellpose_dense.py new file mode 100644 index 000000000..aebb75c44 --- /dev/null +++ b/applications/dynacell/tests/test_cellpose_dense.py @@ -0,0 +1,187 @@ +"""CellposeDense (differentiable Cellpose-SAM) against stock cellpose, and as a SegConsistencyLoss segmenter. + +The network is ViT-L: a 256-px tile takes ~27 s forward on 6 CPU cores (measured), so the +tests that run it need a GPU and the stock ``cpsam`` checkpoint in cellpose's model dir. +""" + +import numpy as np +import pytest +import torch + +pytest.importorskip("cellpose") + +from cellpose import models as cp_models # noqa: E402 +from cellpose import transforms as cp_transforms # noqa: E402 +from generate_spotlight_v2_leaves import SEG_CONSISTENCY_ARGS # noqa: E402 + +from dynacell.cellpose_dense import CellposeDense, normalize99_diff # noqa: E402 +from dynacell.engine import DynacellUNet # noqa: E402 +from viscy_utils.losses import SegConsistencyLoss # noqa: E402 + +needs_cpsam = pytest.mark.skipif( + not torch.cuda.is_available() or not (cp_models.MODEL_DIR / "cpsam").is_file(), + reason="needs a GPU and the cpsam checkpoint in cellpose's model directory", +) + + +def _nuclei_image(shape: tuple[int, int], seed: int = 0) -> np.ndarray: + """Bright disks of varying intensity on noise, in arbitrary units.""" + rng = np.random.default_rng(seed) + yy, xx = np.mgrid[: shape[0], : shape[1]] + img = np.zeros(shape, np.float32) + for _ in range(14): + cy, cx, r = rng.uniform(15, shape[0] - 15), rng.uniform(15, shape[1] - 15), rng.uniform(8, 16) + img += ((yy - cy) ** 2 + (xx - cx) ** 2 < r**2) * rng.uniform(0.5, 1.0) + return (300.0 * img + 100.0 + rng.normal(0, 10.0, shape)).astype(np.float32) + + +@pytest.fixture(scope="module") +def cpsam_bf16() -> CellposeDense: + return CellposeDense("cpsam").cuda() + + +def test_configured_crop_is_one_cellpose_tile(): + """The segcons leaves' crop pads to exactly one tile; a 256-px crop would pad to 2 x 2.""" + crop = SEG_CONSISTENCY_ARGS["crop_size"] + _, coords = CellposeDense.tiling(crop, crop) + assert len(coords) == 1 + assert len(CellposeDense.tiling(256, 256)[1]) == 4 + + +@needs_cpsam +def test_train_keeps_the_network_in_eval(cpsam_bf16): + """Even as a registered submodule, a parent's train() leaves cellpose's network in eval mode.""" + parent = torch.nn.ModuleDict({"seg": cpsam_bf16}) + parent.train() + assert not cpsam_bf16.training and not any(m.training for m in cpsam_bf16.net.modules()) + + +@pytest.mark.skipif(not (cp_models.MODEL_DIR / "cpsam").is_file(), reason="needs the cpsam checkpoint") +def test_construction_leaves_the_global_torch_rng_untouched(): + """Building the segmenter must not shift the seeded torch RNG the datamodule's split draws from. + + ``HCSDataModule`` splits train/val FOVs with ``torch.randperm``; otherwise a segcons arm and + its same-seed control would train on different splits. + """ + torch.manual_seed(42) + reference = torch.randperm(50) + torch.manual_seed(42) + CellposeDense("cpsam") + assert torch.equal(reference, torch.randperm(50)) + + +def test_refuses_the_evaluation_segmenter_and_missing_models(): + with pytest.raises(ValueError, match="evaluation segmenter"): + CellposeDense("cpdino") + with pytest.raises(FileNotFoundError, match="not found"): + CellposeDense("no_such_cellpose_model") + with pytest.raises(ValueError, match="checkpoint_tiles"): + CellposeDense("cpsam", checkpoint_tiles=0) + + +def test_normalization_matches_cellpose_normalize99(): + img = _nuclei_image((96, 80)) + ours = normalize99_diff(torch.from_numpy(img)[None, None])[0, 0].numpy() + np.testing.assert_allclose(ours, cp_transforms.normalize99(img), atol=1e-5) + flat = normalize99_diff(torch.full((1, 1, 8, 8), 3.0)) + assert torch.equal(flat, torch.zeros_like(flat)) + + +@needs_cpsam +def test_dense_outputs_match_stock_cellpose(cpsam_bf16): + """Same flows and cell-probability logits as ``CellposeModel.eval`` on the same normalized input. + + Two overlapping tiles exercise the taper blend. The input is normalized once and fed to + both: in bfloat16 the ~1e-6 float differences between the two normalize99 implementations + move single outputs by up to 0.2 (measured), which says nothing about the tiling path. + """ + img = cp_transforms.normalize99(_nuclei_image((256, 300))).astype(np.float32) + with torch.no_grad(): + out = cpsam_bf16.dense_normalized(torch.from_numpy(img)[None, None].cuda()) + stock = cp_models.CellposeModel(gpu=True, pretrained_model=str(cp_models.MODEL_DIR / "cpsam")) + _, flows, _ = stock.eval(img, normalize=False, diameter=None, channel_axis=None, compute_masks=False) + np.testing.assert_allclose(out["flow"][0].cpu().numpy(), flows[1], atol=1e-5) + np.testing.assert_allclose(out["cellprob_logit"][0, 0].cpu().numpy(), flows[2], atol=1e-5) + assert np.abs(flows[1]).mean() > 0.05 # the image has cells, so the parity is not 0 == 0 + + +@needs_cpsam +def test_input_gradient_matches_finite_difference(): + """Central difference on the post-normalization path (statistics are stop-gradient), TF32 off.""" + seg = CellposeDense("cpsam", use_bfloat16=False).cuda() + tf32 = (torch.backends.cudnn.allow_tf32, torch.backends.cuda.matmul.allow_tf32) + torch.backends.cudnn.allow_tf32 = torch.backends.cuda.matmul.allow_tf32 = False + try: + g = torch.Generator(device="cuda").manual_seed(0) + x = torch.rand(1, 1, 64, 64, device="cuda", generator=g).requires_grad_(True) + out = seg.dense_normalized(x) + w = {k: torch.randn(out[k].shape, device="cuda", generator=g) for k in out} + + def f(z: torch.Tensor) -> torch.Tensor: + o = seg.dense_normalized(z) + return sum((o[k].double() * w[k]).sum() for k in o) + + f(x).backward() + v = torch.randn(x.shape, device="cuda", generator=g) + v = v / v.norm() * x.detach().norm() * 1e-2 + with torch.no_grad(): + num = float((f(x.detach() + v) - f(x.detach() - v)) / 2) + ana = float((x.grad * v).sum()) + finally: + torch.backends.cudnn.allow_tf32, torch.backends.cuda.matmul.allow_tf32 = tf32 + assert torch.isfinite(x.grad).all() + assert abs(num - ana) / abs(num) < 1e-2 + + +@needs_cpsam +def test_activation_checkpointing_keeps_outputs_and_gradient(): + """Checkpointed single tiles equal a stored pass over the same single tiles (2 tiles). + + Same grouping on both sides: regrouping tiles alone moves the input gradient by 5e-4 + (TF32) / 1.2e-5 (fp32) relative norm (measured), from the batched kernels, not checkpointing. + """ + img = torch.from_numpy(_nuclei_image((64, 300)))[None, None].cuda() + results = [] + for tiles in (1, None): + seg = CellposeDense("cpsam", use_bfloat16=False, checkpoint_tiles=tiles).cuda() + x = img.clone().requires_grad_(True) + out = seg(x, batch_size=1) + (out["flow"].square().sum() + out["cellprob_logit"].sum()).backward() + results.append((out["flow"].detach(), x.grad)) + del seg + torch.testing.assert_close(results[0][0], results[1][0], rtol=0, atol=0) + torch.testing.assert_close(results[0][1], results[1][1], rtol=0, atol=0) + assert results[0][1].abs().sum() > 0 + + +@needs_cpsam +def test_seg_consistency_through_cpsam_trains_a_backbone(cpsam_bf16): + """Zero at pred == target; one DynacellUNet step under bf16 autocast puts its gradient on the backbone.""" + target = torch.from_numpy(np.stack([_nuclei_image((64, 64), seed=s) for s in range(4)])).reshape(2, 1, 2, 64, 64) + target = target.repeat_interleave(2, dim=2).cuda() # (2, 1, 4, 64, 64) + loss_fn = SegConsistencyLoss(cpsam_bf16, crop_size=64) + # The loss freezes its segmenter, and the LightningModule's train() cannot reach it. + assert not cpsam_bf16.training and not any(p.requires_grad for p in cpsam_bf16.parameters()) + assert float(loss_fn(target.clone(), target)) == pytest.approx(0.0, abs=1e-6) + + torch.manual_seed(0) + module = DynacellUNet( + architecture="FNet3D", + model_config={"in_channels": 1, "out_channels": 1, "depth": 1, "mult_chan": 8, "in_stack_depth": 4}, + seg_consistency=loss_fn, + seg_consistency_weight=1.0, + log_batches_per_epoch=0, + ).cuda() + module.train() + assert not cpsam_bf16.training + module.log = lambda *args, **kwargs: None + batch = {"source": torch.randn_like(target), "target": target} + with torch.autocast("cuda", dtype=torch.bfloat16): + pred = module(batch["source"]) + segcons = loss_fn(pred, batch["target"]) + assert float(segcons.detach()) > 0 + segcons.backward() + grads = [p.grad for p in module.model.parameters()] + assert all(g is not None and torch.isfinite(g).all() for g in grads) + assert sum(float(g.abs().sum()) for g in grads) > 0 + assert all(p.grad is None for p in cpsam_bf16.parameters()) diff --git a/applications/dynacell/tests/test_paths.py b/applications/dynacell/tests/test_paths.py index 16db143b0..8f285ef5c 100644 --- a/applications/dynacell/tests/test_paths.py +++ b/applications/dynacell/tests/test_paths.py @@ -99,6 +99,12 @@ "fcmae_vscyto3d_scratch_segaux_cpout", "celldiff_2d_bglp", "celldiff_2d_bgflat", + "fnet3d_vscyto3daug_segcons", + "fnet3d_vscyto3daug_segcons_ctrl", + "fcmae_vscyto3d_scratch_segcons", + "fcmae_vscyto3d_scratch_segcons_ctrl", + "pix2pix2d_unetvit_segcons", + "pix2pix2d_unetvit_segcons_ctrl", ) # Spotlight-v2 post-hoc control stores (no training leaf): each 3D nucleus baseline # histogram-matched to its segaux arm. Same longest-prefix hazard as the arms. diff --git a/applications/dynacell/tests/test_seg_consistency_focus.py b/applications/dynacell/tests/test_seg_consistency_focus.py new file mode 100644 index 000000000..57ea23d40 --- /dev/null +++ b/applications/dynacell/tests/test_seg_consistency_focus.py @@ -0,0 +1,37 @@ +"""The torch focus-plane port in SegConsistencyLoss picks the plane and slab the eval segments.""" + +import numpy as np +import torch + +from dynacell.evaluation.focus import nucleus_area_plane, slab_mip +from viscy_utils.losses.seg_consistency import _slab_mip, nucleus_area_planes + + +def _blob_volume(rng: np.random.Generator, z: int) -> np.ndarray: + """Disks peaked at a random plane, on noise, at a random scale and offset.""" + yy, xx = np.mgrid[:48, :48] + zz = np.arange(z)[:, None, None] + vol = np.zeros((z, 48, 48), np.float32) + for _ in range(int(rng.integers(1, 6))): + cy, cx, r = rng.uniform(6, 42), rng.uniform(6, 42), rng.uniform(3, 9) + prof = np.exp(-(((zz - rng.uniform(0, z)) / rng.uniform(0.7, 6.0)) ** 2)) + vol += prof * ((yy - cy) ** 2 + (xx - cx) ** 2 < r**2) + vol += rng.normal(0, rng.uniform(0.01, 0.2), vol.shape) + return (vol * rng.uniform(0.1, 500) + rng.uniform(-10, 10)).astype(np.float32) + + +def test_focus_plane_and_slab_match_the_eval(): + rng = np.random.default_rng(0) + vols = [_blob_volume(rng, int(z)) for z in rng.integers(1, 41, size=48)] + vols += [np.full((7, 48, 48), 3.0, np.float32), _blob_volume(rng, 1)] # constant -> center; Z = 1 + planes_seen = set() + for vol in vols: + expected = nucleus_area_plane(vol) + got = int(nucleus_area_planes(torch.from_numpy(vol)[None])[0]) + assert got == expected, (vol.shape, got, expected) + planes_seen.add((vol.shape[0], got)) + for halfwidth in (0, 1, 2): + ours = _slab_mip(torch.from_numpy(vol)[None], torch.tensor([got]), halfwidth)[0].numpy() + np.testing.assert_array_equal(ours, slab_mip(vol[None], [expected], halfwidth)[0]) + # The volumes exercise more than the center fallback. + assert sum(p != z // 2 for z, p in planes_seen) >= 10 diff --git a/applications/dynacell/tests/test_seg_consistency_training.py b/applications/dynacell/tests/test_seg_consistency_training.py new file mode 100644 index 000000000..cebd254ec --- /dev/null +++ b/applications/dynacell/tests/test_seg_consistency_training.py @@ -0,0 +1,237 @@ +"""SegConsistencyLoss wired into DynacellUNet (FNet3D, UNeXt2-3D) and DynacellGAN (pix2pix2d geometry). + +The engine tests use a tiny stand-in segmenter; ``test_cellpose_dense.py`` covers the +Cellpose-SAM one. +""" + +from pathlib import Path + +import pytest +import torch +from lightning.pytorch import Trainer, seed_everything +from lightning.pytorch.loggers import TensorBoardLogger +from torch import Tensor, nn + +from dynacell.engine import DynacellGAN, DynacellUNet +from viscy_data.hcs import HCSDataModule +from viscy_utils.losses import SegAuxDice, SegConsistencyLoss + +FNET_CONFIG = {"in_channels": 1, "out_channels": 1, "depth": 1, "mult_chan": 8, "in_stack_depth": 4} +FCMAE3D_CONFIG = { + "in_channels": 1, + "out_channels": 1, + "encoder_blocks": [1, 1, 1, 1], + "dims": [8, 16, 32, 64], + "decoder_conv_blocks": 1, + "stem_kernel_size": [5, 4, 4], + "in_stack_depth": 5, + "pretraining": False, +} +# pix2pix2d geometry: a Z=1 UNetViT generator and a Z-preserving discriminator. +VIT2D_CONFIG = { + "input_spatial_size": [1, 32, 32], + "in_channels": 1, + "out_channels": 1, + "dims": [32, 64, 128], + "num_res_block": [2, 2], + "hidden_size": 64, + "num_heads": 4, + "dim_head": 16, + "num_hidden_layers": 1, + "patch_size": [1, 4, 4], +} +DISC2D_CONFIG = {"in_channels": 2, "base_channels": 8, "num_scales": 1, "preserve_z": True} +# (architecture, model_config, z_window_size) +UNET_CASES = {"fnet3d": ("FNet3D", FNET_CONFIG, 4), "unext2_3d": ("fcmae", FCMAE3D_CONFIG, 5)} + + +class TinySegmenter(nn.Module): + """Two conv layers to flow (2 ch) + cell-probability logit (1 ch), per-image min-max normalized.""" + + def __init__(self) -> None: + super().__init__() + self.net = nn.Sequential(nn.Conv2d(1, 8, 3, padding=1), nn.Tanh(), nn.Conv2d(8, 3, 3, padding=1)) + + def forward(self, x: Tensor) -> dict[str, Tensor]: + with torch.no_grad(): + lo = x.amin((1, 2, 3), keepdim=True) + hi = x.amax((1, 2, 3), keepdim=True) + y = self.net((x - lo) / (hi - lo + 1e-6)) + return {"flow": y[:, :2], "cellprob_logit": y[:, 2:3]} + + +def _segcons() -> SegConsistencyLoss: + return SegConsistencyLoss(TinySegmenter()) + + +def _datamodule(zarr_path: Path, z_window_size: int) -> HCSDataModule: + return HCSDataModule( + data_path=str(zarr_path), + source_channel="Phase3D", + target_channel="Fluorescence", + z_window_size=z_window_size, + batch_size=2, + num_workers=0, + split_ratio=0.5, + yx_patch_size=(32, 32), + ) + + +def _cpu_trainer(tmp_path: Path | None = None, **kwargs) -> Trainer: + return Trainer( + accelerator="cpu", + logger=TensorBoardLogger(save_dir=tmp_path) if tmp_path is not None else False, + enable_checkpointing=False, + enable_progress_bar=False, + enable_model_summary=False, + **kwargs, + ) + + +def _assert_finite(metrics: dict, keys: list[str]) -> None: + for key in keys: + assert key in metrics, f"{key} not logged; got {sorted(metrics)}" + assert torch.isfinite(metrics[key]), f"{key}={metrics[key]}" + + +@pytest.mark.parametrize("case", list(UNET_CASES)) +def test_unet_seg_consistency_fast_dev_run(tmp_path, tiny_hcs_zarr, case): + architecture, model_config, z_window_size = UNET_CASES[case] + seed_everything(42) + module = DynacellUNet( + architecture=architecture, + model_config=model_config, + seg_consistency=_segcons(), + seg_consistency_weight=0.5, + log_batches_per_epoch=0, + ) + trainer = _cpu_trainer(tmp_path, fast_dev_run=True) + trainer.fit(module, datamodule=_datamodule(tiny_hcs_zarr, z_window_size)) + assert trainer.state.status == "finished" + metrics = trainer.callback_metrics + keys = ["loss/train", "loss/base_train", "loss/segcons_train", "loss/segcons_flow_train", "loss/segcons_prob_train"] + _assert_finite(metrics, [*keys, "loss/validate", "loss/validate_segcons"]) + assert metrics["loss/segcons_train"] > 0 + expected = metrics["loss/base_train"] + 0.5 * metrics["loss/segcons_train"] + assert float(metrics["loss/train"]) == pytest.approx(float(expected), rel=1e-5) + + +def test_unet_training_step_gradient_includes_the_term(): + """The step's backbone gradient is the base gradient plus weight x the term's own gradient.""" + seed_everything(0) + batch = {"source": torch.randn(2, 1, 4, 32, 32), "target": torch.randn(2, 1, 4, 32, 32)} + term = _segcons() + baseline = DynacellUNet(architecture="FNet3D", model_config=FNET_CONFIG, log_batches_per_epoch=0) + arm = DynacellUNet( + architecture="FNet3D", + model_config=FNET_CONFIG, + seg_consistency=term, + seg_consistency_weight=2.0, + log_batches_per_epoch=0, + ) + arm.load_state_dict(baseline.state_dict()) + + def grad(module: DynacellUNet, fn) -> Tensor: + params = list(module.model.parameters()) + return torch.cat([g.flatten() for g in torch.autograd.grad(fn(), params)]) + + for module in (baseline, arm): + module.log = lambda *args, **kw: None + g_arm = grad(arm, lambda: arm.training_step(batch, 0)) + g_base = grad(baseline, lambda: baseline.training_step(batch, 0)) + g_term = grad(arm, lambda: term(arm(batch["source"]), batch["target"])) + assert g_term.abs().sum() > 0 + torch.testing.assert_close(g_arm, g_base + 2.0 * g_term) + + +def test_unet_validate_excludes_the_term(tiny_hcs_zarr): + """``loss/validate`` matches the baseline's on identical weights, so ``--ckpt best`` keeps its meaning.""" + seed_everything(0) + baseline = DynacellUNet(architecture="FNet3D", model_config=FNET_CONFIG, log_batches_per_epoch=0) + arm = DynacellUNet( + architecture="FNet3D", + model_config=FNET_CONFIG, + seg_consistency=_segcons(), + seg_consistency_weight=5.0, + log_batches_per_epoch=0, + ) + arm.load_state_dict(baseline.state_dict()) + seed_everything(7) + base_metrics = _cpu_trainer().validate(baseline, datamodule=_datamodule(tiny_hcs_zarr, 4))[0] + seed_everything(7) + arm_metrics = _cpu_trainer().validate(arm, datamodule=_datamodule(tiny_hcs_zarr, 4))[0] + assert "loss/validate_segcons" not in base_metrics + assert arm_metrics["loss/validate_segcons"] > 0.0 + assert arm_metrics["loss/validate"] == pytest.approx(base_metrics["loss/validate"], rel=1e-6) + + +def test_gan_seg_consistency_fast_dev_run(tmp_path, tiny_hcs_zarr): + """pix2pix2d geometry (Z=1) with EMA: the term trains and is logged apart from the recon monitors.""" + seed_everything(42) + module = DynacellGAN( + architecture="UNetViT3D", + generator_config=VIT2D_CONFIG, + discriminator_config=DISC2D_CONFIG, + ema_kimg=10.0, + seg_consistency=_segcons(), + seg_consistency_weight=0.5, + log_batches_per_epoch=0, + ) + trainer = _cpu_trainer(tmp_path, fast_dev_run=True) + trainer.fit(module, datamodule=_datamodule(tiny_hcs_zarr, 1)) + assert trainer.state.status == "finished" + _assert_finite( + trainer.callback_metrics, + [ + "loss/g_train", + "loss/g_l1_train", + "loss/g_segcons_train", + "loss/g_segcons_flow_train", + "loss/g_segcons_prob_train", + "loss/validate", + "loss/validate_ema", + "loss/validate_segcons", + "loss/validate_ema_segcons", + ], + ) + assert trainer.callback_metrics["loss/g_segcons_train"] > 0 + + +@pytest.mark.parametrize("family", ["unet", "gan"]) +def test_weight_is_required_with_the_term_and_seg_aux_is_rejected_with_it(family): + if family == "unet": + kwargs = {"architecture": "FNet3D", "model_config": FNET_CONFIG} + cls = DynacellUNet + else: + kwargs = {"generator_config": VIT2D_CONFIG, "discriminator_config": DISC2D_CONFIG} + cls = DynacellGAN + with pytest.raises(ValueError, match="without seg_consistency_weight"): + cls(**kwargs, seg_consistency=_segcons(), seg_consistency_weight=None) + with pytest.raises(ValueError, match="no effect without seg_consistency"): + cls(**kwargs, seg_consistency_weight=1.0) + for weight in (0.0, -1.0): + with pytest.raises(ValueError, match="> 0"): + cls(**kwargs, seg_consistency=_segcons(), seg_consistency_weight=weight) + with pytest.raises(ValueError, match="cannot be combined with seg_aux"): + cls(**kwargs, seg_consistency=_segcons(), seg_consistency_weight=1.0, seg_aux=SegAuxDice(), seg_aux_weight=1.0) + + +@pytest.mark.parametrize("family", ["unet", "gan"]) +def test_fork_loads_a_baseline_checkpoint_strictly_and_saves_no_segmenter(tmp_path, family): + """A fine-tune arm starts from a baseline checkpoint via ``ckpt_path`` and writes the same keys.""" + if family == "unet": + kwargs = {"architecture": "FNet3D", "model_config": FNET_CONFIG} + cls = DynacellUNet + else: + kwargs = {"generator_config": VIT2D_CONFIG, "discriminator_config": DISC2D_CONFIG, "ema_kimg": 10.0} + cls = DynacellGAN + seed_everything(0) + baseline = cls(**kwargs) + ckpt = tmp_path / "baseline.ckpt" + torch.save({"state_dict": baseline.state_dict()}, ckpt) + arm = cls(**kwargs, ckpt_path=str(ckpt), seg_consistency=_segcons(), seg_consistency_weight=1.0) + assert arm.state_dict().keys() == baseline.state_dict().keys() + for key, value in baseline.state_dict().items(): + torch.testing.assert_close(arm.state_dict()[key], value) + assert "seg_consistency" not in arm.hparams + assert arm.hparams["seg_consistency_weight"] == 1.0 diff --git a/applications/dynacell/tools/generate_grouped_eval_configs.py b/applications/dynacell/tools/generate_grouped_eval_configs.py index 231e63269..13504a48d 100644 --- a/applications/dynacell/tools/generate_grouped_eval_configs.py +++ b/applications/dynacell/tools/generate_grouped_eval_configs.py @@ -139,6 +139,12 @@ "fnet3d_paper_histmatch": "fnet3d_histmatch", "pix2pix3d_unetvit_v2_histmatch": "pix2pix3d_v2_histmatch", "fcmae_vscyto3d_scratch_v2_histmatch": "unext2_v2_histmatch", + "fnet3d_vscyto3daug_segcons": "fnet3d_vscyto3daug_segcons", + "fnet3d_vscyto3daug_segcons_ctrl": "fnet3d_vscyto3daug_segcons_ctrl", + "fcmae_vscyto3d_scratch_segcons": "unext2_segcons", + "fcmae_vscyto3d_scratch_segcons_ctrl": "unext2_segcons_ctrl", + "pix2pix2d_unetvit_segcons": "pix2pix2d_segcons", + "pix2pix2d_unetvit_segcons_ctrl": "pix2pix2d_segcons_ctrl", } # iPSC: target key in aics-hipsc manifest. @@ -242,6 +248,12 @@ "fnet3d_paper_histmatch", "pix2pix3d_unetvit_v2_histmatch", "fcmae_vscyto3d_scratch_v2_histmatch", + "fnet3d_vscyto3daug_segcons", + "fnet3d_vscyto3daug_segcons_ctrl", + "fcmae_vscyto3d_scratch_segcons", + "fcmae_vscyto3d_scratch_segcons_ctrl", + "pix2pix2d_unetvit_segcons", + "pix2pix2d_unetvit_segcons_ctrl", ) _CELLDIFF_MODELS: tuple[str, ...] = ( "celldiff_2d_segauxself", diff --git a/applications/dynacell/tools/generate_spotlight_v2_leaves.py b/applications/dynacell/tools/generate_spotlight_v2_leaves.py index c26ac241c..c49d09571 100644 --- a/applications/dynacell/tools/generate_spotlight_v2_leaves.py +++ b/applications/dynacell/tools/generate_spotlight_v2_leaves.py @@ -80,6 +80,18 @@ illumination removed too). The background stays supervised, unlike a masked loss. Validation stays raw, so ``--ckpt best`` selects on the baseline's criterion; predict leaves equal the baseline's. Normalization is the baseline's. +- ``segcons`` / ``segcons_ctrl`` -- P1 of the conditional-generation review + (``experiments/2026-10-15_conditional-generation-review/PROPOSALS.md``), nucleus only, on + FNet-3D vscyto3daug, UNeXt2-3D and pix2pix2d. Both are fine-tunes ("forks") of the + baseline recipe's converged final checkpoint (:data:`FORK_CKPTS`, loaded weights-only + through ``model.init_args.ckpt_path``, so optimizer and schedule start fresh) on a + shortened budget at a lower LR (:data:`FORK_SCHEDULES`). ``segcons`` adds + ``SegConsistencyLoss`` with Cellpose-SAM (``cpsam``) as the frozen segmenter; + ``segcons_ctrl`` is the same fork, steps, LR and seed without it, so the pair is + stage-matched. ``seg_consistency_weight`` is calibrated per arm at the fork checkpoint + (:data:`SEG_CONSISTENCY_WEIGHTS`). Both + predict from ``last.ckpt`` (``--ckpt last``): ``loss/validate`` stays base-only, so a + best-by-val pick would favour the term arm's earliest epoch. Leaves are emitted with ``yaml.safe_dump``, so they carry no inline comments: the recipe rationale stays in the baseline leaf each header names, and the reasons for @@ -227,6 +239,69 @@ ) _WALL_4GPU = "launcher_profiles/hardware_4gpu.yml" _WALL_4GPU_LONG = "launcher_profiles/hardware_4gpu_long.yml" +# P1 segmenter-consistency fine-tunes (2026-10-15): three converged nucleus baselines, each +# forked twice, with and without the term. The fork source is the baseline recipe's final +# checkpoint under its immutable latest-* name (all three runs completed their budget); the +# FNet-3D and UNeXt2-3D forks start from the v2 retrains, the Spotlight-v2 controls on today's code. +SEGCONS_BASELINES: tuple[str, ...] = ("fnet3d_vscyto3daug", "fcmae_vscyto3d_scratch", "pix2pix2d_unetvit") +FORK_CKPTS: dict[tuple[str, str], str] = { + ("nucleus", "fnet3d_vscyto3daug"): f"{MODELS_ROOT}/ipsc/nucleus/fnet3d_vscyto3daug_v2/checkpoints/" + "latest-epoch=76-step=200000.ckpt", + ("nucleus", "fcmae_vscyto3d_scratch"): f"{MODELS_ROOT}/ipsc/nucleus/fcmae_vscyto3d_scratch_v2/checkpoints/" + "latest-epoch=199-step=62400.ckpt", + ("nucleus", "pix2pix2d_unetvit"): f"{MODELS_ROOT}/ipsc/nucleus/pix2pix2d_unetvit/checkpoints/" + "latest-epoch=39-step=20000.ckpt", +} +# Fine-tune budget: 10% of the baseline's (200000 steps, 200 epochs, 40 epochs, read off the fork +# checkpoints' names) at 0.1x its peak LR (FNet-3D keeps its Constant schedule). The WarmupCosine +# families warm up over 5% of the fine-tune's scheduler steps: UNeXt2-3D 20 x 312 = 6240 (62400 +# steps / 200 epochs), pix2pix2d 4 x 250 = 1000 (20000 global steps / 40 epochs, two optimizer +# steps per GAN batch). The scheduler is rebuilt for the fine-tune, so a cosine decays to 0 again. +# pix2pix2d predicts from generator_ema, which starts as the fork's own shadow. At the baseline's +# ema_kimg 10 the fine-tune (1000 G steps x 16 samples: 4 ranks x batch 4) would leave +# 0.5 ** (16000 / 10000) = 33% of the fork-start shadow in it, diluting both arms; ema_kimg 2 +# leaves 0.5 ** (16000 / 2000) = 0.5 ** 8 = 0.4% (decay = 0.5 ** (global batch / (ema_kimg * 1000))). +FORK_SCHEDULES: dict[str, dict[str, dict[str, Any]]] = { + "fnet3d_vscyto3daug": {"trainer": {"max_steps": 20000}, "model": {"lr": 1.0e-4}}, + "fcmae_vscyto3d_scratch": {"trainer": {"max_epochs": 20}, "model": {"lr": 4.0e-5, "warmup_steps": 312}}, + "pix2pix2d_unetvit": { + "trainer": {"max_epochs": 4}, + "model": {"lr_g": 2.0e-5, "lr_d": 2.0e-5, "warmup_steps": 50, "ema_kimg": 2.0}, + }, +} +# S_train from gate G4: cpsam (Kendall tau vs cpdino nucleus Dice/mAP .65/.80, not in the eval +# stack). One slab-MIP pass per sample around the target crop's nuclear-area plane, center +# 240-px crop, which cellpose pads to exactly one 256-px tile, the geometry the G4 bench timed +# (a 256-px crop pads to 272 and runs 2 x 2 tiles, 4x the segmenter compute, measured with +# get_pad_yx); see viscy_utils.losses.seg_consistency. +# Every sample of the batch is scored. The with-gradient pass is checkpointed per tile: without +# it each sample holds ~5.6 GB of activations (G4 H200 bench: 42.9 GB at 8 samples, against the +# FNet-3D fit's ~41 GB), with it ~0.07 GB per sample at ~1.8x the term's time (A40, 2026-10-15). +SEG_CONSISTENCY_ARGS: dict[str, Any] = { + "segmenter": { + "class_path": "dynacell.cellpose_dense.CellposeDense", + "init_args": {"model_name": "cpsam", "checkpoint_tiles": 1}, + }, + "slab_halfwidth": 1, + "crop_size": 240, +} +# seg_consistency_weight per arm: the median over 20 seeded val batches of ||grad base|| / +# ||grad seg_consistency|| at the fork checkpoint, at the leaves' 240-px crop (calibrate_variant.py +# --variants dice segcons --ckpt --tag segcons_crop240, micro-batch 1, A40 2026-10-16), +# 2 significant figures [IQR]: FNet-3D 0.3670 [0.138-0.654], UNeXt2-3D 0.03238 [0.0243-0.0657], +# pix2pix2d 0.1912 [0.0673-0.342]. Source: experiments/2026-09-24_spotlight-v2/calibration/results/ +# nucleus__{fnet3d_vscyto3daug_v2,fcmae_vscyto3d_scratch_v2,pix2pix2d_unetvit}__segcons_crop240{.csv,__meta.json}. +# (The first calibration, at the 4-tile 256-px crop, gave 0.48 / 0.050 / 0.57: results/*__segcons.*.) +# Known limitation, kept for comparability with the segaux weights: the convention is the segaux +# one, so for pix2pix2d the gradients are on generator_ema with base = lambda_l1 * L1 (the +# adversarial term is ignored), and micro-batch gradients are summed to the val batch's mean, not +# averaged over the training batch's composition (per-rank batch, DDP world size). +SEG_CONSISTENCY_WEIGHTS: dict[tuple[str, str], float] = { + ("nucleus", "fnet3d_vscyto3daug_segcons"): 0.37, + ("nucleus", "fcmae_vscyto3d_scratch_segcons"): 0.032, + ("nucleus", "pix2pix2d_unetvit_segcons"): 0.19, +} +_FORK_SUFFIXES: tuple[str, ...] = ("segcons", "segcons_ctrl") SAFE_CROP_SIZE = [1, 384, 384] MASK_VELOCITY_WEIGHT = 1.0 # C-cond predict conditions on the same-dimensionality FNet `_segaux` prediction for the @@ -399,6 +474,8 @@ def store_dir(self) -> str: Arm("fcmae_vscyto3d_scratch", "segaux_cpout", ("membrane",), a549=True), # Track H: CellDiff-2D trained on a background-low-passed target (H1; launched as H2). *(Arm("celldiff_2d", s, ORGANELLES, a549=True) for s in BG_TARGET_ARGS), + # P1 (2026-10-15): segmenter-consistency fine-tunes and their no-term controls, nucleus only. + *(Arm(m, s, ("nucleus",), a549=True) for m in SEGCONS_BASELINES for s in _FORK_SUFFIXES), ) # Second draws: suffix -> the arm whose recipe it re-draws with seed_everything: 1. SEED_SOURCES: dict[str, str] = { @@ -458,6 +535,11 @@ def _long_wall(arm: Arm) -> bool: + "), a training-only target transform; validation and prediction see the raw target" for suffix, args in BG_TARGET_ARGS.items() }, + "segcons_ctrl": "fork: model ckpt_path (the converged baseline checkpoint, weights only) + the fine-tune " + "budget and LR of FORK_SCHEDULES; no other change (the no-term control of segcons)", + "segcons": "the segcons_ctrl fork + model seg_consistency (SegConsistencyLoss, Cellpose-SAM cpsam, slab " + "halfwidth 1, 240-px crop (one cellpose tile), per-tile activation checkpointing) + seg_consistency_weight " + "(gradient-norm calibrated at the fork checkpoint)", } # Keys every fit / predict arm renames so it owns its run dir, wandb run and job. @@ -569,6 +651,12 @@ def allowed_diff(arm: Arm, kind: str) -> tuple[frozenset[str], frozenset[str]]: } elif arm.suffix in BG_TARGET_ARGS: recipe |= {"data.init_args.fg_mask_key", "model.init_args.target_bg_lowpass"} + elif arm.suffix in _FORK_SUFFIXES: + schedule = FORK_SCHEDULES[arm.baseline] + recipe.add("model.init_args.ckpt_path") + recipe |= {f"trainer.{k}" for k in schedule["trainer"]} | {f"model.init_args.{k}" for k in schedule["model"]} + if arm.suffix == "segcons": + recipe |= {"model.init_args.seg_consistency", "model.init_args.seg_consistency_weight"} if _long_wall(arm): recipe.add("base") return _FIT_RENAMES, frozenset(recipe) @@ -663,6 +751,17 @@ def _apply_recipe(arm: Arm, organelle: str, cfg: dict) -> None: "class_path": "viscy_utils.losses.BackgroundLowPass", "init_args": dict(BG_TARGET_ARGS[arm.suffix]), } + elif arm.suffix in _FORK_SUFFIXES: + schedule = FORK_SCHEDULES[arm.baseline] + model_args["ckpt_path"] = FORK_CKPTS[(organelle, arm.baseline)] + model_args.update(schedule["model"]) + cfg.setdefault("trainer", {}).update(schedule["trainer"]) + if arm.suffix == "segcons": + model_args["seg_consistency"] = { + "class_path": "viscy_utils.losses.SegConsistencyLoss", + "init_args": copy.deepcopy(SEG_CONSISTENCY_ARGS), + } + model_args["seg_consistency_weight"] = SEG_CONSISTENCY_WEIGHTS[(organelle, arm.model)] elif arm.suffix != "v2": raise ValueError(f"unknown arm suffix {arm.suffix!r}") # Drop containers the arm did not need (v2/seed1 add no data/model keys). @@ -753,6 +852,8 @@ def _header(arm: Arm, organelle: str, baseline_leaf: Path, kind: str) -> str: lines.append("# Inference is identical to the baseline's; only the checkpoint, store and names differ.") if not arm.fit: lines.append("# Predict-only arm: the baseline's own last.ckpt, own store; submit with `--ckpt last`.") + elif arm.suffix in _FORK_SUFFIXES: + lines.append("# Fine-tune arm: predict the end of the fine-tune; submit with `--ckpt last`.") else: lines.append("# ckpt_path is a PLACEHOLDER (last.ckpt): submit with `--ckpt best`, re-bake afterwards.") lines.append("# Everything else, and the rationale for it, is the baseline leaf's.") diff --git a/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py b/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py index 218c86044..73bd66fb5 100644 --- a/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py +++ b/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py @@ -2,9 +2,12 @@ from __future__ import annotations +import copy from collections import Counter +from pathlib import Path import pytest +import torch import yaml from generate_spotlight_v2_leaves import ( _FIT_RENAMES, @@ -12,11 +15,16 @@ BASELINES, BENCHMARKS, BG_TARGET_ARGS, + FORK_CKPTS, + FORK_SCHEDULES, ORGANELLES, POOL, SEED_SOURCES, SEG_AUX_WEIGHTS, + SEG_CONSISTENCY_ARGS, + SEG_CONSISTENCY_WEIGHTS, SEGAUXSELF_BASELINES, + SEGCONS_BASELINES, WEIGHT_SCALES, Arm, _unseeded, @@ -31,7 +39,7 @@ from dynacell._compose_hook import _dynacell_ref_resolver from dynacell.evaluation.paths import PAPER_KEY, canonical_model_name from viscy_utils.compose import load_composed_config -from viscy_utils.losses import BackgroundLowPass +from viscy_utils.losses import BackgroundLowPass, SegConsistencyLoss def _suffix(leaf_dir_name: str) -> str: @@ -56,7 +64,7 @@ def test_leaf_counts_per_arm(leaves: dict) -> None: segaux 19+76, segauxself 8+32, seed1 9+36, v2 7+28, probes 3+6 (l1 adds A549), cjoint/ccond 2+8 each, last 0+8, l1segaux 1+4, l1seed1 1+4, segaux_seed1 10+40, v2_seed1 7+28, segaux_halfw 2+8, segaux_doublew 2+8, segaux_sauna 5+20, segaux_sauna_seed1 1+4, segaux_cldice 3+12, - segaux_cpout 1+4, bglp 2+8, bgflat 2+8. + segaux_cpout 1+4, bglp 2+8, bgflat 2+8, segcons 3+12, segcons_ctrl 3+12. The jointsteps and safecrop probes are iPSC-only; every other arm also predicts the 3 A549 legs. """ @@ -84,6 +92,8 @@ def test_leaf_counts_per_arm(leaves: dict) -> None: "segaux_cpout": 1, "bglp": 2, "bgflat": 2, + "segcons": 3, + "segcons_ctrl": 3, } assert predicts == { "segaux": 76, @@ -108,6 +118,8 @@ def test_leaf_counts_per_arm(leaves: dict) -> None: "segaux_cpout": 4, "bglp": 8, "bgflat": 8, + "segcons": 12, + "segcons_ctrl": 12, } @@ -304,3 +316,73 @@ def test_bg_target_arms_compose_and_instantiate_with_only_the_target_op_added(su assert isinstance(op, BackgroundLowPass) assert {k: getattr(op, k) for k in BG_TARGET_ARGS[suffix]} == BG_TARGET_ARGS[suffix] assert (op.sigma_lp_z, op.sigma_feather_z, op.dilate_radius_z) == (0.0, 0.0, 0) + + +def _load_nucleus_fit(model: str) -> dict: + return load_composed_config(BENCHMARKS / "nucleus" / model / POOL / "train.yml", resolver=_dynacell_ref_resolver) + + +def _instantiate_model(model_cfg: dict) -> LightningModule: + """Build the model section as LightningCLI does (see the bg-target test for why LightningModule).""" + parser = ArgumentParser() + parser.add_subclass_arguments(LightningModule, "model") + return parser.instantiate(parser.parse_object({"model": model_cfg})).model + + +@pytest.mark.parametrize("baseline", SEGCONS_BASELINES) +def test_segcons_pair_is_one_fork_that_differs_by_the_term_only(baseline: str) -> None: + """Composed: ctrl = baseline + the fork keys; segcons = ctrl + seg_consistency and its weight.""" + base, ctrl, arm = (_load_nucleus_fit(m) for m in (baseline, f"{baseline}_segcons_ctrl", f"{baseline}_segcons")) + renames, _ = allowed_diff(Arm(baseline, "segcons", ("nucleus",), a549=False), "fit") + schedule = FORK_SCHEDULES[baseline] + fork_keys = ( + {"model.init_args.ckpt_path"} + | {f"trainer.{k}" for k in schedule["trainer"]} + | {f"model.init_args.{k}" for k in schedule["model"]} + ) + assert changed_keys(base, ctrl) - renames == fork_keys + term = "model.init_args.seg_consistency" + assert {term if k.startswith(term + ".") else k for k in changed_keys(ctrl, arm) - renames} == { + term, + "model.init_args.seg_consistency_weight", + } + assert ctrl["model"]["init_args"]["ckpt_path"] == FORK_CKPTS[("nucleus", baseline)] + assert arm["model"]["init_args"]["seg_consistency"]["init_args"] == SEG_CONSISTENCY_ARGS + # Same seed, data, loss and precision as the baseline: the pair is stage-matched to it. + for key in ("seed_everything", "data"): + assert arm.get(key) == base.get(key) + assert arm["trainer"]["precision"] == base["trainer"]["precision"] + + +@pytest.mark.parametrize("baseline", SEGCONS_BASELINES) +def test_segcons_ctrl_builds_from_the_real_fork_checkpoint(baseline: str) -> None: + """The fork checkpoint exists and loads strictly into the composed control's model.""" + ckpt = FORK_CKPTS[("nucleus", baseline)] + if not Path(ckpt).is_file(): + pytest.skip(f"fork checkpoint not reachable: {ckpt}") + module = _instantiate_model(_load_nucleus_fit(f"{baseline}_segcons_ctrl")["model"]) + state = torch.load(ckpt, map_location="cpu", weights_only=True, mmap=True)["state_dict"] + loaded = module.state_dict() + assert loaded.keys() == state.keys() + for key, value in state.items(): + assert torch.equal(loaded[key], value), key + + +@pytest.mark.parametrize("baseline", SEGCONS_BASELINES) +def test_segcons_builds_with_its_calibrated_weight(baseline: str) -> None: + """The emitted calibrated weight builds the cpsam term (a null weight is refused: test_seg_consistency_training).""" + cp_models = pytest.importorskip("cellpose.models") + if not (cp_models.MODEL_DIR / "cpsam").is_file(): + pytest.skip("cpsam checkpoint not in cellpose's model directory") + model_cfg = copy.deepcopy(_load_nucleus_fit(f"{baseline}_segcons")["model"]) + weight = SEG_CONSISTENCY_WEIGHTS[("nucleus", f"{baseline}_segcons")] + assert model_cfg["init_args"]["seg_consistency_weight"] == weight + del model_cfg["init_args"]["ckpt_path"] # the strict fork load is covered by the control's test + module = _instantiate_model(model_cfg) + assert module.seg_consistency_weight == weight + term = module.seg_consistency + assert isinstance(term, SegConsistencyLoss) + assert type(term.segmenter).__name__ == "CellposeDense" and term.segmenter.model_name == "cpsam" + assert term.segmenter.checkpoint_tiles == 1 + assert (term.slab_halfwidth, term.crop_size) == (1, 240) + assert not any(k.startswith("seg_consistency") for k in module.state_dict()) diff --git a/packages/viscy-utils/src/viscy_utils/losses/__init__.py b/packages/viscy-utils/src/viscy_utils/losses/__init__.py index b4064e410..9b6b89c25 100644 --- a/packages/viscy-utils/src/viscy_utils/losses/__init__.py +++ b/packages/viscy-utils/src/viscy_utils/losses/__init__.py @@ -1,6 +1,7 @@ from viscy_utils.losses.background_lowpass import BackgroundLowPass from viscy_utils.losses.mixed_loss import MixedLoss from viscy_utils.losses.seg_aux import SegAuxDice +from viscy_utils.losses.seg_consistency import SegConsistencyLoss from viscy_utils.losses.spotlight import SpotlightLoss -__all__ = ["BackgroundLowPass", "MixedLoss", "SegAuxDice", "SpotlightLoss"] +__all__ = ["BackgroundLowPass", "MixedLoss", "SegAuxDice", "SegConsistencyLoss", "SpotlightLoss"] diff --git a/packages/viscy-utils/src/viscy_utils/losses/seg_consistency.py b/packages/viscy-utils/src/viscy_utils/losses/seg_consistency.py new file mode 100644 index 000000000..58e658c10 --- /dev/null +++ b/packages/viscy-utils/src/viscy_utils/losses/seg_consistency.py @@ -0,0 +1,230 @@ +"""Segmenter-consistency loss for virtual staining (ControlNet++ style). + +A frozen, differentiable segmenter ``S`` is run on the prediction and on the target, +and their dense, pre-threshold outputs are matched: + +``L = mean_px ( mean_c (flow_pred - flow_target)**2 + KL(q_target || q_pred) )`` + +- ``flow`` are the segmenter's two flow channels and ``q = sigmoid(cellprob_logit)``. +- The probability term is the binary KL, i.e. BCE-with-logits against the target's + soft probability minus that probability's entropy: the gradient is BCE's, and the + loss is exactly 0 at ``pred == target``. The loss equals the G4 gate's + ``D_total = D_flow + D_prob``. +- ``S(target)`` is computed under ``no_grad`` on the same augmented crop, every step; + ``S(pred)`` carries the gradient into the prediction. ``S`` itself gets none. + +The segmenter sees 2-D images. A crop ``(B, 1, Z, Y, X)`` is reduced to the +``2 * slab_halfwidth + 1``-plane max projection around its nuclear-area focus plane: +the plane of largest Otsu foreground fraction of the target crop, with the outer 15% +of planes excluded and the crop center as the fallback, as +``dynacell.evaluation.focus.nucleus_area_plane`` picks the plane the eval segments +(slab halfwidth 1 there). The same planes are projected from the prediction. One +segmenter pass per sample is the cheapest geometry the G4 bench measured, and it is +the geometry the eval scores. ``Z = 1`` crops (2-D models) give their single plane. +The projection is then center-cropped to ``crop_size``: 240 px, which Cellpose pads to +exactly one 256-px tile (a 256-px crop pads to 272 and costs 2 x 2 tiles). + +The segmenter is held outside the module tree, so its weights never enter a +``state_dict`` (checkpoints stay the backbone's, and ``ckpt_path`` loads of baseline +checkpoints stay strict) and the module's ``train()`` cannot switch it out of eval +mode. It follows the prediction's device on first use. It runs with autocast +disabled, in its own dtype, as stock inference does. +""" + +import torch +import torch.nn.functional as F +from torch import Tensor, nn + +__all__ = ["SegConsistencyLoss", "nucleus_area_planes"] + +_GUARD_FRAC = 0.15 +_MIN_FOREGROUND = 0.01 +_OTSU_BINS = 256 + + +def _otsu_thresholds(flat: Tensor) -> Tensor: + """Per-row Otsu threshold, as ``skimage.filters.threshold_otsu`` (256 bins over each row's range). + + Parameters + ---------- + flat : Tensor + Values, shape ``(R, N)``. A constant row gets a meaningless threshold. + + Returns + ------- + Tensor + Threshold per row, shape ``(R,)``: the center of the bin maximizing the + between-class variance. + """ + lo = flat.amin(1, keepdim=True) + width = ((flat.amax(1, keepdim=True) - lo) / _OTSU_BINS).clamp_min(torch.finfo(flat.dtype).tiny) + idx = ((flat - lo) / width).floor().long().clamp(0, _OTSU_BINS - 1) + # scatter_add_ rather than bincount: bincount sizes its output on the host (a sync). + hist = torch.zeros(flat.shape[0], _OTSU_BINS, device=flat.device, dtype=torch.long) + hist = hist.scatter_add_(1, idx, torch.ones((), device=flat.device, dtype=torch.long).expand_as(idx)).double() + centers = lo.double() + width.double() * (torch.arange(_OTSU_BINS, device=flat.device) + 0.5) + w1 = hist.cumsum(1) + w2 = hist.flip(1).cumsum(1).flip(1) + m1 = (hist * centers).cumsum(1) / w1 + m2 = (hist * centers).flip(1).cumsum(1).flip(1) / w2 + var = w1[:, :-1] * w2[:, 1:] * (m1[:, :-1] - m2[:, 1:]) ** 2 + return centers.gather(1, var.argmax(1, keepdim=True)).squeeze(1).to(flat.dtype) + + +@torch.no_grad() +def nucleus_area_planes(vol: Tensor) -> Tensor: + """Return each volume's plane of largest nuclear foreground area, without a host sync. + + Mirrors ``dynacell.evaluation.focus.nucleus_area_plane``: Otsu over the whole + volume, per-plane foreground fraction, argmax inside a band that drops the outer + 15% of planes; a constant volume, or one whose band peak is below 1% foreground, + falls back to the center plane. + + Parameters + ---------- + vol : Tensor + Volumes, shape ``(B, Z, Y, X)``. + + Returns + ------- + Tensor + Plane index per volume, shape ``(B,)``, int64. + """ + b, z = vol.shape[:2] + center = torch.full((b,), z // 2, device=vol.device, dtype=torch.long) + flat = vol.reshape(b, -1).float() + # A constant row has no threshold; its NaN/arbitrary one is discarded below. + constant = flat.amax(1) <= flat.amin(1) + thr = _otsu_thresholds(flat) + frac = (flat > thr.unsqueeze(1)).reshape(b, z, -1).float().mean(2) + lo = int(round(_GUARD_FRAC * z)) + hi = max(lo + 1, int(round((1.0 - _GUARD_FRAC) * z))) + band = frac[:, lo:hi] + peak, arg = band.max(1) + use_center = constant | (peak < _MIN_FOREGROUND) + return torch.where(use_center, center, lo + arg) + + +def _slab_mip(vol: Tensor, planes: Tensor, halfwidth: int) -> Tensor: + """Max-project ``(B, Z, Y, X)`` onto the planes ``[p - h, p + h]`` clipped to the stack. + + Clipping repeats an edge plane, which leaves the maximum unchanged, so this equals + ``dynacell.evaluation.focus.slab_mip``. Returns ``(B, Y, X)``. + """ + z = vol.shape[1] + offsets = torch.arange(-halfwidth, halfwidth + 1, device=vol.device) + idx = (planes.unsqueeze(1) + offsets).clamp(0, z - 1) + idx = idx[:, :, None, None].expand(-1, -1, *vol.shape[2:]) + return vol.gather(1, idx).amax(1) + + +def _center_crop(img: Tensor, size: int | None) -> Tensor: + """Center-crop ``(..., Y, X)`` to at most ``size`` per side (no-op when ``None``).""" + if size is None: + return img + y, x = img.shape[-2:] + oy, ox = max(0, (y - size) // 2), max(0, (x - size) // 2) + return img[..., oy : oy + min(y, size), ox : ox + min(x, size)] + + +def _binary_kl(p_logit: Tensor, q_logit: Tensor) -> Tensor: + """Per-element binary KL(q || p) for ``q = sigmoid(q_logit)``; 0 iff the logits match.""" + q = torch.sigmoid(q_logit) + bce = F.binary_cross_entropy_with_logits(p_logit, q, reduction="none") + ent = F.binary_cross_entropy_with_logits(q_logit, q, reduction="none") + return bce - ent + + +class SegConsistencyLoss(nn.Module): + """Match a frozen segmenter's flows and cell probability on prediction and target. + + Parameters + ---------- + segmenter : nn.Module + Frozen differentiable segmenter mapping ``(N, 1, Y, X)`` images in any + positive-affine units to ``{"flow": (N, 2, Y, X), "cellprob_logit": (N, 1, Y, X)}`` + (e.g. ``dynacell.cellpose_dense.CellposeDense``). Put in eval mode with + ``requires_grad=False`` here; never the evaluation segmenter. + slab_halfwidth : int + Half-width of the max projection around a 3-D crop's focus plane. + crop_size : int or None + Center-crop side of the projected image, ``None`` for the whole image. + """ + + def __init__( + self, + segmenter: nn.Module, + slab_halfwidth: int = 1, + crop_size: int | None = 240, + ) -> None: + super().__init__() + if slab_halfwidth < 0: + raise ValueError(f"slab_halfwidth must be >= 0, got {slab_halfwidth}") + if crop_size is not None and crop_size < 1: + raise ValueError(f"crop_size must be >= 1 or None, got {crop_size}") + segmenter.eval() + segmenter.requires_grad_(False) + # Bypass nn.Module.__setattr__ so the segmenter is not a registered submodule. + self.__dict__["segmenter"] = segmenter + self.slab_halfwidth = slab_halfwidth + self.crop_size = crop_size + + def project(self, pred: Tensor, target: Tensor) -> tuple[Tensor, Tensor]: + """Reduce ``pred`` and ``target`` to the 2-D images the segmenter scores. + + Parameters + ---------- + pred, target : Tensor + Shape ``(B, 1, Z, Y, X)``. + + Returns + ------- + tuple of (Tensor, Tensor) + Projected ``pred`` and ``target``, each ``(B, 1, Y', X')``. + """ + if pred.shape != target.shape: + raise ValueError(f"pred and target must share a shape; got {tuple(pred.shape)}, {tuple(target.shape)}") + if pred.ndim != 5 or pred.shape[1] != 1: + raise ValueError(f"expected (B, 1, Z, Y, X), got {tuple(pred.shape)}") + planes = nucleus_area_planes(target[:, 0].detach()) + pred = _slab_mip(pred[:, 0], planes, self.slab_halfwidth).unsqueeze(1) + target = _slab_mip(target[:, 0], planes, self.slab_halfwidth).unsqueeze(1) + return _center_crop(pred, self.crop_size), _center_crop(target, self.crop_size) + + def _segment(self, img: Tensor) -> dict[str, Tensor]: + seg = self.segmenter + if next(seg.parameters()).device != img.device: + seg.to(img.device) + with torch.autocast(device_type=img.device.type, enabled=False): + return seg(img.to(torch.promote_types(img.dtype, torch.float32))) + + def forward( + self, pred: Tensor, target: Tensor, return_components: bool = False + ) -> Tensor | tuple[Tensor, dict[str, Tensor]]: + """Compute the consistency loss. + + Parameters + ---------- + pred : Tensor + Prediction, ``(B, 1, Z, Y, X)``. + target : Tensor + Ground truth, same shape. + return_components : bool + Also return ``{"segcons", "segcons_flow", "segcons_prob"}``. + + Returns + ------- + Tensor or tuple of (Tensor, dict of str to Tensor) + Scalar float32 loss, optionally with its flow and probability terms. + """ + pred2d, target2d = self.project(pred, target.detach()) + with torch.no_grad(): + ref = self._segment(target2d) + out = self._segment(pred2d) + flow_sq = ((out["flow"] - ref["flow"]) ** 2).mean(1) + kl = _binary_kl(out["cellprob_logit"], ref["cellprob_logit"])[:, 0] + flow_term, prob_term = flow_sq.mean(), kl.mean() + loss = flow_term + prob_term + if return_components: + return loss, {"segcons": loss, "segcons_flow": flow_term, "segcons_prob": prob_term} + return loss diff --git a/packages/viscy-utils/tests/test_seg_consistency_loss.py b/packages/viscy-utils/tests/test_seg_consistency_loss.py new file mode 100644 index 000000000..6217d02b1 --- /dev/null +++ b/packages/viscy-utils/tests/test_seg_consistency_loss.py @@ -0,0 +1,171 @@ +"""Tests for SegConsistencyLoss on a tiny stand-in segmenter (the Cellpose-SAM one is in dynacell).""" + +import pytest +import torch +from torch import Tensor, nn + +from viscy_utils.losses import SegConsistencyLoss +from viscy_utils.losses.seg_consistency import _OTSU_BINS, _binary_kl, _otsu_thresholds, nucleus_area_planes + + +class TinySegmenter(nn.Module): + """Two conv layers to flow (2 ch) + cell-probability logit (1 ch), per-image min-max normalized.""" + + def __init__(self, seed: int = 0) -> None: + super().__init__() + torch.manual_seed(seed) + self.net = nn.Sequential(nn.Conv2d(1, 8, 3, padding=1), nn.Tanh(), nn.Conv2d(8, 3, 3, padding=1)) + + def forward(self, x: Tensor) -> dict[str, Tensor]: + with torch.no_grad(): + lo = x.amin((1, 2, 3), keepdim=True) + hi = x.amax((1, 2, 3), keepdim=True) + y = self.net((x - lo) / (hi - lo + 1e-6)) + return {"flow": y[:, :2], "cellprob_logit": y[:, 2:3]} + + +def _nuclei(shape: tuple[int, ...], seed: int = 0) -> Tensor: + """Bright blobs on noise, peaked in z around a random plane per sample.""" + g = torch.Generator().manual_seed(seed) + b, z, y, x = shape[0], shape[2], shape[3], shape[4] + yy, xx = torch.meshgrid(torch.arange(y), torch.arange(x), indexing="ij") + out = 0.05 * torch.randn(shape, generator=g) + for i in range(b): + zc = float(torch.randint(z // 4, max(z // 4 + 1, 3 * z // 4), (1,), generator=g)) + for _ in range(4): + cy, cx = (torch.rand(2, generator=g) * torch.tensor([y, x])).tolist() + disk = ((yy - cy) ** 2 + (xx - cx) ** 2 < 36).float() + prof = torch.exp(-(((torch.arange(z) - zc) / 2.0) ** 2)) + out[i, 0] += prof[:, None, None] * disk + return out + + +@pytest.mark.parametrize("shape", [(2, 1, 8, 32, 32), (2, 1, 1, 32, 32)], ids=["3d", "z1"]) +def test_zero_at_pred_equals_target_and_positive_otherwise(shape): + target = _nuclei(shape) + loss_fn = SegConsistencyLoss(TinySegmenter()) + assert float(loss_fn(target.clone(), target)) == pytest.approx(0.0, abs=1e-6) + g = torch.Generator().manual_seed(1) + assert float(loss_fn(target + 0.3 * torch.randn(shape, generator=g), target)) > 1e-4 + + +def test_invariant_to_positive_affine_units_of_both_inputs(): + """The segmenter normalizes per image, so a z-scored pred vs a raw target loses nothing.""" + target = _nuclei((2, 1, 6, 32, 32)) + loss_fn = SegConsistencyLoss(TinySegmenter()) + assert float(loss_fn(3.0 * target + 7.0, target)) == pytest.approx(0.0, abs=1e-5) + + +def test_components_match_their_definition(): + target = _nuclei((2, 1, 6, 32, 32)) + pred = target + 0.2 * torch.randn_like(target) + seg = TinySegmenter() + loss_fn = SegConsistencyLoss(seg, crop_size=None) + loss, comps = loss_fn(pred, target, return_components=True) + pred2d, target2d = loss_fn.project(pred, target) + out, ref = seg(pred2d), seg(target2d) + d_flow = ((out["flow"] - ref["flow"]) ** 2).mean() + d_prob = _binary_kl(out["cellprob_logit"], ref["cellprob_logit"]).mean() + torch.testing.assert_close(comps["segcons_flow"], d_flow) + torch.testing.assert_close(comps["segcons_prob"], d_prob) + torch.testing.assert_close(loss, d_flow + d_prob) + assert comps["segcons"] is loss + + +def test_gradient_reaches_pred_only_and_matches_finite_difference(): + target = _nuclei((1, 1, 6, 16, 16)).double() + seg = TinySegmenter().double() + loss_fn = SegConsistencyLoss(seg) + pred = (target + 0.2 * torch.randn_like(target)).requires_grad_(True) + loss_fn(pred, target).backward() + assert torch.isfinite(pred.grad).all() and pred.grad.abs().sum() > 0 + assert all(p.grad is None for p in seg.parameters()) + # Only the 3 focus-slab planes feed the segmenter. + plane = int(nucleus_area_planes(target[:, 0])[0]) + touched = (pred.grad[0, 0].abs().sum((1, 2)) > 0).nonzero().flatten().tolist() + assert set(touched) <= {plane - 1, plane, plane + 1} + # Directional central difference. The segmenter's min/max statistics are stop-gradient, + # so perturb away from the projected image's extreme pixels, which leaves them fixed. + v = torch.zeros_like(pred) + v[0, 0, plane] = torch.randn(16, 16, dtype=torch.float64) + pred2d, _ = loss_fn.project(pred.detach(), target) + for flat_idx in (pred2d.argmax(), pred2d.argmin()): + y, x = divmod(int(flat_idx), 16) + v[0, 0, :, y, x] = 0.0 + eps = 1e-6 + with torch.no_grad(): + num = (loss_fn(pred + eps * v, target) - loss_fn(pred - eps * v, target)) / (2 * eps) + ana = (pred.grad * v).sum() + assert float(num) == pytest.approx(float(ana), rel=1e-3) + + +def test_segmenter_is_frozen_unregistered_and_stays_in_eval(): + seg = TinySegmenter() + loss_fn = SegConsistencyLoss(seg) + parent = nn.Sequential(nn.Linear(2, 2)) + parent.loss = loss_fn + assert all(not k.startswith("loss") for k in parent.state_dict()) + assert list(loss_fn.parameters()) == [] + parent.train() + assert not seg.training + assert not any(p.requires_grad for p in seg.parameters()) + + +def test_crop_size_shapes_what_the_segmenter_sees(): + target = _nuclei((4, 1, 6, 40, 40)) + loss_fn = SegConsistencyLoss(TinySegmenter(), crop_size=24) + pred2d, target2d = loss_fn.project(target, target) + assert pred2d.shape == target2d.shape == (4, 1, 24, 24) + + +def test_rejects_bad_arguments_and_shapes(): + with pytest.raises(ValueError, match="slab_halfwidth"): + SegConsistencyLoss(TinySegmenter(), slab_halfwidth=-1) + with pytest.raises(ValueError, match="crop_size"): + SegConsistencyLoss(TinySegmenter(), crop_size=0) + loss_fn = SegConsistencyLoss(TinySegmenter()) + with pytest.raises(ValueError, match="share a shape"): + loss_fn(torch.zeros(1, 1, 4, 8, 8), torch.zeros(1, 1, 4, 8, 9)) + with pytest.raises(ValueError, match="expected"): + loss_fn(torch.zeros(1, 2, 4, 8, 8), torch.zeros(1, 2, 4, 8, 8)) + with pytest.raises(ValueError, match="expected"): + loss_fn(torch.zeros(1, 1, 8, 8), torch.zeros(1, 1, 8, 8)) + + +def test_nucleus_area_planes_fallbacks(): + assert nucleus_area_planes(torch.zeros(2, 7, 8, 8)).tolist() == [3, 3] + assert nucleus_area_planes(torch.zeros(1, 1, 8, 8)).tolist() == [0] + + +def _otsu_thresholds_bincount(flat: Tensor) -> Tensor: + """The first (bincount-histogram) implementation of ``_otsu_thresholds``, kept as the reference.""" + r = flat.shape[0] + lo = flat.amin(1, keepdim=True) + width = ((flat.amax(1, keepdim=True) - lo) / _OTSU_BINS).clamp_min(torch.finfo(flat.dtype).tiny) + idx = ((flat - lo) / width).floor().long().clamp(0, _OTSU_BINS - 1) + idx += torch.arange(r, device=flat.device).unsqueeze(1) * _OTSU_BINS + hist = torch.bincount(idx.reshape(-1), minlength=r * _OTSU_BINS).reshape(r, _OTSU_BINS).double() + centers = lo.double() + width.double() * (torch.arange(_OTSU_BINS, device=flat.device) + 0.5) + w1 = hist.cumsum(1) + w2 = hist.flip(1).cumsum(1).flip(1) + m1 = (hist * centers).cumsum(1) / w1 + m2 = (hist * centers).flip(1).cumsum(1).flip(1) / w2 + var = w1[:, :-1] * w2[:, 1:] * (m1[:, :-1] - m2[:, 1:]) ** 2 + return centers.gather(1, var.argmax(1, keepdim=True)).squeeze(1).to(flat.dtype) + + +def test_otsu_thresholds_equal_the_bincount_implementation(): + vols = [_nuclei((3, 1, z, 24, 24), seed=z).reshape(3, -1) for z in (1, 5, 16)] + vols.append(torch.rand(2, 1000) ** 3 * 50 - 7) + for flat in vols: + assert torch.equal(_otsu_thresholds(flat), _otsu_thresholds_bincount(flat)) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="host-sync detection needs CUDA") +def test_focus_planes_make_no_host_sync(): + vol = _nuclei((2, 1, 8, 32, 32))[:, 0].cuda() + torch.cuda.set_sync_debug_mode("error") + try: + nucleus_area_planes(vol) + finally: + torch.cuda.set_sync_debug_mode("default")