update demo

Author: Donglai Wei

README.md

@@ -31,9 +31,7 @@ PyTorch Connectomics (PyTC) helps neuroscientists:
 - ✅ **Train models** without deep ML expertise
 - ✅ **Process** large-scale connectomics datasets efficiently
 
-**Built on:** [PyTorch Lightning](https://lightning.ai/) + [MONAI](https://monai.io/) for modern, scalable deep learning.
-
-**Used by:** Harvard, MIT, Janelia Research Campus, and 100+ labs worldwide.
+**Built on:** [PyTorch Lightning](https://lightning.ai/) + [MONAI](https://monai.io/) + [nnU-Net](https://github.com/MIC-DKFZ/nnUNet) for modern, scalable deep learning.
 
 ---
 
@@ -105,22 +103,29 @@ python scripts/main.py --demo
 Train on real mitochondria segmentation data:
 
 ```bash
-# Download tutorial data (~100 MB)
-mkdir -p datasets/
-wget https://huggingface.co/datasets/pytc/tutorial/resolve/main/Lucchi%2B%2B.zip
-unzip Lucchi++.zip -d datasets/
-rm Lucchi++.zip
+# Download tutorial data (~50 MB)
+just download lucchi++
 
 # Quick test (1 batch)
-python scripts/main.py --config tutorials/monai_lucchi++.yaml --fast-dev-run
+just train lucchi++ monai_unet --fast-dev-run
 
-# Full training
-python scripts/main.py --config tutorials/monai_lucchi++.yaml
+# Full training on a single GPU (choose your architecture: monai_unet, rsunet, mednext)
+just train lucchi++ monai_unet -- system.training.num_gpus=1
 ```
 
 **Monitor progress:**
 ```bash
-tensorboard --logdir outputs/lucchi++_monai_unet
+just tensorboard lucchi++_monai_unet
+```
+
+**Resume training from checkpoint:**
+```bash
+just resume lucchi++ monai_unet outputs/lucchi++_monai_unet/*/checkpoints/last.ckpt
+```
+
+**Run inference:**
+```bash
+just test lucchi++ monai_unet outputs/lucchi++_monai_unet/*/checkpoints/best.ckpt
 ```
 
 ---

connectomics/models/arch/monai_models.py

@@ -102,9 +102,10 @@ def build_basic_unet(cfg) -> ConnectomicsModel:
         )
 
     # BasicUNet requires exactly 6 feature levels
-    base_features = list(cfg.model.filters) if hasattr(cfg.model, 'filters') else [32, 64, 128, 256, 512]
+    # Pad with last value repeated (not doubled) to keep memory usage low
+    base_features = list(cfg.model.filters) if hasattr(cfg.model, 'filters') else [32, 64, 128, 256, 512, 1024]
     while len(base_features) < 6:
-        base_features.append(base_features[-1] * 2)
+        base_features.append(base_features[-1])  # Repeat last value instead of doubling
     features = tuple(base_features[:6])
 
     model = BasicUNet(
@@ -157,8 +158,7 @@ def build_monai_unet(cfg) -> ConnectomicsModel:
             "Use [H, W] for 2D or [D, H, W] for 3D. "
             "spatial_dims will be automatically inferred from input_size length."
         )
-    features = list(cfg.model.filters) if hasattr(cfg.model, 'filters') else [32, 64, 128, 256, 512]
-    channels = features[:5]  # Limit to 5 levels
+    channels = list(cfg.model.filters) if hasattr(cfg.model, 'filters') else [32, 64, 128, 256, 512]
     strides = [2] * (len(channels) - 1)  # 2x downsampling at each level
 
     # Handle normalization type and parameters

connectomics/utils/demo.py

@@ -85,6 +85,8 @@ def create_demo_config():
         ModelConfig,
         DataConfig,
         OptimizationConfig,
+        OptimizerConfig,
+        SchedulerConfig,
         MonitorConfig,
         CheckpointConfig,
         EarlyStoppingConfig,
@@ -111,14 +113,15 @@ def create_demo_config():
             ),
         ),
         model=ModelConfig(
-            architecture="monai_basic_unet3d",
+            architecture="monai_unet",
             in_channels=1,
-            out_channels=2,
+            out_channels=1,  # Single channel binary segmentation
             spatial_dims=3,
             filters=[16, 32, 64, 128],  # Smaller for demo
             dropout=0.1,
             loss_functions=["DiceLoss"],
             loss_weights=[1.0],
+            loss_kwargs=[{"sigmoid": True, "smooth_nr": 1e-5, "smooth_dr": 1e-5}],
         ),
         data=DataConfig(
             train_image=None,  # Will be generated
@@ -142,9 +145,9 @@ def create_demo_config():
             log_every_n_steps=1,
             deterministic=False,
             benchmark=True,
+            optimizer=OptimizerConfig(name="AdamW", lr=1e-3, weight_decay=1e-4),
+            scheduler=SchedulerConfig(name="ConstantLR", warmup_epochs=0),
         ),
