Merge pull request #211 from PytorchConnectomics/feat/malis-gt-passthrough

MalisLoss: pass eroded gt_seg through data pipeline (skip per-step CC, fix crop topology)

Author: donglaiw

connectomics/config/pipeline/config_io.py

@@ -153,8 +153,7 @@ def _raise_unconsumed_keys(yaml_conf: DictConfig) -> None:
         "(sibling of `monitor`, `inference`, `decoding`, `tune`)."
     ),
     "use_timestamp": (
-        "field removed. Train mode is always timestamped; "
-        "test/tune modes are never timestamped."
+        "field removed. Train mode is always timestamped; " "test/tune modes are never timestamped."
     ),
 }
 _MONITOR_CHECKPOINT_ROOTS = (
@@ -201,8 +200,7 @@ def _reject_inference_runtime_alias_paths(explicit_field_paths: set[str]) -> Non
             alias_path = f"{root}.{alias}"
             if any(_path_is_or_descendant(path, alias_path) for path in explicit_field_paths):
                 raise ValueError(
-                    f"`{alias_path}` was renamed. "
-                    f"Use `{root}.{canonical_tail}` instead."
+                    f"`{alias_path}` was renamed. " f"Use `{root}.{canonical_tail}` instead."
                 )
 
     for root in _MONITOR_CHECKPOINT_ROOTS:
@@ -212,8 +210,7 @@ def _reject_inference_runtime_alias_paths(explicit_field_paths: set[str]) -> Non
                 if replacement.startswith("field "):
                     raise ValueError(f"`{alias_path}` {replacement}")
                 raise ValueError(
-                    f"`{alias_path}` was renamed. "
-                    f"Use `{root}.{replacement}` instead."
+                    f"`{alias_path}` was renamed. " f"Use `{root}.{replacement}` instead."
                 )
 
     # tune.output:* sub-block hoisted to tune.save_*
@@ -506,7 +503,6 @@ def validate_config(cfg: Config) -> None:
     if cfg.model.out_channels <= 0:
         raise ValueError("model.out_channels must be positive")
     model_heads = getattr(cfg.model, "heads", None) or {}
-    inference_cfg = getattr(cfg, "inference", None)
     inference_head = get_inference_model_value(cfg, "head", None)
     images_cfg = getattr(getattr(getattr(cfg, "monitor", None), "logging", None), "images", None)
     visualization_head = getattr(images_cfg, "head", None) if images_cfg is not None else None
@@ -556,8 +552,8 @@ def validate_config(cfg: Config) -> None:
             missing = [h for h in inference_head_names if h not in model_heads]
             if missing:
                 raise ValueError(
-                    f"inference.model.head={inference_head_names} references unknown heads {missing}; "
-                    f"available: {sorted(model_heads.keys())}."
+                    f"inference.model.head={inference_head_names} references unknown heads "
+                    f"{missing}; available: {sorted(model_heads.keys())}."
                 )
         if (
             visualization_head is not None
@@ -909,6 +905,11 @@ def _resolve_split_paths(split_cfg):
         split_cfg.image = _combine_path(split_base, split_cfg.image)
         split_cfg.label = _combine_path(split_base, split_cfg.label)
         split_cfg.mask = _combine_path(split_base, split_cfg.mask)
+        split_json_resolved = _combine_path(split_base, split_cfg.json)
+        if isinstance(split_json_resolved, list):
+            split_cfg.json = split_json_resolved[0] if split_json_resolved else None
+        else:
+            split_cfg.json = split_json_resolved
 
     # Resolve inference/test paths from merged runtime cfg.data.
     if getattr(cfg.data, "test", None) is not None:

connectomics/config/schema/data.py

@@ -66,6 +66,7 @@ class LabelTransformConfig:
 
     normalize: bool = True  # Convert labels to 0-1 range
     erosion: int = 0  # Border erosion kernel half-size (0 = disabled, uses seg_widen_border)
+    emit_gt_seg: bool = False  # Emit post-augmentation/post-erosion segmentation for MALIS.
     skeleton_distance: SkeletonDistanceConfig = field(default_factory=SkeletonDistanceConfig)
     edge_mode: EdgeModeConfig = field(
         default_factory=EdgeModeConfig

connectomics/data/augmentation/build.py

@@ -15,6 +15,7 @@
     BorderPadd,
     CenterSpatialCropd,
     Compose,
+    CopyItemsd,
     Lambdad,
     OneOf,
     RandAdjustContrastd,
@@ -182,7 +183,7 @@ def _build_nnunet_preprocess_transform(keys, nnunet_pre_cfg, source_spacing):
 
 
 def build_train_transforms(
-    cfg: Config, keys: list[str] = None, skip_loading: bool = False
+    cfg: Config, keys: list[str] | None = None, skip_loading: bool = False
 ) -> Compose:
     """
     Build training transforms from Hydra config.
@@ -320,6 +321,8 @@ def build_train_transforms(
         # Apply instance erosion first if specified
         if hasattr(label_cfg, "erosion") and label_cfg.erosion > 0:
             transforms.append(SegErosionInstanced(keys=["label"], tsz_h=label_cfg.erosion))
+        if label_cfg.emit_gt_seg:
+            transforms.append(CopyItemsd(keys="label", names="gt_seg"))
 
         # Build label transform pipeline directly from label_transform config
         label_transform = create_label_transform_pipeline(label_cfg)
@@ -348,7 +351,7 @@ def build_train_transforms(
 
 
 def _build_eval_transforms_impl(
-    cfg: Config, mode: str = "val", keys: list[str] = None, skip_loading: bool = False
+    cfg: Config, mode: str = "val", keys: list[str] | None = None, skip_loading: bool = False
 ) -> Compose:
     """
     Internal implementation for building evaluation transforms (validation or test).
@@ -669,6 +672,8 @@ def _resolve_eval_split():
             # Apply instance erosion first if specified
             if hasattr(label_cfg, "erosion") and label_cfg.erosion > 0:
                 transforms.append(SegErosionInstanced(keys=["label"], tsz_h=label_cfg.erosion))
+            if label_cfg.emit_gt_seg:
+                transforms.append(CopyItemsd(keys="label", names="gt_seg"))
 
             # Build label transform pipeline directly from label_transform config
             label_transform = create_label_transform_pipeline(label_cfg)
@@ -695,7 +700,7 @@ def _resolve_eval_split():
 
 
 def build_val_transforms(
-    cfg: Config, keys: list[str] = None, skip_loading: bool = False
+    cfg: Config, keys: list[str] | None = None, skip_loading: bool = False
 ) -> Compose:
     """
     Build validation transforms from Hydra config.
@@ -711,7 +716,9 @@ def build_val_transforms(
     return _build_eval_transforms_impl(cfg, mode="val", keys=keys, skip_loading=skip_loading)
 
 
-def build_test_transforms(cfg: Config, keys: list[str] = None, mode: str = "test") -> Compose:
+def build_test_transforms(
+    cfg: Config, keys: list[str] | None = None, mode: str = "test"
+) -> Compose:
     """
     Build test/tune inference transforms from Hydra config.
 

connectomics/models/losses/malis.py

@@ -26,6 +26,16 @@ class MalisLoss(nn.Module):
     2D tensors are rejected explicitly because the vendored MALIS helpers operate
     on 3D affinity graphs by default. See ``lib/malis/INVESTIGATION.md`` for
     GPU MALIS candidates and algorithm-level speedup follow-ups.
+
+    Performance knobs (see ``docs/source/notes/malis.rst``):
+
+    - ``malis_crop_size`` — random sub-volume crop on each forward call.
+      ``64`` on a ``128^3`` patch gives ~4.6x measured step speedup vs
+      the full-volume baseline (slurm 2505814 vs 2487040).
+    - ``label_transform.emit_gt_seg: true`` (YAML, paired with this
+      loss) — passes the eroded GT segmentation in via ``gt_seg=...``,
+      skipping the per-step ``connected_components_affgraph`` call and
+      preserving global instance IDs when ``malis_crop_size`` is set.
     """
 
     def __init__(
@@ -68,6 +78,7 @@ def forward(
         pred: torch.Tensor,
         target: torch.Tensor,
         mask: torch.Tensor | None = None,
+        gt_seg: torch.Tensor | np.ndarray | None = None,
     ) -> torch.Tensor:
         """Compute MALIS-weighted squared affinity error.
 
@@ -83,21 +94,31 @@ def forward(
                 Masked-out edges are excluded from MALIS pass constraints and
                 zeroed before per-pass normalization, but the mask does not
                 change GT connected-component reconstruction.
+            gt_seg: Optional ground-truth segmentation with shape ``[B, Z, Y, X]``
+                or ``[B, 1, Z, Y, X]``. When supplied, MALIS uses these instance
+                labels directly instead of reconstructing components from
+                ``target`` affinities.
         """
         self._validate_inputs(pred, target)
 
         pred_aff = torch.sigmoid(pred) if self.sigmoid else pred
         target_aff = target.to(device=pred.device, dtype=pred_aff.dtype)
         mask_aff = None if mask is None else self._prepare_mask(mask, pred_aff)
-        pred_aff, target_aff, mask_aff = self._apply_crop_if_configured(
+        gt_seg_tensor = self._prepare_gt_seg(gt_seg, pred_aff)
+        pred_aff, target_aff, mask_aff, gt_seg_tensor = self._apply_crop_if_configured(
             pred_aff,
             target_aff,
             mask_aff,
+            gt_seg_tensor,
         )
+        weight_kwargs = {}
+        if gt_seg_tensor is not None:
+            weight_kwargs["gt_seg"] = gt_seg_tensor.detach()
         weights = self._compute_malis_weights(
             pred_aff.detach(),
             target_aff.detach(),
             None if mask_aff is None else mask_aff.detach(),
+            **weight_kwargs,
         )
 
         edge_loss = (pred_aff - target_aff) ** 2
@@ -211,12 +232,36 @@ def _prepare_mask(self, mask: torch.Tensor, pred: torch.Tensor) -> torch.Tensor:
                 f"mask={tuple(mask.shape)}, pred={tuple(pred.shape)}."
             ) from e
 
+    def _prepare_gt_seg(
+        self,
+        gt_seg: torch.Tensor | np.ndarray | None,
+        pred: torch.Tensor,
+    ) -> torch.Tensor | None:
+        if gt_seg is None:
+            return None
+
+        gt_seg_tensor = torch.as_tensor(gt_seg, device=pred.device).detach()
+        if gt_seg_tensor.ndim == pred.ndim and gt_seg_tensor.shape[1] == 1:
+            gt_seg_tensor = gt_seg_tensor.squeeze(1)
+        elif gt_seg_tensor.ndim == pred.ndim - 2 and pred.shape[0] == 1:
+            gt_seg_tensor = gt_seg_tensor.unsqueeze(0)
+
+        expected_shape = (pred.shape[0],) + tuple(pred.shape[-3:])
+        if tuple(gt_seg_tensor.shape) != expected_shape:
+            raise ValueError(
+                "MalisLoss gt_seg must have shape [B, Z, Y, X] or [B, 1, Z, Y, X] "
+                f"matching pred spatial dims; got gt_seg={tuple(gt_seg_tensor.shape)}, "
+                f"expected={expected_shape}."
+            )
+        return gt_seg_tensor.contiguous()
+
     def _apply_crop_if_configured(
         self,
         pred_aff: torch.Tensor,
         target_aff: torch.Tensor,
         mask_aff: torch.Tensor | None,
-    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
+        gt_seg: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None, torch.Tensor | None]:
         """Apply the configured random sub-volume crop, if any.
 
         Offset sampling stays on CPU. The returned tensors are contiguous copies
@@ -226,11 +271,11 @@ def _apply_crop_if_configured(
         crop=64, and fp16, pred + target + mask copies are about 9 MiB before
         overhead.
 
-        Returns ``(pred_cropped, target_cropped, mask_cropped)``. If no crop is
-        configured the inputs are returned unchanged.
+        Returns ``(pred_cropped, target_cropped, mask_cropped, gt_seg_cropped)``.
+        If no crop is configured the inputs are returned unchanged.
         """
         if self.malis_crop_size is None:
-            return pred_aff, target_aff, mask_aff
+            return pred_aff, target_aff, mask_aff, gt_seg
 
         k_z, k_y, k_x = self.malis_crop_size
         z_dim, y_dim, x_dim = pred_aff.shape[-3:]
@@ -253,29 +298,40 @@ def _apply_crop_if_configured(
             if mask_aff is None
             else mask_aff.narrow(-3, z0, k_z).narrow(-2, y0, k_y).narrow(-1, x0, k_x).contiguous()
         )
-        return pred_c, target_c, mask_c
+        gt_seg_c = (
+            None
+            if gt_seg is None
+            else gt_seg.narrow(-3, z0, k_z).narrow(-2, y0, k_y).narrow(-1, x0, k_x).contiguous()
+        )
+        return pred_c, target_c, mask_c, gt_seg_c
 
     def _compute_malis_weights(
         self,
         pred_aff: torch.Tensor,
         target_aff: torch.Tensor,
         mask: torch.Tensor | None = None,
+        *,
+        gt_seg: torch.Tensor | None = None,
     ) -> torch.Tensor:
         pred_np = pred_aff.to(dtype=torch.float32).cpu().numpy()
         target_np = target_aff.to(dtype=torch.float32).cpu().numpy()
         mask_np = None if mask is None else mask.to(dtype=torch.float32).cpu().numpy()
+        gt_seg_np = None if gt_seg is None else gt_seg.cpu().numpy()
         weights = np.empty_like(pred_np, dtype=np.float32)
         for batch_idx in range(pred_np.shape[0]):
             gt_affs = np.ascontiguousarray(target_np[batch_idx] > 0.5, dtype=np.int32)
             pred_sample = np.ascontiguousarray(pred_np[batch_idx], dtype=np.float32)
             mask_sample = None
             if mask_np is not None:
                 mask_sample = np.ascontiguousarray(mask_np[batch_idx] == 1, dtype=bool)
-            gt_seg, _ = _malis_lib.connected_components_affgraph(gt_affs, self.nhood)
+            if gt_seg_np is None:
+                gt_seg_sample, _ = _malis_lib.connected_components_affgraph(gt_affs, self.nhood)
+            else:
+                gt_seg_sample = np.ascontiguousarray(gt_seg_np[batch_idx], dtype=np.uint64)
             weights[batch_idx] = self._compute_sample_weights(
                 pred_sample,
                 gt_affs,
-                gt_seg,
+                gt_seg_sample,
                 mask_sample,
             )
 

connectomics/models/losses/metadata.py

@@ -19,6 +19,7 @@ class LossMetadata:
     call_kind: LossCallKind = "pred_target"  # pred_target | pred_only | pred_pred | unsupported
     target_kind: TargetKind = "dense"  # dense | class_index | none
     spatial_weight_arg: Optional[str] = None  # weight | mask | None
+    gt_seg_arg: Optional[str] = None  # gt_seg | None
 
 
 _LOSS_METADATA_BY_NAME = {
@@ -44,7 +45,7 @@ class LossMetadata:
     "SoftClDiceLoss": LossMetadata("SoftClDiceLoss", spatial_weight_arg="weight"),
     "WeightedMSELoss": LossMetadata("WeightedMSELoss", spatial_weight_arg="weight"),
     "WeightedMAELoss": LossMetadata("WeightedMAELoss", spatial_weight_arg="weight"),
-    "MalisLoss": LossMetadata("MalisLoss", spatial_weight_arg="mask"),
+    "MalisLoss": LossMetadata("MalisLoss", spatial_weight_arg="mask", gt_seg_arg="gt_seg"),
     # GAN is not compatible with the generic supervised orchestrator path
     "GANLoss": LossMetadata("GANLoss", call_kind="unsupported", target_kind="none"),
     # Regularization losses

connectomics/training/lightning/model.py

@@ -845,6 +845,7 @@ def _compute_loss(
         stage: str,
         mask: Optional[torch.Tensor] = None,
         target_mask: Optional[torch.Tensor] = None,
+        gt_seg: Optional[torch.Tensor] = None,
     ):
         """Compute loss handling both standard and deep supervision outputs."""
         loss_orchestrator = self._require_loss_orchestrator()
@@ -853,10 +854,10 @@ def _compute_loss(
         )
         if is_deep_supervision:
             return loss_orchestrator.compute_deep_supervision_loss(
-                outputs, labels, stage=stage, mask=mask, target_mask=target_mask
+                outputs, labels, stage=stage, mask=mask, target_mask=target_mask, gt_seg=gt_seg
             )
         return loss_orchestrator.compute_standard_loss(
-            outputs, labels, stage=stage, mask=mask, target_mask=target_mask
+            outputs, labels, stage=stage, mask=mask, target_mask=target_mask, gt_seg=gt_seg
         )
 
     def training_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> STEP_OUTPUT:
@@ -868,13 +869,14 @@ def training_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> STEP_
         # Binarize mask: (B, 1, D, H, W) float, 1 = valid, 0 = ignore
         mask = (raw_mask > 0).float() if raw_mask is not None else None
         target_mask = batch.get("label_mask", None)
+        gt_seg = batch.get("gt_seg", None)
 
         # Forward pass
         outputs = self(images)
 
         # Compute loss using the loss orchestrator
         total_loss, loss_dict = self._compute_loss(
-            outputs, labels, stage="train", mask=mask, target_mask=target_mask
+            outputs, labels, stage="train", mask=mask, target_mask=target_mask, gt_seg=gt_seg
         )
 
         # Keep full training curves in TensorBoard while avoiding console spam.
@@ -896,13 +898,14 @@ def validation_step(self, batch: Dict[str, torch.Tensor], batch_idx: int) -> STE
         raw_mask = batch.get("mask", None)
         mask = (raw_mask > 0).float() if raw_mask is not None else None
         target_mask = batch.get("label_mask", None)
+        gt_seg = batch.get("gt_seg", None)
 
         # Forward pass
         outputs = self(images)
 
         # Compute loss using the loss orchestrator
         total_loss, loss_dict = self._compute_loss(
-            outputs, labels, stage="val", mask=mask, target_mask=target_mask
+            outputs, labels, stage="val", mask=mask, target_mask=target_mask, gt_seg=gt_seg
         )
 
         # Compute evaluation metrics if enabled