fix optuna

Author: Donglai Wei

connectomics/config/auto_config.py

@@ -199,15 +199,15 @@ def plan(
             )
             result.planning_notes.append(
                 f"Estimated memory: {result.estimated_gpu_memory_gb:.2f} GB "
-                f"({result.estimated_gpu_memory_gb/gpu_memory_gb*100:.1f}% of GPU)"
+                f"({result.estimated_gpu_memory_gb / gpu_memory_gb * 100:.1f}% of GPU)"
             )
             result.planning_notes.append(f"Batch size: {batch_size}")
 
             # Gradient accumulation if batch size is very small
             if batch_size == 1:
                 result.accumulate_grad_batches = 4
                 result.planning_notes.append(
-                    f"Using gradient accumulation (4 batches) for effective batch_size=4"
+                    "Using gradient accumulation (4 batches) for effective batch_size=4"
                 )
 
         # Step 5: Determine num_workers
@@ -311,7 +311,7 @@ def print_plan(self, result: AutoPlanResult):
             print(f"  Available: {result.available_gpu_memory_gb:.2f} GB")
             print(
                 f"  Estimated Usage: {result.estimated_gpu_memory_gb:.2f} GB "
-                f"({result.estimated_gpu_memory_gb/result.available_gpu_memory_gb*100:.1f}%)"
+                f"({result.estimated_gpu_memory_gb / result.available_gpu_memory_gb * 100:.1f}%)"
             )
             print(f"  Per Sample: {result.gpu_memory_per_sample_gb:.2f} GB")
             print()
