refactor training/inference

Author: Donglai Wei

.gitignore

@@ -149,6 +149,10 @@ build/
 .DS_Store
 Thumbs.db
 
+# external libraries
+lib/
+
 # Development logs and documentation
 .claude/
+.codex/
 .CLAUDE.md

CLAUDE.md

@@ -156,11 +156,16 @@ connectomics/                    # Main Python package (77 files, ~23K lines)
 │       ├── build.py             # Optimizer/scheduler factory
 │       └── lr_scheduler.py      # Custom LR schedulers
 │
-├── lightning/                   # PyTorch Lightning integration (PRIMARY)
-│   ├── lit_data.py              # LightningDataModule (Volume/Tile/Cloud datasets)
-│   ├── lit_model.py             # LightningModule (1.8K lines - deep supervision, TTA)
-│   ├── lit_trainer.py           # Trainer creation utilities
-│   └── callbacks.py             # Custom Lightning callbacks
+├── training/                    # Training utilities
+│   ├── lit/                     # PyTorch Lightning integration (PRIMARY)
+│   │   ├── data.py              # LightningDataModule (Volume/Tile/Cloud datasets)
+│   │   ├── model.py             # LightningModule (deep supervision, TTA)
+│   │   ├── trainer.py           # Trainer creation utilities
+│   │   ├── callbacks.py         # Custom Lightning callbacks
+│   │   ├── config.py            # Factory functions for training setup
+│   │   └── utils.py             # CLI/config helpers
+│   ├── deep_supervision.py      # Deep supervision utilities
+│   └── debugging.py             # NaN detection and debugging utilities
 │
 ├── data/                        # Data loading and preprocessing
 │   ├── dataset/                 # Dataset classes (HDF5, TIFF, Zarr, Cloud)
@@ -399,7 +404,7 @@ Multiple losses can be combined with weights in the config.
 
 ## PyTorch Lightning Integration
 
-### LightningModule (`lightning/lit_model.py`)
+### LightningModule (`training/lit/model.py`)
 Wraps models with automatic training features:
 - Distributed training (DDP)
 - Mixed precision training (AMP)
@@ -410,7 +415,7 @@ Wraps models with automatic training features:
 - **Deep supervision**: Multi-scale loss computation with automatic target resizing
 
 ```python
-from connectomics.lightning import ConnectomicsModule
+from connectomics.training.lit import ConnectomicsModule
 
 # Create Lightning module
 lit_model = ConnectomicsModule(cfg)
@@ -419,23 +424,23 @@ lit_model = ConnectomicsModule(cfg)
 lit_model = ConnectomicsModule(cfg, model=custom_model)
 ```
 
-### LightningDataModule (`lightning/lit_data.py`)
+### LightningDataModule (`training/lit/data.py`)
 Handles data loading with MONAI transforms:
 - Train/val/test splits
 - MONAI CacheDataset for fast loading
 - Automatic augmentation pipeline
 - Persistent workers for efficiency
 
 ```python
-from connectomics.lightning import ConnectomicsDataModule
+from connectomics.training.lit import ConnectomicsDataModule
 
 datamodule = ConnectomicsDataModule(cfg)
 ```
 
