Lung Nodule Segmentation — 3D U-Net (DynUNet)

Voxel-level segmentation of pulmonary nodules in 3D chest CT volumes, cropped to the lung region. MONAI DynUNet — the nnU-Net-style dynamic 3D U-Net (encoder-decoder with symmetric skip connections, instance normalisation, per-level kernel/stride configuration).

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

Model details

Architecture: MONAI DynUNet (nnU-Net-style 3D U-Net)
Trainable parameters: 31,181,730
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
Depth: 6 encoder levels
Kernels: [3, 3, 3] at every level
Strides: [[1,1,1], [2,2,2], [2,2,2], [2,2,2], [2,2,2], [2,2,2]]
Framework: PyTorch + MONAI

Data

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

NLST
NSCLC-Radiomics
LIDC-IDRI

Split sizes: 1 683 train / 297 val / 325 test (held out). Nodule voxels are on the order of ~10⁻⁵ of the total — see the note on metric interpretation below.

Validation metrics

Evaluated on 297 val volumes (≈ 5 × 10⁹ voxels), micro-averaged at the argmax of the 2-class softmax output.

Metric	Value
mIoU	0.7431
Accuracy	0.9995
Precision	0.6479
Recall	0.6618
Dice / F1 (micro)	0.6548
IoU (nodule class)	0.4867
Dice per case, mean	0.5380
Dice per case, med.	0.6105

Because ~10⁻⁵ of voxels are nodule, Accuracy is trivially near 1.0 regardless of model quality and mIoU is dominated by IoU_background. The operationally meaningful numbers are Precision, Recall, and the per-case Dice distribution.

Compared to the SegResNet 3D-small baseline (nodule-segresnet-3d-small), DynUNet trades some recall (0.66 vs 0.74) for higher precision (0.65 vs 0.61) — depending on the downstream use case, one or the other may be preferable.

How to load & run inference

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

cfg = yaml.safe_load(open("config.yaml"))["model"]
default_strides = [[1, 1, 1]] + [[2, 2, 2]] * 5
strides = cfg.get("strides", default_strides)
kernel  = cfg.get("kernel_size", [[3, 3, 3]] * len(strides))
model = DynUNet(
    spatial_dims         = cfg["spatial_dims"],
    in_channels          = cfg["in_channels"],
    out_channels         = cfg["out_channels"],
    kernel_size          = kernel,
    strides              = strides,
    upsample_kernel_size = cfg.get("upsample_kernel_size", strides[1:]),
    norm_name            = cfg.get("norm_name", "instance"),
    deep_supervision     = cfg.get("deep_supervision", False),
    deep_supr_num        = cfg.get("deep_supr_num", 1),
    dropout              = cfg.get("dropout_prob", 0.0),
)
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 = 400, η_min = 1e-6)
Batch size: 2
Epochs: 400 | best checkpoint at epoch 354 / 400
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: ≈ 3.5 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.

Downloads last month: 6