@@ -450,7 +450,6 @@ def auto_plan_config(
 if __name__ == "__main__":
     # Test auto planning
     from connectomics.config import Config
-    from omegaconf import OmegaConf
 
     # Create test config
     cfg = OmegaConf.structured(Config())

connectomics/config/hydra_config.py

@@ -150,15 +150,15 @@ class ModelConfig:
     activation: str = "relu"
 
     # UNet-specific parameters (MONAI UNet)
-    spatial_dims: int = (
-        3  # Spatial dimensions: 2 for 2D, 3 for 3D (auto-inferred from input_size length, not used directly)
-    )
+    spatial_dims: int = 3  # Spatial dimensions: 2 for 2D, 3 for 3D
+    # (auto-inferred from input_size length, not used directly)
     num_res_units: int = 2  # Number of residual units per block
     kernel_size: int = 3  # Convolution kernel size
     strides: Optional[List[int]] = None  # Downsampling strides (e.g., [2, 2, 2, 2] for 4 levels)
     act: str = "relu"  # Activation function: 'relu', 'prelu', 'elu', etc.
     upsample: str = (
-        "deconv"  # Upsampling mode for MONAI BasicUNet: 'deconv' (transposed conv), 'nontrainable' (interpolation + conv), or 'pixelshuffle'
+        "deconv"  # Upsampling mode: 'deconv' (transposed conv),
+        # 'nontrainable' (interpolation + conv), or 'pixelshuffle'
     )
 
     # Transformer-specific (UNETR, etc.)
@@ -911,19 +911,16 @@ class TestTimeAugmentationConfig:
     - rotation90_axes: Uses spatial-only indices (e.g., [1, 2] for H-W plane where 0=D, 1=H, 2=W)
     """
 
-    enabled: bool = False
-    flip_axes: Any = (
-        None  # TTA flip strategy: "all" (8 flips), null (no aug), or list like [[2], [3]] (full tensor indices)
-    )
-    rotation90_axes: Any = (
-        None  # TTA rotation90 strategy: "all" (3 planes × 4 rotations), null, or list like [[1, 2]] (spatial indices: 0=D, 1=H, 2=W)
-    )
+    flip_axes: Any = None  # TTA flip strategy: "all" (8 flips), null (no aug),
+    # or list like [[2], [3]] (full tensor indices)
+    rotation90_axes: Any = None  # TTA rotation90 strategy: "all" (3 planes × 4 rotations),
+    # null, or list like [[1, 2]] (spatial indices: 0=D, 1=H, 2=W)
     channel_activations: Optional[List[Any]] = (
-        None  # Per-channel activations: [[start_ch, end_ch, 'activation'], ...] e.g., [[0, 2, 'softmax'], [2, 3, 'sigmoid'], [3, 4, 'tanh']]
-    )
-    select_channel: Any = (
-        None  # Channel selection: null (all), [1] (foreground), -1 (all) (applied even with null flip_axes)
+        None  # Per-channel activations: [[start_ch, end_ch, 'activation'], ...]
+        # e.g., [[0, 2, 'softmax'], [2, 3, 'sigmoid'], [3, 4, 'tanh']]
     )
+    select_channel: Any = None  # Channel selection: null (all), [1] (foreground), -1 (all)
+    # (applied even with null flip_axes)
     ensemble_mode: str = "mean"  # Ensemble mode for TTA: 'mean', 'min', 'max'
     apply_mask: bool = False  # Multiply each channel by corresponding test_mask after ensemble
 
@@ -943,8 +940,9 @@ class SavePredictionConfig:
 
     enabled: bool = True  # Enable saving intermediate predictions
     intensity_scale: float = (
-        -1.0
-    )  # If < 0, keep raw predictions (no normalization/scaling). If > 0, normalize to [0,1] then scale.
+        -1.0  # If < 0, keep raw predictions (no normalization/scaling).
+        # If > 0, normalize to [0,1] then scale.
+    )
     intensity_dtype: str = (
         "uint8"  # Save as uint8 for visualization (ignored if intensity_scale < 0)
     )
@@ -983,9 +981,8 @@ class DecodeBinaryContourDistanceWatershedConfig:
 class DecodeModeConfig:
     """Configuration for a single decode mode/function."""
 
-    name: str = (
-        "decode_binary_watershed"  # Function name: decode_binary_cc, decode_binary_watershed, decode_binary_contour_distance_watershed, etc.
-    )
+    name: str = "decode_binary_watershed"  # Function name: decode_binary_cc,
+    # decode_binary_watershed, decode_binary_contour_distance_watershed, etc.
     kwargs: Dict[str, Any] = field(
         default_factory=dict
     )  # Keyword arguments for the decode function
@@ -1105,6 +1102,8 @@ class TestDataConfig:
     cache_suffix: str = "_prediction.h5"
     # Image transformation (applied to test images during inference)
     image_transform: ImageTransformConfig = field(default_factory=ImageTransformConfig)
+    # Label transformation (optional, typically not used for test mode to preserve raw labels for evaluation)
+    label_transform: Optional[LabelTransformConfig] = None
 
 
 @dataclass
@@ -1128,6 +1127,8 @@ class TuneDataConfig:
     tune_resolution: Optional[List[int]] = None
     # Image transformation (applied to tune images during inference)
     image_transform: ImageTransformConfig = field(default_factory=ImageTransformConfig)
+    # Label transformation (optional, typically not used for tune mode to preserve raw labels for evaluation)
+    label_transform: Optional[LabelTransformConfig] = None
 
 
 @dataclass
@@ -1210,6 +1211,24 @@ class TuneConfig:
     parameter_space: ParameterSpaceConfig = field(default_factory=ParameterSpaceConfig)
 
 
+# Allow safe loading of checkpoints with PyTorch 2.6+ weights_only defaults
+try:
+    import torch
+
+    if hasattr(torch, "serialization") and hasattr(torch.serialization, "add_safe_globals"):
+        torch.serialization.add_safe_globals(
+            [
+                ParameterConfig,
+                DecodingParameterSpace,
+                PostprocessingParameterSpace,
+                ParameterSpaceConfig,
+            ]
+        )
+except Exception:
+    # Best-effort registration; ignore if torch not available at import time
+    pass
+
+
 @dataclass
 class Config:
     """Main configuration for PyTorch Connectomics.

connectomics/config/hydra_utils.py

@@ -269,7 +269,11 @@ def resolve_data_paths(cfg: Config) -> Config:
         >>> cfg.data.train_image = ["PT37/*_raw.tif", "file.tif"]
         >>> resolve_data_paths(cfg)
         >>> print(cfg.data.train_image)
-        ['/data/barcode/PT37/img1_raw.tif', '/data/barcode/PT37/img2_raw.tif', '/data/barcode/file.tif']
+        [
+            '/data/barcode/PT37/img1_raw.tif',
+            '/data/barcode/PT37/img2_raw.tif',
+            '/data/barcode/file.tif'
+        ]
 
         >>> cfg.test.data.test_path = "/data/test/"
         >>> cfg.test.data.test_image = ["volume_*.tif"]
@@ -325,7 +329,8 @@ def _combine_path(
                 index = int(selector)
                 if index < -len(expanded) or index >= len(expanded):
                     print(
-                        f"Warning: Index {index} out of range for {len(expanded)} files, using first"
+                        f"Warning: Index {index} out of range for {len(expanded)} files, "
+                        f"using first"
                     )
                     return expanded[0]
                 return expanded[index]
@@ -343,7 +348,8 @@ def _combine_path(
                     return matching[0]
                 else:
                     print(
-                        f"Warning: No file matches selector '{selector}', using first of {len(expanded)} files"
+                        f"Warning: No file matches selector '{selector}', "
+                        f"using first of {len(expanded)} files"
                     )
                     return expanded[0]
 

connectomics/config/slurm_utils.py

@@ -275,7 +275,6 @@ def filter_partitions(
                 and resources.gpus >= min_gpus
                 and resources.memory_gb >= min_memory_gb
             ):
-
                 # Check GPU type if specified
                 if gpu_type and resources.gpu_type != gpu_type:
                     continue

connectomics/data/augment/build.py

@@ -12,13 +12,10 @@
     RandRotate90d,
     RandFlipd,
     RandAffined,
-    RandZoomd,
     RandGaussianNoised,
     RandShiftIntensityd,
-    RandGaussianSmoothd,
     RandAdjustContrastd,
     RandSpatialCropd,
-    ScaleIntensityRanged,
     ToTensord,
     CenterSpatialCropd,
     SpatialPadd,
@@ -55,7 +52,8 @@ def build_train_transforms(
 
     Args:
         cfg: Hydra Config object
-        keys: Keys to transform (default: ['image', 'label'] or ['image', 'label', 'mask'] if masks are used)
+        keys: Keys to transform (default: ['image', 'label'] or
+            ['image', 'label', 'mask'] if masks are used)
         skip_loading: Skip LoadVolumed (for pre-cached datasets)
 
     Returns:
@@ -295,7 +293,8 @@ def _build_eval_transforms_impl(cfg: Config, mode: str = "val", keys: list[str]
         transforms.append(ApplyVolumetricSplitd(keys=keys))
 
     # Apply resize if configured (before cropping)
-    # For test/tune mode, only use test.data.image_transform or tune.data.image_transform (no fallback)
+    # For test/tune mode, only use test.data.image_transform or
+    # tune.data.image_transform (no fallback)
     # For val mode, use data.image_transform
     resize_factors = None
     if mode == "test":
@@ -370,7 +369,8 @@ def _build_eval_transforms_impl(cfg: Config, mode: str = "val", keys: list[str]
     # else: mode == "test" -> no cropping for sliding window inference
 
     # Normalization - use smart normalization
-    # For test/tune mode, only use test.data.image_transform or tune.data.image_transform (no fallback)
+    # For test/tune mode, only use test.data.image_transform or
+    # tune.data.image_transform (no fallback)
     # For val mode, use data.image_transform
     image_transform = None
     if (
@@ -407,42 +407,20 @@ def _build_eval_transforms_impl(cfg: Config, mode: str = "val", keys: list[str]
         if getattr(cfg.data, "normalize_labels", False):
             transforms.append(NormalizeLabelsd(keys=["label"]))
 
-        # Check if we should skip label transforms (test/tune mode)
-        # Skip label transforms if test.data or tune.data has evaluation.enabled=True
-        # This preserves original instance labels for metric computation
-        skip_label_transform = False
-        if mode == "test":
-            # Check if test.evaluation or tune.evaluation is enabled (for adapted_rand, etc.)
-            evaluation_config = None
-            if hasattr(cfg, "test") and hasattr(cfg.test, "evaluation"):
-                evaluation_config = cfg.test.evaluation
-            elif hasattr(cfg, "tune") and cfg.tune and hasattr(cfg.tune, "optimization"):
-                # For tune mode, check if we're optimizing metrics that need instance labels
-                if hasattr(cfg.tune.optimization, "single_objective"):
-                    metric = getattr(cfg.tune.optimization.single_objective, "metric", None)
-                    if metric == "adapted_rand":
-                        skip_label_transform = True
-                        print(
-                            f"  ⚠️  Skipping label transforms for Optuna tuning (keeping original labels for {metric})"
-                        )
-
-            if evaluation_config:
-                evaluation_enabled = getattr(evaluation_config, "enabled", False)
-                metrics = getattr(evaluation_config, "metrics", [])
-                if evaluation_enabled and metrics:
-                    skip_label_transform = True
-                    print(
-                        f"  ⚠️  Skipping label transforms for metric evaluation (keeping original labels for {metrics})"
-                    )
-
         # Label transformations (affinity, distance transform, etc.)
-        # Only apply if not skipped AND label_transform is configured
-        if hasattr(cfg.data, "label_transform") and not skip_label_transform:
+        # For test/tune modes: NEVER apply label transforms (keep raw instance labels for evaluation)
+        # For val mode: use training label_transform config
+        label_cfg = None
+        if mode == "val":
+            # Validation always uses training label_transform
+            if hasattr(cfg.data, "label_transform"):
+                label_cfg = cfg.data.label_transform
+
+        # Apply label transforms if configured
+        if label_cfg is not None:
             from ..process.build import create_label_transform_pipeline
             from ..process.monai_transforms import SegErosionInstanced
 
-            label_cfg = cfg.data.label_transform
-
             # 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))
@@ -642,27 +620,32 @@ def should_augment(aug_name: str, aug_enabled: Optional[bool]) -> bool:
             scale_range = aug_cfg.affine.scale_range
             shear_range = aug_cfg.affine.shear_range
 
+        # Interpolation per key: bilinear for images, nearest for labels/masks
+        affine_modes = ["bilinear" if k == "image" else "nearest" for k in keys]
+
         transforms.append(
             RandAffined(
                 keys=keys,
                 prob=aug_cfg.affine.prob,
                 rotate_range=rotate_range,
                 scale_range=scale_range,
                 shear_range=shear_range,
-                mode="bilinear",
+                mode=affine_modes,
                 padding_mode="reflection",
             )
         )
 
     if should_augment("elastic", aug_cfg.elastic.enabled):
         # Unified elastic deformation that supports both 2D and 3D
+        elastic_modes = ["bilinear" if k == "image" else "nearest" for k in keys]
         transforms.append(
             RandElasticd(
                 keys=keys,
                 do_2d=do_2d,
                 prob=aug_cfg.elastic.prob,
                 sigma_range=aug_cfg.elastic.sigma_range,
                 magnitude_range=aug_cfg.elastic.magnitude_range,
+                mode=elastic_modes,
             )
         )
 

connectomics/data/augment/monai_transforms.py

@@ -6,11 +6,13 @@
 """
 
 from __future__ import annotations
-from typing import Dict, Any, List, Optional, Union, Tuple
+
 import math
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import cv2
 import numpy as np
 import torch
-import cv2
 from monai.config import KeysCollection
 from monai.transforms import MapTransform, RandomizableTransform
 
@@ -1068,7 +1070,8 @@ def __init__(
                 self.mode = "divide"
             except ValueError:
                 raise ValueError(
-                    f"Invalid divide mode '{mode}'. Format should be 'divide-K' where K is a number (e.g., 'divide-255')"
+                    f"Invalid divide mode '{mode}'. Format should be 'divide-K' "
+                    f"where K is a number (e.g., 'divide-255')"
                 )
         elif mode not in ["none", "normal", "0-1"]:
             raise ValueError(
@@ -1212,7 +1215,8 @@ def __init__(
 
         if orientation not in ["horizontal", "vertical", "random"]:
             raise ValueError(
-                f"Invalid orientation '{orientation}'. Must be 'horizontal', 'vertical', or 'random'"
+                f"Invalid orientation '{orientation}'. Must be 'horizontal', "
+                f"'vertical', or 'random'"
             )
         self.orientation = orientation
 
@@ -1356,7 +1360,8 @@ class ResizeByFactord(MapTransform):
 
     Args:
         keys: Keys to transform
-        scale_factors: Scale factors for each spatial dimension (e.g., [0.25, 0.25] for 2D, [0.5, 0.5, 0.5] for 3D)
+        scale_factors: Scale factors for each spatial dimension
+            (e.g., [0.25, 0.25] for 2D, [0.5, 0.5, 0.5] for 3D)
         mode: Interpolation mode ('bilinear', 'nearest', 'area', etc.)
         align_corners: Whether to align corners (True for bilinear, None for nearest)
         allow_missing_keys: Whether to allow missing keys