Kakimaki00/nodule-segresnet-3d-wide

Lung Nodule Segmentation — SegResNet (3D, wide) — paper's best model

Voxel-level segmentation of pulmonary nodules in 3D chest CT volumes, cropped to the lung region. Wide-capacity SegResNet (init_filters = 32, ~83 M params), best-performing model reported in the paper.

Also known as v7 in the paper's model table.

Model details

• Architecture: MONAI SegResNet, 3D residual U-Net
• Trainable parameters: 82,637,282
• Input: (1, 256, 256, 256) CT crop, intensity-normalised to [0, 1], resampled from a per-series lung bbox (padding = 20 vox)
• Output: (2, 256, 256, 256) softmax logits — class 0 = background, class 1 = nodule
• Framework: PyTorch + MONAI

Data

Trained on the unified split (patient-grouped, dataset-stratified):

• NLST
• NSCLC-Radiomics
• LIDC-IDRI

Split sizes: 1 609 train / 345 val / 351 test (held out).

Note on split. This model uses the earlier unified.json split, not the unified_v2 split used by the other nodule models (nodule-segresnet-3d-small, nodule-dynunet-3d). It reached its peak val Dice of 0.589 at epoch 545 on this split, well past the 400-epoch budget the retrained models use.

Nodule voxels are on the order of ~10⁻⁵ of the total.

Validation metrics

Best val Dice on the unified split val set:

Metric	Value
Dice (per-case mean)	0.589
Recall (sensitivity)	0.774
Precision	0.693

How to load & run inference

import yaml, torch
from monai.networks.nets import SegResNet

cfg = yaml.safe_load(open("config.yaml"))["model"]
model = SegResNet(
    spatial_dims = cfg["spatial_dims"],
    in_channels  = cfg["in_channels"],
    out_channels = cfg["out_channels"],
    init_filters = cfg["init_filters"],
    blocks_down  = tuple(cfg["blocks_down"]),
    blocks_up    = tuple(cfg["blocks_up"]),
    dropout_prob = cfg["dropout_prob"],
)
state = torch.load("model.pth", map_location="cpu", weights_only=True)
model.load_state_dict(state)
model.eval()

with torch.no_grad():
    # Input: (B, 1, 256, 256, 256), pre-cropped to the lung bbox and normalised to [0, 1]
    x = torch.randn(1, 1, 256, 256, 256)
    logits = model(x)                          # (B, 2, D, H, W)
    pred_class = logits.argmax(dim=1)          # (B, D, H, W)
    nodule_mask = (pred_class == 1).to(torch.uint8)

The model expects lung-bbox-cropped input. A separate 2D ROI model is needed to produce that bbox — see the accompanying ROI checkpoints (roi-segresnet-2d or roi-swinunetr-2d).

Training recipe

• Loss: Focal Tversky + weighted CE (α=0.3, β=0.7, γ=2.0, λ_CE=0.1, nodule class weight = 100)
• Optimizer: Adam (lr = 1e-5, weight decay = 1e-5)
• Scheduler: CosineAnnealingLR (T_max = 1000, η_min = 1e-6)
• Batch size: 2 (larger model, less headroom)
• Epochs: 1000 | best checkpoint around epoch 545 / 1000
• Mixed precision: bf16
• Augmentation: 3D flips, 90° rotations, elastic rotation, zoom, intensity scale/shift, Gaussian noise/blur, contrast
• Seed: 42
• Hardware: 1 × NVIDIA H100 94 GB
• Wall-clock: ≈ 12 days

Full config is included in this repo as config.yaml.

Reproducing training

Training code lives in an accompanying reproduction demo (published separately). Once available, reproduce with:

export DATA_ROOT=/path/to/unified          # dir containing ct_3d/ and nodule_sem_seg_3d/
python train.py --config config.yaml

License & intended use

Model weights released under Apache 2.0. Training data was public but covered by dataset-specific terms (NLST, NSCLC-Radiomics, LIDC-IDRI) — users must comply with those separately when using the model on comparable data.

Not a medical device. Not intended for clinical use. Research only.

Citation

Paper in preparation.