nodule-segresnet-3d-wide

Lung Nodule Segmentation — SegResNet (3D, WIDE)

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) retrained on the unified split.

This is the "wide SegResNet" architecture from the paper, directly comparable to the other two-stage nodule models (segresnet_small, dynunet) — all of which share the same train / val / test split.

Peaked at val Dice 0.5711 (per-case mean) at epoch 489 / 1000.

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):

Split sizes: 1 683 train / 297 val / 325 test (held out).

Note on bboxes. This model was originally trained with a regenerated bbox JSON (tighter lung crops) that differs from the bboxes_unified.json bundled with the small and DynUNet variants. Reproducing from the bundled bboxes gives small variance from the released checkpoint but does not materially affect training.

Validation & test metrics

Test split (325 series, ~5.5 × 10⁹ voxels, micro-averaged):

Per-case Dice distribution: mean 0.5706, median 0.6304, p25 0.4069, p75 0.7834.

Comparison to the other two-stage nodule models on the same test set

On voxel-level metrics (mIoU / Dice_micro), DynUNet wins; on per-case Dice mean (equal weight per patient regardless of nodule size), this wide-SegResNet variant wins. Which is "best" depends on the metric: large-nodule cases dominate the micro-averaged numbers, while per-case-averaged metrics weigh each patient equally.

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"],      # 2
    init_filters = cfg["init_filters"],      # 32 (wide)
    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), lung-bbox-cropped CT resampled to 256³
    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) in {0, 1}
    nodule_mask = (pred_class == 1).to(torch.uint8)

This model expects lung-bbox-cropped input. A separate 2D ROI model is needed to produce that bbox — see the accompanying ROI checkpoints (szabopeter/roi-segresnet-2d or szabopeter/roi-swinunetr-2d), or use the demo's inference.py which chains them for you.

Training recipe

• Loss: Focal Tversky + weighted CE (α=0.3, β=0.7, γ=2.0, λ_CE=0.3, ce_nodule_weight = 100)
• Optimizer: Adam (lr=1e-5, wd=1e-5)
• Scheduler: CosineAnnealingLR (T_max=1000, η_min=1e-6)
• Batch size: 2
• Epochs: 1000 | best checkpoint at epoch 489 / 1000
• Mixed precision: bf16
• Seed: 42
• Hardware: 1 × NVIDIA H100 94 GB
• Wall-clock: ≈ 12 days

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

Companion models

• szabopeter/nodule-segresnet-3d-small — small SegResNet (init_filters=16, ~20 M params)
• szabopeter/nodule-dynunet-3d — DynUNet / 3D U-Net (~31 M params, best voxel-level metrics)

License & intended use

Apache 2.0. Training data was public but subject to dataset-specific terms (NLST, NSCLC-Radiomics, LIDC-IDRI).

Not a medical device. Research use only.

Citation

Paper in preparation.