-### Trainer (`lightning/lit_trainer.py`)
+### Trainer (`training/lit/trainer.py`)
 Convenience function for creating Lightning Trainer:
 ```python
-from connectomics.lightning import create_trainer
+from connectomics.training.lit import create_trainer
 
 trainer = create_trainer(cfg)
 trainer.fit(lit_model, datamodule=datamodule)
@@ -475,15 +480,15 @@ from pytorch_lightning import seed_everything
 seed_everything(cfg.system.seed)
 
 # 3. Create data module
-from connectomics.lightning import ConnectomicsDataModule
+from connectomics.training.lit import ConnectomicsDataModule
 datamodule = ConnectomicsDataModule(cfg)
 
 # 4. Create model
-from connectomics.lightning import ConnectomicsModule
+from connectomics.training.lit import ConnectomicsModule
 model = ConnectomicsModule(cfg)
 
 # 5. Create trainer
-from connectomics.lightning import create_trainer
+from connectomics.training.lit import create_trainer
 trainer = create_trainer(cfg)
 
 # 6. Train
@@ -555,9 +560,9 @@ scheduler:
 - `connectomics/models/solver/build.py`: Optimizer/scheduler factory
 
 ### Lightning
-- `connectomics/lightning/lit_model.py`: Lightning module wrapper
-- `connectomics/lightning/lit_data.py`: Data module
-- `connectomics/lightning/lit_trainer.py`: Trainer utilities
+- `connectomics/training/lit/model.py`: Lightning module wrapper
+- `connectomics/training/lit/data.py`: Data module
+- `connectomics/training/lit/trainer.py`: Trainer utilities
 
 ### Entry Points
 - `scripts/main.py`: Primary training script (Lightning + Hydra)
@@ -635,7 +640,7 @@ The codebase has **fully migrated** from legacy systems:
 
 **Current development stack:**
 - Hydra/OmegaConf configs (`tutorials/*.yaml`)
-- PyTorch Lightning modules (`connectomics/lightning/`)
+- PyTorch Lightning modules (`connectomics/training/lit/`)
 - `scripts/main.py` entry point
 - MONAI models and transforms
 - Type-safe dataclass configurations
@@ -687,7 +692,7 @@ Install via `pip install -e .[extra_name]` where `extra_name` is:
 
 #### `[wandb]` - Weights & Biases Integration
 - **wandb** (>=0.13.0): Experiment tracking and monitoring
-- Used in: `connectomics.lightning.lit_trainer` (optional logger)
+- Used in: `connectomics.training.lit.trainer` (optional logger)
 
 #### `[tiff]` - TIFF File Support
 - **tifffile** (>=2021.11.2): Advanced TIFF reading/writing
@@ -817,4 +822,4 @@ python -c "from connectomics.models.arch import list_architectures; print(list_a
 - [MONAI Docs](https://docs.monai.io/en/stable/) - Medical imaging toolkit
 - [Hydra Docs](https://hydra.cc/) - Configuration management
 - [Project Documentation](https://zudi-lin.github.io/pytorch_connectomics/build/html/index.html) - Full docs
-- [Slack Community](https://join.slack.com/t/pytorchconnectomics/shared_invite/zt-obufj5d1-v5_NndNS5yog8vhxy4L12w) - Get help
+- [Slack Community](https://join.slack.com/t/pytorchconnectomics/shared_invite/zt-obufj5d1-v5_NndNS5yog8vhxy4L12w) - Get help

conda-recipe/meta.yaml

@@ -55,7 +55,7 @@ test:
     - connectomics
     - connectomics.config
     - connectomics.data
-    - connectomics.lightning
+    - connectomics.training.lit
     - connectomics.models
     - connectomics.utils
   commands:

connectomics/data/augment/build.py

@@ -376,19 +376,34 @@ def _build_eval_transforms_impl(
         if getattr(cfg.data, "normalize_labels", False):
             transforms.append(NormalizeLabelsd(keys=["label"]))
 
-        # Check if we should skip label transforms (test mode with evaluation metrics)
+        # 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":
-            if hasattr(cfg, "inference") and hasattr(cfg.inference, "evaluation"):
-                evaluation_enabled = getattr(cfg.inference.evaluation, "enabled", False)
-                metrics = getattr(cfg.inference.evaluation, "metrics", [])
+            # 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:
             from ..process.build import create_label_transform_pipeline
             from ..process.monai_transforms import SegErosionInstanced

connectomics/decoding/optuna_tuner.py

@@ -620,7 +620,7 @@ def run_tuning(model, trainer, cfg, checkpoint_path=None):
         6. Save best parameters to YAML file
 
     Example:
-        >>> from connectomics.lightning import ConnectomicsModule, create_trainer
+        >>> from connectomics.training.lit import ConnectomicsModule, create_trainer
         >>> from connectomics.decoding import run_tuning
         >>> model = ConnectomicsModule(cfg)
         >>> trainer = create_trainer(cfg)
@@ -655,7 +655,7 @@ def run_tuning(model, trainer, cfg, checkpoint_path=None):
     print(f"Output directory: {output_dir}")
 
     # Step 1: Run inference on tune dataset
-    from connectomics.lightning import create_datamodule
+    from connectomics.training.lit import create_datamodule
     from connectomics.data.io import read_volume
     import glob
 

connectomics/decoding/segmentation.py

@@ -202,9 +202,6 @@ def decode_instance_binary_contour_distance(
             seed = remove_small_objects(seed, min_seed_size)
 
         # step 3: compute the segmentation mask
-        import pdb
-
-        pdb.set_trace()
         distance[distance < 0] = 0
         segmentation = mahotas.cwatershed(-distance.astype(np.float64), seed)
         segmentation[~foreground] = (

connectomics/inference/__init__.py

@@ -0,0 +1,18 @@
+"""Inference utilities package."""
+
+from .manager import InferenceManager
+from .io import apply_postprocessing, apply_decode_mode, resolve_output_filenames, write_outputs
+from .sliding import build_sliding_inferer, resolve_inferer_roi_size, resolve_inferer_overlap
+from .tta import TTAPredictor
+
+__all__ = [
+    "InferenceManager",
+    "apply_postprocessing",
+    "apply_decode_mode",
+    "resolve_output_filenames",
+    "write_outputs",
+    "build_sliding_inferer",
+    "resolve_inferer_roi_size",
+    "resolve_inferer_overlap",
+    "TTAPredictor",
+]