hydra-lv decoding

Author: Donglai Wei

tutorials/FINETUNING_DESIGN.md .claude/refactor/eng/FINETUNING_DESIGN.mdtutorials/FINETUNING_DESIGN.md renamed to .claude/refactor/eng/FINETUNING_DESIGN.md

@@ -37,8 +37,9 @@
 """
 
 from __future__ import annotations
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, is_dataclass
 from typing import Dict, List, Optional, Tuple, Any, Union
+import inspect
 
 # Note: MISSING can be imported from omegaconf if needed for required fields
 
@@ -1269,26 +1270,14 @@ class Config:
     import torch
 
     if hasattr(torch, "serialization") and hasattr(torch.serialization, "add_safe_globals"):
-        torch.serialization.add_safe_globals(
-            [
-                ParameterConfig,
-                DecodingParameterSpace,
-                PostprocessingParameterSpace,
-                ParameterSpaceConfig,
-                # Core config dataclasses (for Lightning checkpoints)
-                Config,
-                SystemConfig,
-                SystemTrainingConfig,
-                SystemInferenceConfig,
-                ModelConfig,
-                DataConfig,
-                OptimizationConfig,
-                MonitorConfig,
-                InferenceConfig,
-                TestConfig,
-                TuneConfig,
-            ]
-        )
+        # Register every dataclass defined in this module so Lightning checkpoints
+        # can be loaded safely when torch.load defaults to weights_only=True.
+        safe_dataclasses = [
+            obj
+            for obj in globals().values()
+            if inspect.isclass(obj) and obj.__module__ == __name__ and is_dataclass(obj)
+        ]
+        torch.serialization.add_safe_globals(safe_dataclasses)
 except Exception:
     # Best-effort registration; ignore if torch not available at import time
     pass

connectomics/config/hydra_config.py

@@ -83,6 +83,10 @@ def __init__(
         num_tasks = len(cfg.model.multi_task_config) if hasattr(cfg.model, 'multi_task_config') and cfg.model.multi_task_config else 1
         self.loss_weighter = build_loss_weighter(cfg, num_tasks, model=self.model)
 
+        # Track multi-task configuration state for downstream logic/tests
+        self.multi_task_config = getattr(cfg.model, 'multi_task_config', None)
+        self.multi_task_enabled = bool(self.multi_task_config)
+
         # Enable inline NaN detection (can be disabled via config)
         self.enable_nan_detection = getattr(cfg.model, 'enable_nan_detection', True)
         self.debug_on_nan = getattr(cfg.model, 'debug_on_nan', True)

connectomics/training/lit/model.py

@@ -279,10 +279,10 @@ test:
   decoding:
     - name: decode_instance_binary_contour_distance
       kwargs:
-        binary_threshold: [0.9, 0.1]
-        contour_threshold: [0.8, 1.1]
-        distance_threshold: [0.5, 0.0]
-        min_seed_size: 8
+        binary_threshold: [0.5, 0.7]
+        contour_threshold: [0.8, 13.1]
+        distance_threshold: [0.5, -0.5]
+        min_seed_size: 4
 
   evaluation:
     enabled: true