-        optimizer={"name": "AdamW", "lr": 1e-3, "weight_decay": 1e-4},
-        scheduler={"name": "ConstantLR", "warmup_epochs": 0},
         monitor=MonitorConfig(
             checkpoint=CheckpointConfig(
                 dirpath="outputs/demo/checkpoints",
@@ -323,13 +326,18 @@ def run_demo():
     except (ImportError, ModuleNotFoundError):
         pass
 
+    # Use TensorBoard logger to avoid "no logger" warnings
+    from pytorch_lightning.loggers import TensorBoardLogger
+    demo_logger = TensorBoardLogger(save_dir=str(temp_dir), name="demo_logs", version="")
+
     trainer = pl.Trainer(
         max_epochs=cfg.optimization.max_epochs,
         accelerator="gpu" if cfg.system.training.num_gpus > 0 else "cpu",
         devices=cfg.system.training.num_gpus if cfg.system.training.num_gpus > 0 else 1,
         precision=cfg.optimization.precision,
         callbacks=callbacks,
-        logger=False,  # Disable logging for demo
+        logger=demo_logger,
+        log_every_n_steps=cfg.optimization.log_every_n_steps,
         enable_checkpointing=True,
         enable_progress_bar=True,
         enable_model_summary=True,
@@ -352,13 +360,18 @@ def run_demo():
         print("=" * 60)
         print("\nYour installation is working correctly! 🎉")
         print("\n📚 Next steps:")
-        print("  1. Try a tutorial:")
-        print("     python scripts/main.py --config tutorials/lucchi.yaml --fast-dev-run")
-        print("\n  2. Download tutorial data:")
-        print("     just download-data lucchi  # Or see README for manual download")
+        print("  1. Download tutorial data:")
+        print("     just download lucchi++")
+        print("     just download-list  # See all available datasets")
+        print("\n  2. Try a fast dev run:")
+        print("     just train lucchi++ monai_unet -- --fast-dev-run")
         print("\n  3. Train on real data:")
-        print("     python scripts/main.py --config tutorials/lucchi.yaml")
-        print("\n  4. Read the documentation:")
+        print("     just train lucchi++ monai_unet")
+        print("     just train lucchi++ rsunet")
+        print("     just train lucchi++ mednext")
+        print("\n  4. Monitor training:")
+        print("     just tensorboard lucchi++_monai_unet")
+        print("\n  5. Read the documentation:")
         print("     https://connectomics.readthedocs.io")
         print("\n" + "=" * 60 + "\n")
 

justfile

@@ -6,28 +6,38 @@ default:
     @just --list
 
 # ============================================================================
-# SLURM/HPC Setup
+# Setup & Data
 # ============================================================================
 
 # Setup SLURM environment: detect CUDA/cuDNN and install PyTorch with correct versions
 setup-slurm:
     bash connectomics/utils/setup_slurm.sh
 
+# Download dataset(s) (e.g., just download lucchi++, just download all)
+# Available: lucchi++, snemi, mitoem, cremi
+download +datasets:
+    python scripts/download_data.py {{datasets}}
+
+# List available datasets
+download-list:
+    python scripts/download_data.py --list
+
 # ============================================================================
 # Training Commands
 # ============================================================================
 
-# Train on Lucchi dataset (use '+' to pass extra args: just train monai lucchi++ -- data.batch_size=8, --fast-dev-run)
-train model dataset *ARGS='':
-    python scripts/main.py --config tutorials/{{model}}_{{dataset}}.yaml {{ARGS}}
+# Train with unified config (e.g., just train lucchi++ rsunet -- data.batch_size=8)
+# Architecture options: monai_basic_unet3d, rsunet, mednext
+train dataset arch *ARGS='':
+    python scripts/main.py --config tutorials/{{dataset}}.yaml model.architecture={{arch}} {{ARGS}}
 
-# Continue training from checkpoint (use '+' for extra args: just resume monai lucchi++ ckpt.pt -- --reset-optimizer)
-resume model dataset checkpoint *ARGS='':
-    python scripts/main.py --config tutorials/{{model}}_{{dataset}}.yaml --checkpoint {{checkpoint}} {{ARGS}}
+# Continue training from checkpoint (e.g., just resume lucchi++ rsunet ckpt.pt -- --reset-optimizer)
+resume dataset arch checkpoint *ARGS='':
+    python scripts/main.py --config tutorials/{{dataset}}.yaml model.architecture={{arch}} --checkpoint {{checkpoint}} {{ARGS}}
 
-# Test on Lucchi++ dataset (provide path to checkpoint)
-test model dataset checkpoint *ARGS='':
-    python scripts/main.py --config tutorials/{{model}}_{{dataset}}.yaml --mode test --checkpoint {{checkpoint}} {{ARGS}}
+# Test model (e.g., just test lucchi++ rsunet ckpt.pt)
+test dataset arch checkpoint *ARGS='':
+    python scripts/main.py --config tutorials/{{dataset}}.yaml model.architecture={{arch}} --mode test --checkpoint {{checkpoint}} {{ARGS}}
 
 # ============================================================================
 # Monitoring Commands