Fix label transform resolution and batched affinity masks

Author: Donglai Wei

connectomics/config/schema/data.py

@@ -78,6 +78,7 @@ class LabelTransformConfig:
     allow_missing_keys: bool = False
     segment_id: Optional[List[int]] = None
     boundary_thickness: int = 1
+    resolution: Optional[List[float]] = None  # Forwarded into compatible label targets.
     targets: List[Any] = field(default_factory=list)
 
 

connectomics/models/loss/losses.py

@@ -31,6 +31,14 @@ def _reduce_weighted_tensor(
         return loss_tensor.mean()
 
     valid = weight > 0
+    if valid.shape != loss_tensor.shape:
+        try:
+            valid = torch.broadcast_to(valid, loss_tensor.shape)
+        except RuntimeError as e:
+            raise ValueError(
+                "Weight mask shape is not broadcastable to loss tensor shape: "
+                f"weight={tuple(weight.shape)}, loss={tuple(loss_tensor.shape)}"
+            ) from e
     if not torch.any(valid):
         return loss_tensor.new_tensor(0.0)
     return loss_tensor[valid].mean()

tests/unit/test_hydra_config.py

@@ -30,6 +30,7 @@
 from connectomics.config.schema import TestConfig as HydraTestConfig  # noqa: E402
 from connectomics.config.schema import TuneConfig  # noqa: E402
 from connectomics.data.augment.build import build_test_transforms  # noqa: E402
+from connectomics.data.process import create_label_transform_pipeline  # noqa: E402
 
 
 def test_default_config_creation():
@@ -467,6 +468,35 @@ def test_yaml_shared_profile_selectors(tmp_path):
     print("[OK]YAML shared profile selectors work")
 
 
+def test_label_transform_profile_can_inherit_top_level_resolution(tmp_path):
+    repo_root = Path(__file__).resolve().parents[2]
+    base_profiles = repo_root / "tutorials" / "bases" / "all_profiles.yaml"
+    config_yaml = tmp_path / "config.yaml"
+    config_yaml.write_text(
+        f"""
+_base_:
+  - {base_profiles}
+
+default:
+  data:
+    label_transform:
+      profile: label_affinity_9_sdt
+      resolution: [30, 6, 6]
+""".strip()
+    )
+
+    cfg = load_config(config_yaml)
+
+    targets = cfg.default.data.label_transform.targets
+    assert len(targets) == 2
+    assert targets[0]["name"] == "affinity"
+    assert targets[1]["name"] == "skeleton_aware_edt"
+    assert cfg.default.data.label_transform.resolution == [30.0, 6.0, 6.0]
+
+    pipeline = create_label_transform_pipeline(cfg.default.data.label_transform)
+    assert pipeline.task_specs[1]["kwargs"]["resolution"] == [30.0, 6.0, 6.0]
+
+
 def test_arch_profile_rejects_non_model_sections(tmp_path):
     """Arch profiles with invalid keys are rejected by OmegaConf schema merge."""
     base_yaml = tmp_path / "base.yaml"

tests/unit/test_loss_orchestrator.py

@@ -325,6 +325,41 @@ def test_standard_loss_supports_affinity_deepem_crop():
     assert weight[0, 2, :, :, 1].sum() > 0  # x=1 valid for channel 2
 
 
+def test_standard_loss_supports_batched_affinity_deepem_crop_with_weighted_bce():
+    cfg = _cfg(
+        losses=[
+            {
+                "weight": 1.0,
+            },
+        ]
+    )
+    cfg.data.label_transform.targets = [
+        {
+            "name": "affinity",
+            "kwargs": {
+                "offsets": ["1-0-0", "0-1-0", "0-0-1"],
+                "deepem_crop": True,
+            },
+        }
+    ]
+    orchestrator = LossOrchestrator(
+        cfg=cfg,
+        loss_functions=nn.ModuleList([WeightedBCEWithLogitsLoss(reduction="mean")]),
+        loss_weights=[1.0],
+        enable_nan_detection=False,
+        debug_on_nan=False,
+    )
+
+    outputs = torch.zeros(4, 3, 5, 5, 5)
+    labels = torch.zeros(4, 3, 5, 5, 5)
+
+    total_loss, loss_dict = orchestrator.compute_standard_loss(outputs, labels, stage="train")
+
+    assert torch.isfinite(total_loss)
+    assert "train_loss_total" in loss_dict
+    assert torch.isclose(total_loss, torch.log(torch.tensor(2.0)))
+
+
 def test_standard_loss_rejects_invalid_runtime_negative_channel_slice():
     weighted_loss = WeightAwareSpyLoss()
     orchestrator = LossOrchestrator(
@@ -630,6 +665,18 @@ def test_weighted_bce_returns_zero_when_no_valid_weights():
     assert torch.isclose(loss, torch.tensor(0.0))
 
 
+def test_weighted_bce_mean_supports_broadcast_batchless_weight_mask():
+    loss_fn = WeightedBCEWithLogitsLoss(reduction="mean")
+    pred = torch.zeros(4, 1, 1, 1, 2)
+    target = torch.tensor([[[[[1.0, 0.0]]]]]).expand_as(pred)
+    weight = torch.tensor([[[[[2.0, 0.0]]]]])
+
+    loss = loss_fn(pred, target, weight=weight)
+
+    expected = 2.0 * torch.log(torch.tensor(2.0))
+    assert torch.isclose(loss, expected)
+
+
 def test_multitask_single_scale_routes_class_index_and_dense_targets():
     ce_loss = CrossEntropyLossWrapperSpy()
     reg_loss = WeightAwareSpyLoss()