From a1da77e6b3cbe52c4e6ac28d009c82c18c51bcf9 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Thu, 15 Oct 2026 22:08:10 -0700 Subject: [PATCH 01/20] feat(viscy-utils): add a segmenter-consistency loss SegConsistencyLoss matches a frozen segmenter's dense outputs on the prediction and the target: MSE on the two flow channels plus binary KL (BCE with logits minus the target entropy) on the cell probability, so the loss is 0 at pred == target. S(target) runs under no_grad on the same crop every step. A 3-D crop is reduced to the slab MIP around the target's nuclear-area plane, as the eval segments it; an optional Ctrl-U weight down-weights pixels where S disagrees with its flipped view. The segmenter stays out of the module tree, so checkpoints and strict ckpt_path loads are unchanged. Co-Authored-By: Claude Opus 5.5 --- .../src/viscy_utils/losses/__init__.py | 3 +- .../src/viscy_utils/losses/seg_consistency.py | 271 ++++++++++++++++++ .../tests/test_seg_consistency_loss.py | 144 ++++++++++ 3 files changed, 417 insertions(+), 1 deletion(-) create mode 100644 packages/viscy-utils/src/viscy_utils/losses/seg_consistency.py create mode 100644 packages/viscy-utils/tests/test_seg_consistency_loss.py 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..c2a15c414 --- /dev/null +++ b/packages/viscy-utils/src/viscy_utils/losses/seg_consistency.py @@ -0,0 +1,271 @@ +"""Segmenter-consistency loss for virtual staining (ControlNet++ / Ctrl-U 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 w * ( 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``. With ``w = 1`` 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. +- Optional Ctrl-U uncertainty weight: ``w = exp(-d / uncertainty_scale)``, with ``d`` + the per-pixel disagreement of ``S`` on the target and on its horizontal flip + (flipped back): the mean squared difference over the two flows and the probability. + It down-weights pixels where the segmenter itself is unsure. + Off by default (``w = 1``); it costs one more no-grad pass. + +The segmenter sees 2-D images. A 3-D 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) use their single plane. +The projection is then center-cropped to ``crop_size`` (one 256-px Cellpose tile). + +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. + """ + 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) + + +@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) + if z <= 1: + return torch.zeros(b, 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. + max_samples : int or None + Score only the first ``max_samples`` samples of each batch (the batch is + shuffled), ``None`` for all. + uncertainty_weight : bool + Weight pixels by ``exp(-d / uncertainty_scale)`` (Ctrl-U); see the module docstring. + uncertainty_scale : float + Disagreement scale of the uncertainty weight. + """ + + def __init__( + self, + segmenter: nn.Module, + slab_halfwidth: int = 1, + crop_size: int | None = 256, + max_samples: int | None = None, + uncertainty_weight: bool = False, + uncertainty_scale: float = 1.0, + ) -> 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}") + if max_samples is not None and max_samples < 1: + raise ValueError(f"max_samples must be >= 1 or None, got {max_samples}") + if uncertainty_scale <= 0: + raise ValueError(f"uncertainty_scale must be > 0, got {uncertainty_scale}") + 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 + self.max_samples = max_samples + self.uncertainty_weight = uncertainty_weight + self.uncertainty_scale = uncertainty_scale + + 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)`` or ``(B, 1, Y, X)``. + + Returns + ------- + tuple of (Tensor, Tensor) + Projected ``pred`` and ``target``, each ``(N, 1, Y', X')`` with + ``N = min(B, max_samples)``. + """ + 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 not in (4, 5) or pred.shape[1] != 1: + raise ValueError(f"expected (B, 1, Z, Y, X) or (B, 1, Y, X), got {tuple(pred.shape)}") + if self.max_samples is not None: + pred, target = pred[: self.max_samples], target[: self.max_samples] + if pred.ndim == 5: + 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))) + + @torch.no_grad() + def _uncertainty(self, target: Tensor, ref: dict[str, Tensor]) -> Tensor: + """Return ``exp(-d / scale)`` from the segmenter's disagreement with its flipped view.""" + flipped = self._segment(target.flip(-1)) + flow = flipped["flow"].flip(-1) + flow = torch.cat([flow[:, :1], -flow[:, 1:]], dim=1) # dX changes sign under an X flip + prob = torch.sigmoid(flipped["cellprob_logit"].flip(-1)) + prob_ref = torch.sigmoid(ref["cellprob_logit"]) + d = torch.cat([(ref["flow"] - flow) ** 2, (prob_ref - prob) ** 2], dim=1).mean(1) + return torch.exp(-d / self.uncertainty_scale) + + 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)`` or ``(B, 1, 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 (weighted) 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] + if self.uncertainty_weight: + w = self._uncertainty(target2d, ref) + flow_sq, kl = w * flow_sq, w * kl + 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..fd28ab435 --- /dev/null +++ b/packages/viscy-utils/tests/test_seg_consistency_loss.py @@ -0,0 +1,144 @@ +"""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 _binary_kl, 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_max_samples_and_crop_size_shape_what_the_segmenter_sees(): + target = _nuclei((4, 1, 6, 40, 40)) + loss_fn = SegConsistencyLoss(TinySegmenter(), crop_size=24, max_samples=3) + pred2d, target2d = loss_fn.project(target, target) + assert pred2d.shape == target2d.shape == (3, 1, 24, 24) + + +def test_uncertainty_weight_lowers_the_loss_and_keeps_zero_at_target(): + target = _nuclei((2, 1, 6, 32, 32)) + pred = target + 0.3 * torch.randn_like(target) + plain = SegConsistencyLoss(TinySegmenter()) + weighted = SegConsistencyLoss(TinySegmenter(), uncertainty_weight=True, uncertainty_scale=0.01) + assert 0.0 < float(weighted(pred, target)) < float(plain(pred, target)) + assert float(weighted(target, target)) == pytest.approx(0.0, abs=1e-6) + + +def test_rejects_bad_arguments_and_shapes(): + with pytest.raises(ValueError, match="slab_halfwidth"): + SegConsistencyLoss(TinySegmenter(), slab_halfwidth=-1) + with pytest.raises(ValueError, match="max_samples"): + SegConsistencyLoss(TinySegmenter(), max_samples=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)) + + +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] From ebf3b2fadf00affb7cee42c230e13571f4402cfe Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Thu, 15 Oct 2026 22:08:21 -0700 Subject: [PATCH 02/20] feat(dynacell): wire the segmenter-consistency term into the UNet and GAN engines DynacellUNet and DynacellGAN take seg_consistency + seg_consistency_weight, added next to seg_aux (training loss / generator objective). loss/validate stays base-only; the term is logged as loss/segcons_* (UNet) and loss/g_segcons_* (GAN) in training and loss/validate[_ema]_segcons in validation. The weight has no default: a term with weight None (the leaf placeholder) is refused at build, as is a weight without the term. Co-Authored-By: Claude Opus 5.5 --- applications/dynacell/src/dynacell/engine.py | 94 +++++++- .../tests/test_seg_consistency_training.py | 223 ++++++++++++++++++ 2 files changed, 314 insertions(+), 3 deletions(-) create mode 100644 applications/dynacell/tests/test_seg_consistency_training.py diff --git a/applications/dynacell/src/dynacell/engine.py b/applications/dynacell/src/dynacell/engine.py index 8b8425299..1ad6cd103 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,26 @@ def _check_seg_aux_args( ) +def _check_seg_consistency_args(seg_consistency: SegConsistencyLoss | None, weight: float | None) -> None: + """Reject a ``seg_consistency`` weight without the term, and the term without a weight. + + The weight has no default: it is calibrated per cell (median over validation + batches of ``||grad base|| / ||grad seg_consistency||`` at the starting + checkpoint), and a leaf that still carries the ``null`` placeholder must not train. + """ + 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_weight is unset (placeholder): calibrate it as the median gradient-norm ratio " + "||grad base|| / ||grad seg_consistency|| at the starting checkpoint and set it in the leaf." + ) + if weight < 0: + raise ValueError(f"seg_consistency_weight must be >= 0, got {weight}") + + 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 +613,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 +660,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 +670,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 +699,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) + 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 +711,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 +859,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 +879,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 +914,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 +977,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 +1682,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 +1723,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 +1736,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 +1774,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) + self.seg_consistency = seg_consistency + self.seg_consistency_weight = seg_consistency_weight self.lambda_l1 = lambda_l1 self.loss_type = loss_type @@ -1762,6 +1829,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 +2156,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 +2209,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 +2266,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 +2285,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 +2336,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/tests/test_seg_consistency_training.py b/applications/dynacell/tests/test_seg_consistency_training.py new file mode 100644 index 000000000..42cb4dfc5 --- /dev/null +++ b/applications/dynacell/tests/test_seg_consistency_training.py @@ -0,0 +1,223 @@ +"""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 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 weighted term changes the backbone gradient by exactly weight x its own gradient.""" + seed_everything(0) + batch = {"source": torch.randn(2, 1, 4, 32, 32), "target": torch.randn(2, 1, 4, 32, 32)} + baseline = DynacellUNet(architecture="FNet3D", model_config=FNET_CONFIG, log_batches_per_epoch=0) + grads = {} + for weight in (None, 0.0, 2.0): + kwargs = {} if weight is None else {"seg_consistency": _segcons(), "seg_consistency_weight": weight} + module = DynacellUNet(architecture="FNet3D", model_config=FNET_CONFIG, log_batches_per_epoch=0, **kwargs) + module.load_state_dict(baseline.state_dict()) + module.log = lambda *args, **kw: None + module.training_step(batch, 0).backward() + grads[weight] = torch.cat([p.grad.flatten() for p in module.model.parameters()]) + torch.testing.assert_close(grads[None], grads[0.0]) + assert not torch.allclose(grads[None], grads[2.0]) + + +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_rejected_without_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="placeholder"): + 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) + with pytest.raises(ValueError, match=">= 0"): + cls(**kwargs, seg_consistency=_segcons(), seg_consistency_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 From 5ab832acbab82d1a1163360f0913eb5c5896f0d4 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Thu, 15 Oct 2026 22:08:22 -0700 Subject: [PATCH 03/20] feat(dynacell): add a differentiable Cellpose-SAM wrapper CellposeDense returns cellpose flows and cell-probability logits that are differentiable w.r.t. the input, ported from the G4 gate: normalize99 with no-grad statistics, 256-px tiles, taper blend, loaded by absolute path from cellpose's model dir (no silent fallback or download), cpdino refused. With a gradient, tiles run under activation checkpointing (checkpoint_tiles, default 1): a stored fwd/bwd holds ~5.6 GB per tile. Tests: parity with CellposeModel.eval on identical normalized input, finite-difference gradient, checkpointing parity, and one DynacellUNet step through SegConsistencyLoss (GPU + cpsam only). Co-Authored-By: Claude Opus 5.5 --- .../dynacell/src/dynacell/cellpose_dense.py | 185 ++++++++++++++++++ .../dynacell/tests/test_cellpose_dense.py | 149 ++++++++++++++ 2 files changed, 334 insertions(+) create mode 100644 applications/dynacell/src/dynacell/cellpose_dense.py create mode 100644 applications/dynacell/tests/test_cellpose_dense.py diff --git a/applications/dynacell/src/dynacell/cellpose_dense.py b/applications/dynacell/src/dynacell/cellpose_dense.py new file mode 100644 index 000000000..256818dad --- /dev/null +++ b/applications/dynacell/src/dynacell/cellpose_dense.py @@ -0,0 +1,185 @@ +"""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, lower: float = 1.0, upper: float = 99.0) -> Tensor: + """Apply cellpose ``normalize99`` per image, with the percentiles under ``no_grad``. + + Parameters + ---------- + x : Tensor + Images, shape ``(B, ...)``. + lower, upper : float + Percentiles mapped to 0 and 1. + + Returns + ------- + Tensor + ``(x - p_lower) / (p_upper - p_lower)`` 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() + q = torch.quantile(flat, torch.tensor([lower / 100, upper / 100], device=x.device), dim=1) + lo, hi = q[0], q[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``. + + The network always stays in eval mode with ``requires_grad=False``; ``train()`` + is a no-op. It is loaded on the CPU and follows the module's ``.to(device)``. + + 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}") + 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 + self.net.requires_grad_(False) + self.model_name = model_name + self.checkpoint_tiles = checkpoint_tiles + self.train(False) + + def train(self, mode: bool = True) -> "CellposeDense": + """Keep the frozen network in eval mode whatever ``mode`` is.""" + return super().train(False) + + 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 = max(224, 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 + yp1, yp2, xp1, xp2 = cp_transforms.get_pad_yx(ly0, lx0, min_size=(bs, bs)) + xn = F.pad(xn, (xp1, xp2, yp1, yp2)) + ly, lx = xn.shape[-2:] + # make_tiles (non-augmented); ly, lx >= bs after get_pad_yx. + ny = 1 if ly <= bs else int(np.ceil((1.0 + 2 * self.TILE_OVERLAP) * ly / bs)) + nx = 1 if lx <= bs else int(np.ceil((1.0 + 2 * self.TILE_OVERLAP) * lx / bs)) + ys = np.linspace(0, ly - bs, ny).astype(int) + xs = np.linspace(0, lx - bs, nx).astype(int) + coords = [(int(a), int(c)) for a in ys for c in xs] + 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 torch.is_grad_enabled() 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/tests/test_cellpose_dense.py b/applications/dynacell/tests/test_cellpose_dense.py new file mode 100644 index 000000000..aa08c7f73 --- /dev/null +++ b/applications/dynacell/tests/test_cellpose_dense.py @@ -0,0 +1,149 @@ +"""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 +from cellpose import models as cp_models +from cellpose import transforms as cp_transforms + +from dynacell.cellpose_dense import CellposeDense, normalize99_diff +from dynacell.engine import DynacellUNet +from viscy_utils.losses import SegConsistencyLoss + +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_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) + 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.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()) From dc25ad5ada4bd082c9236af29ea99f0877d2d65d Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Thu, 15 Oct 2026 22:08:32 -0700 Subject: [PATCH 04/20] feat(dynacell): add P1 segmenter-consistency fine-tune arms (nucleus) Spotlight-v2 generator arms segcons / segcons_ctrl on FNet-3D vscyto3daug, UNeXt2-3D and pix2pix2d nucleus: both fork the baseline recipe's final checkpoint weights-only (model.init_args.ckpt_path) and fine-tune for 10% of its budget at 0.1x its peak LR; segcons adds SegConsistencyLoss with cpsam (G4's S_train). seg_consistency_weight is a null placeholder until calibrated, which the engines refuse. Predict leaves read last.ckpt. Registers the six tokens in PAPER_KEY and the grouped parser, and admits the pix2pix2d forks to the <80 GB cards like their baseline. Co-Authored-By: Claude Opus 5.5 --- .../predict__a549_mantis_denv.yml | 44 ++++++++ .../predict__a549_mantis_mock.yml | 44 ++++++++ .../predict__a549_mantis_zikv.yml | 44 ++++++++ .../ipsc_confocal/predict__ipsc_confocal.yml | 42 ++++++++ .../ipsc_confocal/train.yml | 72 +++++++++++++ .../predict__a549_mantis_denv.yml | 44 ++++++++ .../predict__a549_mantis_mock.yml | 44 ++++++++ .../predict__a549_mantis_zikv.yml | 44 ++++++++ .../ipsc_confocal/predict__ipsc_confocal.yml | 42 ++++++++ .../ipsc_confocal/train.yml | 60 +++++++++++ .../predict__a549_mantis_h2b_denv.yml | 44 ++++++++ .../predict__a549_mantis_h2b_mock.yml | 44 ++++++++ .../predict__a549_mantis_h2b_zikv.yml | 44 ++++++++ .../ipsc_confocal/predict__ipsc_confocal.yml | 42 ++++++++ .../ipsc_confocal/train.yml | 100 ++++++++++++++++++ .../predict__a549_mantis_h2b_denv.yml | 44 ++++++++ .../predict__a549_mantis_h2b_mock.yml | 44 ++++++++ .../predict__a549_mantis_h2b_zikv.yml | 44 ++++++++ .../ipsc_confocal/predict__ipsc_confocal.yml | 42 ++++++++ .../ipsc_confocal/train.yml | 88 +++++++++++++++ .../predict__a549_mantis_denv.yml | 45 ++++++++ .../predict__a549_mantis_mock.yml | 45 ++++++++ .../predict__a549_mantis_zikv.yml | 45 ++++++++ .../ipsc_confocal/predict__ipsc_confocal.yml | 43 ++++++++ .../ipsc_confocal/train.yml | 83 +++++++++++++++ .../predict__a549_mantis_denv.yml | 45 ++++++++ .../predict__a549_mantis_mock.yml | 45 ++++++++ .../predict__a549_mantis_zikv.yml | 45 ++++++++ .../ipsc_confocal/predict__ipsc_confocal.yml | 43 ++++++++ .../ipsc_confocal/train.yml | 71 +++++++++++++ .../dynacell/src/dynacell/evaluation/paths.py | 9 ++ .../test_benchmark_config_composition.py | 7 ++ applications/dynacell/tests/test_paths.py | 6 ++ .../tools/generate_grouped_eval_configs.py | 12 +++ .../tools/generate_spotlight_v2_leaves.py | 90 ++++++++++++++++ .../generate_spotlight_v2_leaves_test.py | 85 ++++++++++++++- 36 files changed, 1733 insertions(+), 2 deletions(-) create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc_confocal/predict__a549_mantis_denv.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc_confocal/predict__a549_mantis_mock.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc_confocal/predict__a549_mantis_zikv.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc_confocal/predict__ipsc_confocal.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/ipsc_confocal/train.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc_confocal/predict__a549_mantis_denv.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc_confocal/predict__a549_mantis_mock.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc_confocal/predict__a549_mantis_zikv.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc_confocal/predict__ipsc_confocal.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons_ctrl/ipsc_confocal/train.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/predict__a549_mantis_h2b_denv.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/predict__a549_mantis_h2b_mock.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/predict__a549_mantis_h2b_zikv.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/predict__ipsc_confocal.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/train.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc_confocal/predict__a549_mantis_h2b_denv.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc_confocal/predict__a549_mantis_h2b_mock.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc_confocal/predict__a549_mantis_h2b_zikv.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc_confocal/predict__ipsc_confocal.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons_ctrl/ipsc_confocal/train.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons/ipsc_confocal/predict__a549_mantis_denv.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons/ipsc_confocal/predict__a549_mantis_mock.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons/ipsc_confocal/predict__a549_mantis_zikv.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons/ipsc_confocal/predict__ipsc_confocal.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons/ipsc_confocal/train.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/predict__a549_mantis_denv.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/predict__a549_mantis_mock.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/predict__a549_mantis_zikv.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/predict__ipsc_confocal.yml create mode 100644 applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/train.yml 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..4ee286836 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fcmae_vscyto3d_scratch_segcons/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` 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, 256-px crop, per-tile activation checkpointing) + seg_consistency_weight (PLACEHOLDER null until calibrated). +# seg_consistency_weight is a PLACEHOLDER (null): calibrate it and set SEG_CONSISTENCY_WEIGHTS in the generator before submitting; the engine refuses to build with it unset. +# 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: 256 + seg_consistency_weight: null 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..d7d812b41 --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/fnet3d_vscyto3daug_segcons/ipsc_confocal/train.yml @@ -0,0 +1,100 @@ +# 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, 256-px crop, per-tile activation checkpointing) + seg_consistency_weight (PLACEHOLDER null until calibrated). +# seg_consistency_weight is a PLACEHOLDER (null): calibrate it and set SEG_CONSISTENCY_WEIGHTS in the generator before submitting; the engine refuses to build with it unset. +# 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: 256 + seg_consistency_weight: null 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..c7226d293 --- /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, 256-px crop, per-tile activation checkpointing) + seg_consistency_weight (PLACEHOLDER null until calibrated). +# seg_consistency_weight is a PLACEHOLDER (null): calibrate it and set SEG_CONSISTENCY_WEIGHTS in the generator before submitting; the engine refuses to build with it unset. +# 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 + 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: 256 + seg_consistency_weight: null 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..2fb76eeff --- /dev/null +++ b/applications/dynacell/configs/benchmarks/virtual_staining/nucleus/pix2pix2d_unetvit_segcons_ctrl/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_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 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_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/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..d61f51305 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 a ``null`` PLACEHOLDER until calibrated + (:data:`SEG_CONSISTENCY_WEIGHTS`); the engines refuse to build with it unset. 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,53 @@ ) _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. +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}, + }, +} +# 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 +# 256-px tile (the geometry the G4 bench timed); 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": 256, +} +# PLACEHOLDER weights (None -> `seg_consistency_weight: null`, which the engines refuse). Calibrate +# each as in calibration/calibrate_variant.py: the median over 20 seeded val batches of +# ||grad base|| / ||grad seg_consistency|| at the fork checkpoint, 2 significant figures. +SEG_CONSISTENCY_WEIGHTS: dict[tuple[str, str], float | None] = { + ("nucleus", f"{m}_segcons"): None for m in SEGCONS_BASELINES +} +_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 +458,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 +519,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, 256-px crop, per-tile activation checkpointing) + seg_consistency_weight (PLACEHOLDER null " + "until calibrated)", } # Keys every fit / predict arm renames so it owns its run dir, wandb run and job. @@ -569,6 +635,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 +735,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). @@ -744,6 +827,11 @@ def _header(arm: Arm, organelle: str, baseline_leaf: Path, kind: str) -> str: ] if kind == "fit": lines.append(f"# Recipe delta vs that leaf: {_DESCRIPTION[arm.suffix]}.") + if arm.suffix == "segcons" and SEG_CONSISTENCY_WEIGHTS[(organelle, arm.model)] is None: + lines.append( + "# seg_consistency_weight is a PLACEHOLDER (null): calibrate it and set SEG_CONSISTENCY_WEIGHTS" + " in the generator before submitting; the engine refuses to build with it unset." + ) if _long_wall(arm): lines.append( "# Wall: hardware_4gpu_long (7 d); 200 epochs at the measured ~2.0 ep/h is ~100 h > 4 d" @@ -753,6 +841,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..ac86d8d6d 100644 --- a/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py +++ b/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py @@ -2,21 +2,29 @@ from __future__ import annotations +import copy from collections import Counter +from pathlib import Path import pytest +import torch import yaml +from cellpose import models as cp_models from generate_spotlight_v2_leaves import ( _FIT_RENAMES, ARMS, BASELINES, BENCHMARKS, BG_TARGET_ARGS, + FORK_CKPTS, + FORK_SCHEDULES, ORGANELLES, POOL, SEED_SOURCES, SEG_AUX_WEIGHTS, + SEG_CONSISTENCY_ARGS, 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,72 @@ 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"] + assert module.state_dict().keys() == state.keys() + key = next(iter(state)) + torch.testing.assert_close(module.state_dict()[key], state[key]) + + +@pytest.mark.parametrize("baseline", SEGCONS_BASELINES) +def test_segcons_refuses_the_placeholder_weight_and_builds_once_set(baseline: str) -> None: + """The emitted ``seg_consistency_weight: null`` fails at build; a set weight builds the cpsam term.""" + 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"]) + assert model_cfg["init_args"]["seg_consistency_weight"] is None + del model_cfg["init_args"]["ckpt_path"] # the strict fork load is covered by the control's test + with pytest.raises(ValueError, match="placeholder"): + _instantiate_model(model_cfg) + model_cfg["init_args"]["seg_consistency_weight"] = 1.0 + module = _instantiate_model(model_cfg) + 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, 256) + assert not any(k.startswith("seg_consistency") for k in module.state_dict()) From 3825b42acc40d231b80d2390a5c1d095131af70b Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Thu, 15 Oct 2026 23:01:54 -0700 Subject: [PATCH 05/20] feat(dynacell): calibrate the P1 segcons weights seg_consistency_weight per arm is the median over 20 seeded val batches of ||grad base|| / ||grad seg_consistency|| at the fork checkpoint (calibrate_variant.py segcons variant, A40): FNet-3D 0.48 [0.25-0.93], UNeXt2-3D 0.050 [0.024-0.088], pix2pix2d 0.57 [0.32-0.67]. The tool's plain-Dice row reproduces the recorded FNet-3D vscyto3daug 4.5 on its July checkpoint (4.504 [2.58-6.72]). Co-Authored-By: Claude Opus 5.5 --- .../ipsc_confocal/train.yml | 5 ++--- .../ipsc_confocal/train.yml | 5 ++--- .../ipsc_confocal/train.yml | 5 ++--- .../tools/generate_spotlight_v2_leaves.py | 22 ++++++++++++------- .../generate_spotlight_v2_leaves_test.py | 13 ++++++----- 5 files changed, 28 insertions(+), 22 deletions(-) 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 index 4ee286836..8a15d152e 100644 --- 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 @@ -1,8 +1,7 @@ # 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, 256-px crop, per-tile activation checkpointing) + seg_consistency_weight (PLACEHOLDER null until calibrated). -# seg_consistency_weight is a PLACEHOLDER (null): calibrate it and set SEG_CONSISTENCY_WEIGHTS in the generator before submitting; the engine refuses to build with it unset. +# Recipe delta vs that leaf: the segcons_ctrl fork + model seg_consistency (SegConsistencyLoss, Cellpose-SAM cpsam, slab halfwidth 1, 256-px crop, 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 @@ -69,4 +68,4 @@ model: checkpoint_tiles: 1 slab_halfwidth: 1 crop_size: 256 - seg_consistency_weight: null + seg_consistency_weight: 0.05 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 index d7d812b41..10c2a61ef 100644 --- 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 @@ -1,8 +1,7 @@ # 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, 256-px crop, per-tile activation checkpointing) + seg_consistency_weight (PLACEHOLDER null until calibrated). -# seg_consistency_weight is a PLACEHOLDER (null): calibrate it and set SEG_CONSISTENCY_WEIGHTS in the generator before submitting; the engine refuses to build with it unset. +# Recipe delta vs that leaf: the segcons_ctrl fork + model seg_consistency (SegConsistencyLoss, Cellpose-SAM cpsam, slab halfwidth 1, 256-px crop, 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 @@ -97,4 +96,4 @@ model: checkpoint_tiles: 1 slab_halfwidth: 1 crop_size: 256 - seg_consistency_weight: null + seg_consistency_weight: 0.48 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 index c7226d293..7653d7b83 100644 --- 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 @@ -1,8 +1,7 @@ # 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, 256-px crop, per-tile activation checkpointing) + seg_consistency_weight (PLACEHOLDER null until calibrated). -# seg_consistency_weight is a PLACEHOLDER (null): calibrate it and set SEG_CONSISTENCY_WEIGHTS in the generator before submitting; the engine refuses to build with it unset. +# Recipe delta vs that leaf: the segcons_ctrl fork + model seg_consistency (SegConsistencyLoss, Cellpose-SAM cpsam, slab halfwidth 1, 256-px crop, 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 @@ -80,4 +79,4 @@ model: checkpoint_tiles: 1 slab_halfwidth: 1 crop_size: 256 - seg_consistency_weight: null + seg_consistency_weight: 0.57 diff --git a/applications/dynacell/tools/generate_spotlight_v2_leaves.py b/applications/dynacell/tools/generate_spotlight_v2_leaves.py index d61f51305..b91b59c47 100644 --- a/applications/dynacell/tools/generate_spotlight_v2_leaves.py +++ b/applications/dynacell/tools/generate_spotlight_v2_leaves.py @@ -88,8 +88,8 @@ 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 a ``null`` PLACEHOLDER until calibrated - (:data:`SEG_CONSISTENCY_WEIGHTS`); the engines refuse to build with it unset. Both + stage-matched. ``seg_consistency_weight`` is calibrated per arm at the fork checkpoint + (:data:`SEG_CONSISTENCY_WEIGHTS`; the engines refuse a ``null`` one). 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. @@ -279,11 +279,17 @@ "slab_halfwidth": 1, "crop_size": 256, } -# PLACEHOLDER weights (None -> `seg_consistency_weight: null`, which the engines refuse). Calibrate -# each as in calibration/calibrate_variant.py: the median over 20 seeded val batches of -# ||grad base|| / ||grad seg_consistency|| at the fork checkpoint, 2 significant figures. +# seg_consistency_weight per arm: the median over 20 seeded val batches of ||grad base|| / +# ||grad seg_consistency|| at the fork checkpoint (calibrate_variant.py --variants segcons --ckpt +# , micro-batch 1, A40 2026-10-15), 2 significant figures [IQR]. FNet-3D 0.4806 [0.246-0.928], +# UNeXt2-3D 0.0501 [0.0245-0.0883], pix2pix2d (generator_ema, base = lambda_l1 * L1) 0.5707 +# [0.321-0.669]. Source: experiments/2026-09-24_spotlight-v2/calibration/results/ +# nucleus__{fnet3d_vscyto3daug_v2,fcmae_vscyto3d_scratch_v2,pix2pix2d_unetvit}__segcons{.csv,__meta.json}. +# None would emit `seg_consistency_weight: null`, which the engines refuse. SEG_CONSISTENCY_WEIGHTS: dict[tuple[str, str], float | None] = { - ("nucleus", f"{m}_segcons"): None for m in SEGCONS_BASELINES + ("nucleus", "fnet3d_vscyto3daug_segcons"): 0.48, + ("nucleus", "fcmae_vscyto3d_scratch_segcons"): 0.050, + ("nucleus", "pix2pix2d_unetvit_segcons"): 0.57, } _FORK_SUFFIXES: tuple[str, ...] = ("segcons", "segcons_ctrl") SAFE_CROP_SIZE = [1, 384, 384] @@ -522,8 +528,8 @@ def _long_wall(arm: Arm) -> bool: "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, 256-px crop, per-tile activation checkpointing) + seg_consistency_weight (PLACEHOLDER null " - "until calibrated)", + "halfwidth 1, 256-px crop, 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. diff --git a/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py b/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py index ac86d8d6d..77552ec7d 100644 --- a/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py +++ b/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py @@ -23,6 +23,7 @@ SEED_SOURCES, SEG_AUX_WEIGHTS, SEG_CONSISTENCY_ARGS, + SEG_CONSISTENCY_WEIGHTS, SEGAUXSELF_BASELINES, SEGCONS_BASELINES, WEIGHT_SCALES, @@ -368,17 +369,19 @@ def test_segcons_ctrl_builds_from_the_real_fork_checkpoint(baseline: str) -> Non @pytest.mark.parametrize("baseline", SEGCONS_BASELINES) -def test_segcons_refuses_the_placeholder_weight_and_builds_once_set(baseline: str) -> None: - """The emitted ``seg_consistency_weight: null`` fails at build; a set weight builds the cpsam term.""" +def test_segcons_builds_with_its_calibrated_weight_and_refuses_a_null_one(baseline: str) -> None: + """The emitted calibrated weight builds the cpsam term; a ``null`` weight fails at build.""" 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"]) - assert model_cfg["init_args"]["seg_consistency_weight"] is None + weight = SEG_CONSISTENCY_WEIGHTS[("nucleus", f"{baseline}_segcons")] + assert weight is not None and 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 + model_cfg["init_args"]["seg_consistency_weight"] = None with pytest.raises(ValueError, match="placeholder"): _instantiate_model(model_cfg) - model_cfg["init_args"]["seg_consistency_weight"] = 1.0 - module = _instantiate_model(model_cfg) term = module.seg_consistency assert isinstance(term, SegConsistencyLoss) assert type(term.segmenter).__name__ == "CellposeDense" and term.segmenter.model_name == "cpsam" From 394211b3b49438c4504d1bb609e7f8b72bdd9113 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Fri, 16 Oct 2026 15:05:51 -0700 Subject: [PATCH 06/20] fix(dynacell): reject seg_consistency combined with seg_aux With both terms the logged loss/base_train would be the seg_aux base, not the loss before the consistency term. No arm combines them. Co-Authored-By: Claude Opus 5.5 --- applications/dynacell/src/dynacell/engine.py | 13 +++++++++---- .../dynacell/tests/test_seg_consistency_training.py | 6 ++++-- 2 files changed, 13 insertions(+), 6 deletions(-) diff --git a/applications/dynacell/src/dynacell/engine.py b/applications/dynacell/src/dynacell/engine.py index 1ad6cd103..70618e199 100644 --- a/applications/dynacell/src/dynacell/engine.py +++ b/applications/dynacell/src/dynacell/engine.py @@ -122,12 +122,15 @@ def _check_seg_aux_args( ) -def _check_seg_consistency_args(seg_consistency: SegConsistencyLoss | None, weight: float | None) -> None: - """Reject a ``seg_consistency`` weight without the term, and the term without a weight. +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), and a leaf that still carries the ``null`` placeholder must not train. + 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: @@ -140,6 +143,8 @@ def _check_seg_consistency_args(seg_consistency: SegConsistencyLoss | None, weig ) if weight < 0: raise ValueError(f"seg_consistency_weight must be >= 0, got {weight}") + 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: @@ -699,7 +704,7 @@ 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) + _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 @@ -1774,7 +1779,7 @@ 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) + _check_seg_consistency_args(seg_consistency, seg_consistency_weight, seg_aux) self.seg_consistency = seg_consistency self.seg_consistency_weight = seg_consistency_weight diff --git a/applications/dynacell/tests/test_seg_consistency_training.py b/applications/dynacell/tests/test_seg_consistency_training.py index 42cb4dfc5..99b9719da 100644 --- a/applications/dynacell/tests/test_seg_consistency_training.py +++ b/applications/dynacell/tests/test_seg_consistency_training.py @@ -14,7 +14,7 @@ from dynacell.engine import DynacellGAN, DynacellUNet from viscy_data.hcs import HCSDataModule -from viscy_utils.losses import SegConsistencyLoss +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 = { @@ -187,7 +187,7 @@ def test_gan_seg_consistency_fast_dev_run(tmp_path, tiny_hcs_zarr): @pytest.mark.parametrize("family", ["unet", "gan"]) -def test_weight_is_required_with_the_term_and_rejected_without_it(family): +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 @@ -200,6 +200,8 @@ def test_weight_is_required_with_the_term_and_rejected_without_it(family): cls(**kwargs, seg_consistency_weight=1.0) with pytest.raises(ValueError, match=">= 0"): cls(**kwargs, seg_consistency=_segcons(), seg_consistency_weight=-1.0) + 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"]) From d1dbdcd580c3c525493367df23b9c63490194985 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Fri, 16 Oct 2026 15:06:56 -0700 Subject: [PATCH 07/20] refactor(viscy-utils): drop unused SegConsistencyLoss options and the 4-D path No arm uses max_samples or the Ctrl-U uncertainty weight, and every engine hands the loss 5-D (B, 1, Z, Y, X) crops (Z = 1 for 2-D models), so the loss now requires 5-D. The z <= 1 early return in nucleus_area_planes is redundant: the guard band of a one-plane stack is plane 0. Co-Authored-By: Claude Opus 5.5 --- .../src/viscy_utils/losses/seg_consistency.py | 72 ++++--------------- .../tests/test_seg_consistency_loss.py | 21 ++---- 2 files changed, 22 insertions(+), 71 deletions(-) diff --git a/packages/viscy-utils/src/viscy_utils/losses/seg_consistency.py b/packages/viscy-utils/src/viscy_utils/losses/seg_consistency.py index c2a15c414..bae9c9b01 100644 --- a/packages/viscy-utils/src/viscy_utils/losses/seg_consistency.py +++ b/packages/viscy-utils/src/viscy_utils/losses/seg_consistency.py @@ -1,31 +1,26 @@ -"""Segmenter-consistency loss for virtual staining (ControlNet++ / Ctrl-U style). +"""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 w * ( mean_c (flow_pred - flow_target)**2 + KL(q_target || q_pred) )`` +``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``. With ``w = 1`` the loss equals the G4 - gate's ``D_total = D_flow + D_prob``. + 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. -- Optional Ctrl-U uncertainty weight: ``w = exp(-d / uncertainty_scale)``, with ``d`` - the per-pixel disagreement of ``S`` on the target and on its horizontal flip - (flipped back): the mean squared difference over the two flows and the probability. - It down-weights pixels where the segmenter itself is unsure. - Off by default (``w = 1``); it costs one more no-grad pass. -The segmenter sees 2-D images. A 3-D crop ``(B, 1, Z, Y, X)`` is reduced to the +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) use their single plane. +the geometry the eval scores. ``Z = 1`` crops (2-D models) give their single plane. The projection is then center-cropped to ``crop_size`` (one 256-px Cellpose tile). The segmenter is held outside the module tree, so its weights never enter a @@ -96,8 +91,6 @@ def nucleus_area_planes(vol: Tensor) -> Tensor: """ b, z = vol.shape[:2] center = torch.full((b,), z // 2, device=vol.device, dtype=torch.long) - if z <= 1: - return torch.zeros(b, 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) @@ -155,13 +148,6 @@ class SegConsistencyLoss(nn.Module): 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. - max_samples : int or None - Score only the first ``max_samples`` samples of each batch (the batch is - shuffled), ``None`` for all. - uncertainty_weight : bool - Weight pixels by ``exp(-d / uncertainty_scale)`` (Ctrl-U); see the module docstring. - uncertainty_scale : float - Disagreement scale of the uncertainty weight. """ def __init__( @@ -169,28 +155,18 @@ def __init__( segmenter: nn.Module, slab_halfwidth: int = 1, crop_size: int | None = 256, - max_samples: int | None = None, - uncertainty_weight: bool = False, - uncertainty_scale: float = 1.0, ) -> 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}") - if max_samples is not None and max_samples < 1: - raise ValueError(f"max_samples must be >= 1 or None, got {max_samples}") - if uncertainty_scale <= 0: - raise ValueError(f"uncertainty_scale must be > 0, got {uncertainty_scale}") 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 - self.max_samples = max_samples - self.uncertainty_weight = uncertainty_weight - self.uncertainty_scale = uncertainty_scale def project(self, pred: Tensor, target: Tensor) -> tuple[Tensor, Tensor]: """Reduce ``pred`` and ``target`` to the 2-D images the segmenter scores. @@ -198,24 +174,20 @@ def project(self, pred: Tensor, target: Tensor) -> tuple[Tensor, Tensor]: Parameters ---------- pred, target : Tensor - Shape ``(B, 1, Z, Y, X)`` or ``(B, 1, Y, X)``. + Shape ``(B, 1, Z, Y, X)``. Returns ------- tuple of (Tensor, Tensor) - Projected ``pred`` and ``target``, each ``(N, 1, Y', X')`` with - ``N = min(B, max_samples)``. + 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 not in (4, 5) or pred.shape[1] != 1: - raise ValueError(f"expected (B, 1, Z, Y, X) or (B, 1, Y, X), got {tuple(pred.shape)}") - if self.max_samples is not None: - pred, target = pred[: self.max_samples], target[: self.max_samples] - if pred.ndim == 5: - 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) + 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]: @@ -225,17 +197,6 @@ def _segment(self, img: Tensor) -> dict[str, Tensor]: with torch.autocast(device_type=img.device.type, enabled=False): return seg(img.to(torch.promote_types(img.dtype, torch.float32))) - @torch.no_grad() - def _uncertainty(self, target: Tensor, ref: dict[str, Tensor]) -> Tensor: - """Return ``exp(-d / scale)`` from the segmenter's disagreement with its flipped view.""" - flipped = self._segment(target.flip(-1)) - flow = flipped["flow"].flip(-1) - flow = torch.cat([flow[:, :1], -flow[:, 1:]], dim=1) # dX changes sign under an X flip - prob = torch.sigmoid(flipped["cellprob_logit"].flip(-1)) - prob_ref = torch.sigmoid(ref["cellprob_logit"]) - d = torch.cat([(ref["flow"] - flow) ** 2, (prob_ref - prob) ** 2], dim=1).mean(1) - return torch.exp(-d / self.uncertainty_scale) - def forward( self, pred: Tensor, target: Tensor, return_components: bool = False ) -> Tensor | tuple[Tensor, dict[str, Tensor]]: @@ -244,7 +205,7 @@ def forward( Parameters ---------- pred : Tensor - Prediction, ``(B, 1, Z, Y, X)`` or ``(B, 1, Y, X)``. + Prediction, ``(B, 1, Z, Y, X)``. target : Tensor Ground truth, same shape. return_components : bool @@ -253,7 +214,7 @@ def forward( Returns ------- Tensor or tuple of (Tensor, dict of str to Tensor) - Scalar float32 loss, optionally with its (weighted) flow and probability terms. + Scalar float32 loss, optionally with its flow and probability terms. """ pred2d, target2d = self.project(pred, target.detach()) with torch.no_grad(): @@ -261,9 +222,6 @@ def forward( out = self._segment(pred2d) flow_sq = ((out["flow"] - ref["flow"]) ** 2).mean(1) kl = _binary_kl(out["cellprob_logit"], ref["cellprob_logit"])[:, 0] - if self.uncertainty_weight: - w = self._uncertainty(target2d, ref) - flow_sq, kl = w * flow_sq, w * kl flow_term, prob_term = flow_sq.mean(), kl.mean() loss = flow_term + prob_term if return_components: diff --git a/packages/viscy-utils/tests/test_seg_consistency_loss.py b/packages/viscy-utils/tests/test_seg_consistency_loss.py index fd28ab435..ada1c6230 100644 --- a/packages/viscy-utils/tests/test_seg_consistency_loss.py +++ b/packages/viscy-utils/tests/test_seg_consistency_loss.py @@ -111,32 +111,25 @@ def test_segmenter_is_frozen_unregistered_and_stays_in_eval(): assert not any(p.requires_grad for p in seg.parameters()) -def test_max_samples_and_crop_size_shape_what_the_segmenter_sees(): +def test_crop_size_shapes_what_the_segmenter_sees(): target = _nuclei((4, 1, 6, 40, 40)) - loss_fn = SegConsistencyLoss(TinySegmenter(), crop_size=24, max_samples=3) + loss_fn = SegConsistencyLoss(TinySegmenter(), crop_size=24) pred2d, target2d = loss_fn.project(target, target) - assert pred2d.shape == target2d.shape == (3, 1, 24, 24) - - -def test_uncertainty_weight_lowers_the_loss_and_keeps_zero_at_target(): - target = _nuclei((2, 1, 6, 32, 32)) - pred = target + 0.3 * torch.randn_like(target) - plain = SegConsistencyLoss(TinySegmenter()) - weighted = SegConsistencyLoss(TinySegmenter(), uncertainty_weight=True, uncertainty_scale=0.01) - assert 0.0 < float(weighted(pred, target)) < float(plain(pred, target)) - assert float(weighted(target, target)) == pytest.approx(0.0, abs=1e-6) + 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="max_samples"): - SegConsistencyLoss(TinySegmenter(), max_samples=0) + 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(): From 5f81ff64727a0ac0c7e7da0f285dcc0631ad02c2 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Fri, 16 Oct 2026 15:08:06 -0700 Subject: [PATCH 08/20] perf(viscy-utils): build the Otsu histogram without a host sync torch.bincount sizes its output on the host, so each focus-plane pick synced the GPU. scatter_add_ into a fixed (rows, 256) histogram does not. Thresholds are bit-identical to the bincount version (new test), and a CUDA sync-debug test pins the no-sync claim; it trips on the old version. Co-Authored-By: Claude Opus 5.5 --- .../src/viscy_utils/losses/seg_consistency.py | 6 ++-- .../tests/test_seg_consistency_loss.py | 36 ++++++++++++++++++- 2 files changed, 38 insertions(+), 4 deletions(-) diff --git a/packages/viscy-utils/src/viscy_utils/losses/seg_consistency.py b/packages/viscy-utils/src/viscy_utils/losses/seg_consistency.py index bae9c9b01..bebcb0f38 100644 --- a/packages/viscy-utils/src/viscy_utils/losses/seg_consistency.py +++ b/packages/viscy-utils/src/viscy_utils/losses/seg_consistency.py @@ -55,12 +55,12 @@ def _otsu_thresholds(flat: Tensor) -> Tensor: Threshold per row, shape ``(R,)``: the center of the bin maximizing the between-class variance. """ - 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() + # 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) diff --git a/packages/viscy-utils/tests/test_seg_consistency_loss.py b/packages/viscy-utils/tests/test_seg_consistency_loss.py index ada1c6230..6217d02b1 100644 --- a/packages/viscy-utils/tests/test_seg_consistency_loss.py +++ b/packages/viscy-utils/tests/test_seg_consistency_loss.py @@ -5,7 +5,7 @@ from torch import Tensor, nn from viscy_utils.losses import SegConsistencyLoss -from viscy_utils.losses.seg_consistency import _binary_kl, nucleus_area_planes +from viscy_utils.losses.seg_consistency import _OTSU_BINS, _binary_kl, _otsu_thresholds, nucleus_area_planes class TinySegmenter(nn.Module): @@ -135,3 +135,37 @@ def test_rejects_bad_arguments_and_shapes(): 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") From 88e892d26db2edba58c598520c570101c210cd3b Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Fri, 16 Oct 2026 15:11:12 -0700 Subject: [PATCH 09/20] refactor(dynacell): simplify CellposeDense normalize99 is hard-coded to 1/99 with float quantiles (no per-call host-to-device tensor), the taper uses BSIZE, checkpointing keys on the tiles' requires_grad alone, and the train() override and requires_grad_ call go: SegConsistencyLoss freezes its segmenter, and the cpsam test now checks that it stays in eval with no grads inside the loss. The net is still put in eval once at construction: cellpose builds it in train mode, whose flows differ from stock inference by up to 1.7 (parity test). Co-Authored-By: Claude Opus 5.5 --- .../dynacell/src/dynacell/cellpose_dense.py | 29 +++++++------------ .../dynacell/tests/test_cellpose_dense.py | 4 +++ 2 files changed, 15 insertions(+), 18 deletions(-) diff --git a/applications/dynacell/src/dynacell/cellpose_dense.py b/applications/dynacell/src/dynacell/cellpose_dense.py index 256818dad..9fce67978 100644 --- a/applications/dynacell/src/dynacell/cellpose_dense.py +++ b/applications/dynacell/src/dynacell/cellpose_dense.py @@ -27,26 +27,23 @@ __all__ = ["CellposeDense", "normalize99_diff"] -def normalize99_diff(x: Tensor, lower: float = 1.0, upper: float = 99.0) -> Tensor: +def normalize99_diff(x: Tensor) -> Tensor: """Apply cellpose ``normalize99`` per image, with the percentiles under ``no_grad``. Parameters ---------- x : Tensor Images, shape ``(B, ...)``. - lower, upper : float - Percentiles mapped to 0 and 1. Returns ------- Tensor - ``(x - p_lower) / (p_upper - p_lower)`` per image, unclipped; 0 where the - percentile range is below 1e-3 (cellpose's flat-image rule). + ``(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() - q = torch.quantile(flat, torch.tensor([lower / 100, upper / 100], device=x.device), dim=1) - lo, hi = q[0], q[1] + 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))) @@ -61,8 +58,10 @@ class CellposeDense(nn.Module): 256-px tiles at ``tile_overlap=0.1`` and blended with cellpose's taper mask. No diameter rescale, matching the eval's ``diameter=None``. - The network always stays in eval mode with ``requires_grad=False``; ``train()`` - is a no-op. It is loaded on the CPU and follows the module's ``.to(device)``. + 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). Keeping it frozen afterwards (eval + mode, no parameter gradients) is the caller's job; ``SegConsistencyLoss`` does it. Parameters ---------- @@ -93,15 +92,9 @@ def __init__(self, model_name: str, use_bfloat16: bool = True, checkpoint_tiles: 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 - self.net.requires_grad_(False) + self.net = cpm.net.eval() self.model_name = model_name self.checkpoint_tiles = checkpoint_tiles - self.train(False) - - def train(self, mode: bool = True) -> "CellposeDense": - """Keep the frozen network in eval mode whatever ``mode`` is.""" - return super().train(False) def _net_flows(self, tiles: Tensor) -> Tensor: """Run the network on ``(N, 3, 256, 256)`` tiles; return its ``(N, 3, 256, 256)`` output.""" @@ -109,7 +102,7 @@ def _net_flows(self, tiles: Tensor) -> Tensor: def _taper(self, device: torch.device) -> Tensor: """Return cellpose's ``transforms._taper_mask`` for one tile.""" - bsize = max(224, self.BSIZE) + 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)) @@ -167,7 +160,7 @@ def dense_normalized(self, xn: Tensor, batch_size: int = 8) -> dict[str, Tensor] xs = np.linspace(0, lx - bs, nx).astype(int) coords = [(int(a), int(c)) for a in ys for c in xs] 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 torch.is_grad_enabled() and tiles.requires_grad + 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): diff --git a/applications/dynacell/tests/test_cellpose_dense.py b/applications/dynacell/tests/test_cellpose_dense.py index aa08c7f73..a4769e3d8 100644 --- a/applications/dynacell/tests/test_cellpose_dense.py +++ b/applications/dynacell/tests/test_cellpose_dense.py @@ -126,6 +126,8 @@ def test_seg_consistency_through_cpsam_trains_a_backbone(cpsam_bf16): 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) @@ -136,6 +138,8 @@ def test_seg_consistency_through_cpsam_trains_a_backbone(cpsam_bf16): 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): From 671a3cf94ade1d45555bb0a90d56ba96ca80474a Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Fri, 16 Oct 2026 15:11:49 -0700 Subject: [PATCH 10/20] test(dynacell): pin SegConsistencyLoss's focus port to the eval nucleus_area_planes / _slab_mip against dynacell.evaluation.focus nucleus_area_plane / slab_mip on 50 blob volumes (Z = 1..40, a constant volume, random scale and offset): planes and slabs (halfwidth 0-2) equal. Co-Authored-By: Claude Opus 5.5 --- .../tests/test_seg_consistency_focus.py | 37 +++++++++++++++++++ 1 file changed, 37 insertions(+) create mode 100644 applications/dynacell/tests/test_seg_consistency_focus.py 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 From 0043b4e4d437b55df3df492533fe562f62776554 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Fri, 16 Oct 2026 15:13:20 -0700 Subject: [PATCH 11/20] refactor(dynacell): drop the segcons null-weight placeholder path All three weights are calibrated, so SEG_CONSISTENCY_WEIGHTS is typed float and the placeholder header branch goes. The generator test no longer rebuilds each arm with a null weight (3 extra ViT-L loads); the engine refusal is covered by test_seg_consistency_training. --check: 467 leaves, 0 stale, no leaf changes. Co-Authored-By: Claude Opus 5.5 --- .../dynacell/tools/generate_spotlight_v2_leaves.py | 8 +------- .../dynacell/tools/generate_spotlight_v2_leaves_test.py | 9 +++------ 2 files changed, 4 insertions(+), 13 deletions(-) diff --git a/applications/dynacell/tools/generate_spotlight_v2_leaves.py b/applications/dynacell/tools/generate_spotlight_v2_leaves.py index b91b59c47..d1a3da2c0 100644 --- a/applications/dynacell/tools/generate_spotlight_v2_leaves.py +++ b/applications/dynacell/tools/generate_spotlight_v2_leaves.py @@ -285,8 +285,7 @@ # UNeXt2-3D 0.0501 [0.0245-0.0883], pix2pix2d (generator_ema, base = lambda_l1 * L1) 0.5707 # [0.321-0.669]. Source: experiments/2026-09-24_spotlight-v2/calibration/results/ # nucleus__{fnet3d_vscyto3daug_v2,fcmae_vscyto3d_scratch_v2,pix2pix2d_unetvit}__segcons{.csv,__meta.json}. -# None would emit `seg_consistency_weight: null`, which the engines refuse. -SEG_CONSISTENCY_WEIGHTS: dict[tuple[str, str], float | None] = { +SEG_CONSISTENCY_WEIGHTS: dict[tuple[str, str], float] = { ("nucleus", "fnet3d_vscyto3daug_segcons"): 0.48, ("nucleus", "fcmae_vscyto3d_scratch_segcons"): 0.050, ("nucleus", "pix2pix2d_unetvit_segcons"): 0.57, @@ -833,11 +832,6 @@ def _header(arm: Arm, organelle: str, baseline_leaf: Path, kind: str) -> str: ] if kind == "fit": lines.append(f"# Recipe delta vs that leaf: {_DESCRIPTION[arm.suffix]}.") - if arm.suffix == "segcons" and SEG_CONSISTENCY_WEIGHTS[(organelle, arm.model)] is None: - lines.append( - "# seg_consistency_weight is a PLACEHOLDER (null): calibrate it and set SEG_CONSISTENCY_WEIGHTS" - " in the generator before submitting; the engine refuses to build with it unset." - ) if _long_wall(arm): lines.append( "# Wall: hardware_4gpu_long (7 d); 200 epochs at the measured ~2.0 ep/h is ~100 h > 4 d" diff --git a/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py b/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py index 77552ec7d..ea2deecb0 100644 --- a/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py +++ b/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py @@ -369,19 +369,16 @@ def test_segcons_ctrl_builds_from_the_real_fork_checkpoint(baseline: str) -> Non @pytest.mark.parametrize("baseline", SEGCONS_BASELINES) -def test_segcons_builds_with_its_calibrated_weight_and_refuses_a_null_one(baseline: str) -> None: - """The emitted calibrated weight builds the cpsam term; a ``null`` weight fails at build.""" +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).""" 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 weight is not None and model_cfg["init_args"]["seg_consistency_weight"] == weight + 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 - model_cfg["init_args"]["seg_consistency_weight"] = None - with pytest.raises(ValueError, match="placeholder"): - _instantiate_model(model_cfg) term = module.seg_consistency assert isinstance(term, SegConsistencyLoss) assert type(term.segmenter).__name__ == "CellposeDense" and term.segmenter.model_name == "cpsam" From 7ae9099e73909f5554c589a6e56907723a47cec5 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Fri, 16 Oct 2026 15:13:49 -0700 Subject: [PATCH 12/20] test(dynacell): skip cellpose tests without cellpose; fix a docstring cellpose is an optional dynacell extra: test_cellpose_dense importorskips it at module top, and the generator's cpsam build test does so in-test so the rest of that module still runs. The gradient test's docstring now says what it asserts. Co-Authored-By: Claude Opus 5.5 --- applications/dynacell/tests/test_cellpose_dense.py | 13 ++++++++----- .../dynacell/tests/test_seg_consistency_training.py | 2 +- .../tools/generate_spotlight_v2_leaves_test.py | 2 +- 3 files changed, 10 insertions(+), 7 deletions(-) diff --git a/applications/dynacell/tests/test_cellpose_dense.py b/applications/dynacell/tests/test_cellpose_dense.py index a4769e3d8..e67310a4d 100644 --- a/applications/dynacell/tests/test_cellpose_dense.py +++ b/applications/dynacell/tests/test_cellpose_dense.py @@ -7,12 +7,15 @@ import numpy as np import pytest import torch -from cellpose import models as cp_models -from cellpose import transforms as cp_transforms -from dynacell.cellpose_dense import CellposeDense, normalize99_diff -from dynacell.engine import DynacellUNet -from viscy_utils.losses import SegConsistencyLoss +pytest.importorskip("cellpose") + +from cellpose import models as cp_models # noqa: E402 +from cellpose import transforms as cp_transforms # 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(), diff --git a/applications/dynacell/tests/test_seg_consistency_training.py b/applications/dynacell/tests/test_seg_consistency_training.py index 99b9719da..58521a36c 100644 --- a/applications/dynacell/tests/test_seg_consistency_training.py +++ b/applications/dynacell/tests/test_seg_consistency_training.py @@ -117,7 +117,7 @@ def test_unet_seg_consistency_fast_dev_run(tmp_path, tiny_hcs_zarr, case): def test_unet_training_step_gradient_includes_the_term(): - """The weighted term changes the backbone gradient by exactly weight x its own gradient.""" + """A zero-weight term leaves the backbone gradient equal to the no-term one; a nonzero weight changes it.""" seed_everything(0) batch = {"source": torch.randn(2, 1, 4, 32, 32), "target": torch.randn(2, 1, 4, 32, 32)} baseline = DynacellUNet(architecture="FNet3D", model_config=FNET_CONFIG, log_batches_per_epoch=0) diff --git a/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py b/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py index ea2deecb0..c82c25908 100644 --- a/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py +++ b/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py @@ -9,7 +9,6 @@ import pytest import torch import yaml -from cellpose import models as cp_models from generate_spotlight_v2_leaves import ( _FIT_RENAMES, ARMS, @@ -371,6 +370,7 @@ def test_segcons_ctrl_builds_from_the_real_fork_checkpoint(baseline: str) -> Non @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"]) From 3314fa34ffc5701baca0f590114cd5202fbe4fd1 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Fri, 16 Oct 2026 15:30:42 -0700 Subject: [PATCH 13/20] fix(dynacell): crop the segcons projection to 240 px, one cellpose tile cellpose get_pad_yx pads a 256-px image to 272, which make_tiles cuts into 2 x 2 tiles: 4x the segmenter compute per sample. A 240-px crop pads to exactly 256, one tile. SEG_CONSISTENCY_ARGS and the loss default move to 240; CellposeDense.tiling exposes the padding/tile layout, and a test pins one tile at the configured crop (and 4 at 256). Co-Authored-By: Claude Opus 5.5 --- .../ipsc_confocal/train.yml | 4 +-- .../ipsc_confocal/train.yml | 4 +-- .../ipsc_confocal/train.yml | 4 +-- .../dynacell/src/dynacell/cellpose_dense.py | 31 ++++++++++++++----- .../dynacell/tests/test_cellpose_dense.py | 9 ++++++ .../tools/generate_spotlight_v2_leaves.py | 10 +++--- .../generate_spotlight_v2_leaves_test.py | 2 +- .../src/viscy_utils/losses/seg_consistency.py | 5 +-- 8 files changed, 49 insertions(+), 20 deletions(-) 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 index 8a15d152e..401807f1f 100644 --- 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 @@ -1,7 +1,7 @@ # 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, 256-px crop, per-tile activation checkpointing) + seg_consistency_weight (gradient-norm calibrated at the fork checkpoint). +# 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 @@ -67,5 +67,5 @@ model: model_name: cpsam checkpoint_tiles: 1 slab_halfwidth: 1 - crop_size: 256 + crop_size: 240 seg_consistency_weight: 0.05 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 index 10c2a61ef..94ae13b2a 100644 --- 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 @@ -1,7 +1,7 @@ # 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, 256-px crop, per-tile activation checkpointing) + seg_consistency_weight (gradient-norm calibrated at the fork checkpoint). +# 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 @@ -95,5 +95,5 @@ model: model_name: cpsam checkpoint_tiles: 1 slab_halfwidth: 1 - crop_size: 256 + crop_size: 240 seg_consistency_weight: 0.48 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 index 7653d7b83..ca204390f 100644 --- 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 @@ -1,7 +1,7 @@ # 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, 256-px crop, per-tile activation checkpointing) + seg_consistency_weight (gradient-norm calibrated at the fork checkpoint). +# 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 @@ -78,5 +78,5 @@ model: model_name: cpsam checkpoint_tiles: 1 slab_halfwidth: 1 - crop_size: 256 + crop_size: 240 seg_consistency_weight: 0.57 diff --git a/applications/dynacell/src/dynacell/cellpose_dense.py b/applications/dynacell/src/dynacell/cellpose_dense.py index 9fce67978..91a97faae 100644 --- a/applications/dynacell/src/dynacell/cellpose_dense.py +++ b/applications/dynacell/src/dynacell/cellpose_dense.py @@ -96,6 +96,28 @@ def __init__(self, model_name: str, use_bfloat16: bool = True, checkpoint_tiles: self.model_name = model_name self.checkpoint_tiles = checkpoint_tiles + @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] @@ -150,15 +172,10 @@ def dense_normalized(self, xn: Tensor, batch_size: int = 8) -> dict[str, Tensor] net_dtype = next(self.net.parameters()).dtype xn = torch.cat([xn, torch.zeros_like(xn), torch.zeros_like(xn)], dim=1) bs = self.BSIZE - yp1, yp2, xp1, xp2 = cp_transforms.get_pad_yx(ly0, lx0, min_size=(bs, bs)) + pads, coords = self.tiling(ly0, lx0) + yp1, yp2, xp1, xp2 = pads xn = F.pad(xn, (xp1, xp2, yp1, yp2)) ly, lx = xn.shape[-2:] - # make_tiles (non-augmented); ly, lx >= bs after get_pad_yx. - ny = 1 if ly <= bs else int(np.ceil((1.0 + 2 * self.TILE_OVERLAP) * ly / bs)) - nx = 1 if lx <= bs else int(np.ceil((1.0 + 2 * self.TILE_OVERLAP) * lx / bs)) - ys = np.linspace(0, ly - bs, ny).astype(int) - xs = np.linspace(0, lx - bs, nx).astype(int) - coords = [(int(a), int(c)) for a in ys for c in xs] 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 diff --git a/applications/dynacell/tests/test_cellpose_dense.py b/applications/dynacell/tests/test_cellpose_dense.py index e67310a4d..de80c2508 100644 --- a/applications/dynacell/tests/test_cellpose_dense.py +++ b/applications/dynacell/tests/test_cellpose_dense.py @@ -12,6 +12,7 @@ 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 @@ -39,6 +40,14 @@ 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 + + def test_refuses_the_evaluation_segmenter_and_missing_models(): with pytest.raises(ValueError, match="evaluation segmenter"): CellposeDense("cpdino") diff --git a/applications/dynacell/tools/generate_spotlight_v2_leaves.py b/applications/dynacell/tools/generate_spotlight_v2_leaves.py index d1a3da2c0..661bd536b 100644 --- a/applications/dynacell/tools/generate_spotlight_v2_leaves.py +++ b/applications/dynacell/tools/generate_spotlight_v2_leaves.py @@ -267,7 +267,9 @@ } # 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 -# 256-px tile (the geometry the G4 bench timed); see viscy_utils.losses.seg_consistency. +# 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). @@ -277,7 +279,7 @@ "init_args": {"model_name": "cpsam", "checkpoint_tiles": 1}, }, "slab_halfwidth": 1, - "crop_size": 256, + "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 (calibrate_variant.py --variants segcons --ckpt @@ -527,8 +529,8 @@ def _long_wall(arm: Arm) -> bool: "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, 256-px crop, per-tile activation checkpointing) + seg_consistency_weight (gradient-norm " - "calibrated at the fork checkpoint)", + "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. diff --git a/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py b/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py index c82c25908..b3aafcd94 100644 --- a/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py +++ b/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py @@ -383,5 +383,5 @@ def test_segcons_builds_with_its_calibrated_weight(baseline: str) -> None: 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, 256) + 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/seg_consistency.py b/packages/viscy-utils/src/viscy_utils/losses/seg_consistency.py index bebcb0f38..58e658c10 100644 --- a/packages/viscy-utils/src/viscy_utils/losses/seg_consistency.py +++ b/packages/viscy-utils/src/viscy_utils/losses/seg_consistency.py @@ -21,7 +21,8 @@ (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`` (one 256-px Cellpose tile). +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 @@ -154,7 +155,7 @@ def __init__( self, segmenter: nn.Module, slab_halfwidth: int = 1, - crop_size: int | None = 256, + crop_size: int | None = 240, ) -> None: super().__init__() if slab_halfwidth < 0: From 1141c54887bd558bcf0a4a8bf39f63055dfbdf52 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Fri, 16 Oct 2026 15:31:44 -0700 Subject: [PATCH 14/20] fix(dynacell): shorten the pix2pix2d fork's EMA half-life Both pix2pix2d arms predict from generator_ema, which starts as the fork checkpoint's shadow. At the baseline's ema_kimg 10, the 1000-step x 16-sample fine-tune leaves 0.5^(16000/10000) = 33% of that start shadow in it; ema_kimg 2 leaves 0.5^8 = 0.4%. Set for both arms of the pair. Co-Authored-By: Claude Opus 5.5 --- .../pix2pix2d_unetvit_segcons/ipsc_confocal/train.yml | 1 + .../pix2pix2d_unetvit_segcons_ctrl/ipsc_confocal/train.yml | 1 + applications/dynacell/tools/generate_spotlight_v2_leaves.py | 6 +++++- 3 files changed, 7 insertions(+), 1 deletion(-) 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 index ca204390f..e06a70474 100644 --- 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 @@ -69,6 +69,7 @@ model: 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: 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 index 2fb76eeff..3bde11e21 100644 --- 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 @@ -69,3 +69,4 @@ model: lr_g: 2.0e-05 lr_d: 2.0e-05 warmup_steps: 50 + ema_kimg: 2.0 diff --git a/applications/dynacell/tools/generate_spotlight_v2_leaves.py b/applications/dynacell/tools/generate_spotlight_v2_leaves.py index 661bd536b..f67a7bdc7 100644 --- a/applications/dynacell/tools/generate_spotlight_v2_leaves.py +++ b/applications/dynacell/tools/generate_spotlight_v2_leaves.py @@ -257,12 +257,16 @@ # 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}, + "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 From 7d52315c3254e032172b32689dc6fd8f3a2a13ea Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Fri, 16 Oct 2026 15:33:29 -0700 Subject: [PATCH 15/20] fix(dynacell): keep CellposeDense in eval under train() Restores the train() override next to the construction-time eval(): the network's train mode changes its outputs (flows off by up to 1.7 vs stock inference), so a parent's train() must not reach it if the wrapper is ever registered as a submodule. Tested through a ModuleDict parent. Co-Authored-By: Claude Opus 5.5 --- applications/dynacell/src/dynacell/cellpose_dense.py | 9 +++++++-- applications/dynacell/tests/test_cellpose_dense.py | 8 ++++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/applications/dynacell/src/dynacell/cellpose_dense.py b/applications/dynacell/src/dynacell/cellpose_dense.py index 91a97faae..b5c1f7f33 100644 --- a/applications/dynacell/src/dynacell/cellpose_dense.py +++ b/applications/dynacell/src/dynacell/cellpose_dense.py @@ -60,8 +60,9 @@ class CellposeDense(nn.Module): 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). Keeping it frozen afterwards (eval - mode, no parameter gradients) is the caller's job; ``SegConsistencyLoss`` does it. + 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 ---------- @@ -96,6 +97,10 @@ def __init__(self, model_name: str, use_bfloat16: bool = True, checkpoint_tiles: 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. diff --git a/applications/dynacell/tests/test_cellpose_dense.py b/applications/dynacell/tests/test_cellpose_dense.py index de80c2508..40b36ab7a 100644 --- a/applications/dynacell/tests/test_cellpose_dense.py +++ b/applications/dynacell/tests/test_cellpose_dense.py @@ -48,6 +48,14 @@ def test_configured_crop_is_one_cellpose_tile(): 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()) + + def test_refuses_the_evaluation_segmenter_and_missing_models(): with pytest.raises(ValueError, match="evaluation segmenter"): CellposeDense("cpdino") From 40af11d1b5421a1ff37bcf8c1d7f8c2aff5ac488 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Fri, 16 Oct 2026 15:35:21 -0700 Subject: [PATCH 16/20] fix(dynacell): reject a non-positive seg_consistency_weight A zero weight builds the segmenter and pays its cost for no term; turning the term off means dropping seg_consistency. The gradient test now checks the step's backbone gradient equals the base gradient plus weight x the term's own gradient, instead of comparing against a zero-weight arm. Co-Authored-By: Claude Opus 5.5 --- applications/dynacell/src/dynacell/engine.py | 4 +-- .../tests/test_seg_consistency_training.py | 36 ++++++++++++------- 2 files changed, 26 insertions(+), 14 deletions(-) diff --git a/applications/dynacell/src/dynacell/engine.py b/applications/dynacell/src/dynacell/engine.py index 70618e199..3f735133d 100644 --- a/applications/dynacell/src/dynacell/engine.py +++ b/applications/dynacell/src/dynacell/engine.py @@ -141,8 +141,8 @@ def _check_seg_consistency_args( "seg_consistency_weight is unset (placeholder): calibrate it as the median gradient-norm ratio " "||grad base|| / ||grad seg_consistency|| at the starting checkpoint and set it in the leaf." ) - if weight < 0: - raise ValueError(f"seg_consistency_weight must be >= 0, got {weight}") + 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.") diff --git a/applications/dynacell/tests/test_seg_consistency_training.py b/applications/dynacell/tests/test_seg_consistency_training.py index 58521a36c..24136fb33 100644 --- a/applications/dynacell/tests/test_seg_consistency_training.py +++ b/applications/dynacell/tests/test_seg_consistency_training.py @@ -117,20 +117,31 @@ def test_unet_seg_consistency_fast_dev_run(tmp_path, tiny_hcs_zarr, case): def test_unet_training_step_gradient_includes_the_term(): - """A zero-weight term leaves the backbone gradient equal to the no-term one; a nonzero weight changes it.""" + """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) - grads = {} - for weight in (None, 0.0, 2.0): - kwargs = {} if weight is None else {"seg_consistency": _segcons(), "seg_consistency_weight": weight} - module = DynacellUNet(architecture="FNet3D", model_config=FNET_CONFIG, log_batches_per_epoch=0, **kwargs) - module.load_state_dict(baseline.state_dict()) + 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 - module.training_step(batch, 0).backward() - grads[weight] = torch.cat([p.grad.flatten() for p in module.model.parameters()]) - torch.testing.assert_close(grads[None], grads[0.0]) - assert not torch.allclose(grads[None], grads[2.0]) + 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): @@ -198,8 +209,9 @@ def test_weight_is_required_with_the_term_and_seg_aux_is_rejected_with_it(family 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) - with pytest.raises(ValueError, match=">= 0"): - cls(**kwargs, seg_consistency=_segcons(), 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) From 2691f980fa1a00a3f091b4e9515680cd41142b95 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Fri, 16 Oct 2026 15:37:09 -0700 Subject: [PATCH 17/20] docs(dynacell): drop the stale seg_consistency weight placeholder wording The weights are calibrated, so the engine's docstring and error describe a missing weight (no default) rather than a null leaf placeholder, and the generator docstring no longer mentions a null weight. Co-Authored-By: Claude Opus 5.5 --- applications/dynacell/src/dynacell/engine.py | 6 +++--- .../dynacell/tests/test_seg_consistency_training.py | 2 +- applications/dynacell/tools/generate_spotlight_v2_leaves.py | 2 +- 3 files changed, 5 insertions(+), 5 deletions(-) diff --git a/applications/dynacell/src/dynacell/engine.py b/applications/dynacell/src/dynacell/engine.py index 3f735133d..1eb75db8c 100644 --- a/applications/dynacell/src/dynacell/engine.py +++ b/applications/dynacell/src/dynacell/engine.py @@ -129,7 +129,7 @@ def _check_seg_consistency_args( The weight has no default: it is calibrated per cell (median over validation batches of ``||grad base|| / ||grad seg_consistency||`` at the starting - checkpoint), and a leaf that still carries the ``null`` placeholder must not train. + 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: @@ -138,8 +138,8 @@ def _check_seg_consistency_args( return if weight is None: raise ValueError( - "seg_consistency_weight is unset (placeholder): calibrate it as the median gradient-norm ratio " - "||grad base|| / ||grad seg_consistency|| at the starting checkpoint and set it in the leaf." + "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") diff --git a/applications/dynacell/tests/test_seg_consistency_training.py b/applications/dynacell/tests/test_seg_consistency_training.py index 24136fb33..cebd254ec 100644 --- a/applications/dynacell/tests/test_seg_consistency_training.py +++ b/applications/dynacell/tests/test_seg_consistency_training.py @@ -205,7 +205,7 @@ def test_weight_is_required_with_the_term_and_seg_aux_is_rejected_with_it(family else: kwargs = {"generator_config": VIT2D_CONFIG, "discriminator_config": DISC2D_CONFIG} cls = DynacellGAN - with pytest.raises(ValueError, match="placeholder"): + 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) diff --git a/applications/dynacell/tools/generate_spotlight_v2_leaves.py b/applications/dynacell/tools/generate_spotlight_v2_leaves.py index f67a7bdc7..cf0894888 100644 --- a/applications/dynacell/tools/generate_spotlight_v2_leaves.py +++ b/applications/dynacell/tools/generate_spotlight_v2_leaves.py @@ -89,7 +89,7 @@ ``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`; the engines refuse a ``null`` one). Both + (: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. From e722cff776f17fe7f347d32f24247b479eef1eb1 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Fri, 16 Oct 2026 15:38:14 -0700 Subject: [PATCH 18/20] test(dynacell): check every fork-checkpoint tensor in the control build The strict-load test compared only the first state_dict tensor; it now requires every tensor of the built control to equal the fork checkpoint's. Co-Authored-By: Claude Opus 5.5 --- .../dynacell/tools/generate_spotlight_v2_leaves_test.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py b/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py index b3aafcd94..73bd66fb5 100644 --- a/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py +++ b/applications/dynacell/tools/generate_spotlight_v2_leaves_test.py @@ -362,9 +362,10 @@ def test_segcons_ctrl_builds_from_the_real_fork_checkpoint(baseline: str) -> Non 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"] - assert module.state_dict().keys() == state.keys() - key = next(iter(state)) - torch.testing.assert_close(module.state_dict()[key], state[key]) + 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) From 031d75f23b8b7fe7f15cad0c1f96fd7e7d5de69c Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Fri, 16 Oct 2026 15:54:31 -0700 Subject: [PATCH 19/20] feat(dynacell): recalibrate the segcons weights at the 240-px crop Same protocol at the leaves' one-tile crop (results/*__segcons_crop240): FNet-3D 0.37 [0.14-0.65], UNeXt2-3D 0.032 [0.024-0.066], pix2pix2d 0.19 [0.067-0.34], from 0.48 / 0.050 / 0.57 at the 4-tile 256-px crop. The plain-Dice instrument rows are unchanged to every printed digit. The comment records the convention's known limitation: segaux-style base (pix2pix2d: generator_ema, lambda_l1 * L1, no adversarial term) and no training-batch averaging, kept for comparability with the segaux weights. Co-Authored-By: Claude Opus 5.5 --- .../ipsc_confocal/train.yml | 2 +- .../ipsc_confocal/train.yml | 2 +- .../ipsc_confocal/train.yml | 2 +- .../tools/generate_spotlight_v2_leaves.py | 21 ++++++++++++------- 4 files changed, 16 insertions(+), 11 deletions(-) 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 index 401807f1f..c00b1ee30 100644 --- 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 @@ -68,4 +68,4 @@ model: checkpoint_tiles: 1 slab_halfwidth: 1 crop_size: 240 - seg_consistency_weight: 0.05 + seg_consistency_weight: 0.032 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 index 94ae13b2a..d043e890b 100644 --- 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 @@ -96,4 +96,4 @@ model: checkpoint_tiles: 1 slab_halfwidth: 1 crop_size: 240 - seg_consistency_weight: 0.48 + seg_consistency_weight: 0.37 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 index e06a70474..9b84b4839 100644 --- 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 @@ -80,4 +80,4 @@ model: checkpoint_tiles: 1 slab_halfwidth: 1 crop_size: 240 - seg_consistency_weight: 0.57 + seg_consistency_weight: 0.19 diff --git a/applications/dynacell/tools/generate_spotlight_v2_leaves.py b/applications/dynacell/tools/generate_spotlight_v2_leaves.py index cf0894888..c49d09571 100644 --- a/applications/dynacell/tools/generate_spotlight_v2_leaves.py +++ b/applications/dynacell/tools/generate_spotlight_v2_leaves.py @@ -286,15 +286,20 @@ "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 (calibrate_variant.py --variants segcons --ckpt -# , micro-batch 1, A40 2026-10-15), 2 significant figures [IQR]. FNet-3D 0.4806 [0.246-0.928], -# UNeXt2-3D 0.0501 [0.0245-0.0883], pix2pix2d (generator_ema, base = lambda_l1 * L1) 0.5707 -# [0.321-0.669]. Source: experiments/2026-09-24_spotlight-v2/calibration/results/ -# nucleus__{fnet3d_vscyto3daug_v2,fcmae_vscyto3d_scratch_v2,pix2pix2d_unetvit}__segcons{.csv,__meta.json}. +# ||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.48, - ("nucleus", "fcmae_vscyto3d_scratch_segcons"): 0.050, - ("nucleus", "pix2pix2d_unetvit_segcons"): 0.57, + ("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] From 797a85bc49e798bbfb078c7e0c56ee300197b7f7 Mon Sep 17 00:00:00 2001 From: Alexandr Kalinin Date: Fri, 16 Oct 2026 16:21:10 -0700 Subject: [PATCH 20/20] fix(dynacell): keep the global RNG unchanged when building CellposeDense CellposeModel random-initializes its ViT-L before loading the checkpoint, which advanced the global torch RNG between seed_everything and HCSDataModule's torch.randperm train/val split. A segcons arm therefore split its FOVs differently from its same-seed segcons_ctrl, confounding the comparison (measured: seed 42 randperm differs after construction). Construction now runs inside torch.random.fork_rng; a test pins the torch state (it fails on the previous code). Co-Authored-By: Claude Opus 5.5 Claude-Session: https://claude.ai/code/session_01Gwn9hwgT8uJSoWaCxCAm8J --- .../dynacell/src/dynacell/cellpose_dense.py | 8 +++++++- applications/dynacell/tests/test_cellpose_dense.py | 14 ++++++++++++++ 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/applications/dynacell/src/dynacell/cellpose_dense.py b/applications/dynacell/src/dynacell/cellpose_dense.py index b5c1f7f33..260ed9750 100644 --- a/applications/dynacell/src/dynacell/cellpose_dense.py +++ b/applications/dynacell/src/dynacell/cellpose_dense.py @@ -90,7 +90,13 @@ def __init__(self, model_name: str, use_bfloat16: bool = True, checkpoint_tiles: path = cp_models.MODEL_DIR / model_name if not path.is_file(): raise FileNotFoundError(f"cellpose model {model_name!r} not found at {path}") - cpm = cp_models.CellposeModel(device=torch.device("cpu"), pretrained_model=str(path), use_bfloat16=use_bfloat16) + # 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() diff --git a/applications/dynacell/tests/test_cellpose_dense.py b/applications/dynacell/tests/test_cellpose_dense.py index 40b36ab7a..aebb75c44 100644 --- a/applications/dynacell/tests/test_cellpose_dense.py +++ b/applications/dynacell/tests/test_cellpose_dense.py @@ -56,6 +56,20 @@ def test_train_keeps_the_network_in_eval(cpsam_bf16): 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")