Furen Xiao
e7721760f2
- Updated `05_monai_run_iterative.py` to include legacy mode and improved configuration options for training. - Introduced `UNetBAligned` class in `monai_common.py` for deep supervision and instance normalization. - Modified `train_transform` and `make_dataloader` in `monai_common.py` to support B-aligned data protocol with Z-score normalization and augmentation. - Added new loss functions `nn_dc_and_ce` and `b_dc_and_ce` in `losses.py` for improved training performance. - Enhanced `Unet3D` architecture in `unet3d.py` to support leaky ReLU and deep supervision. - Created new augmentation functions in `augment_b.py` for Pipeline B's training, including rotation, scaling, and noise. - Implemented new training scripts `run_b_100ep_stages.sh` and `run_c_bloss_100ep_after_b.sh` for structured training and evaluation. - Updated dataset handling in `dataset.py` to accommodate new augmentation and normalization strategies.
359 lines
No EOL
14 KiB
Python
359 lines
No EOL
14 KiB
Python
"""Shared helpers for the MONAI iterative pseudo-labeling pipeline (scripts_monai/).
|
||
|
||
Backbone: MONAI 1.6.x (pip) — monai.networks.nets.UNet, MONAI transforms /
|
||
Dataset / DataLoader, and MONAI sliding-window inference. Everything
|
||
study-specific (splits, gates, loss weighting, longitudinal consistency) is
|
||
shared with the other pipelines so round results stay directly comparable:
|
||
|
||
* per-sample weighted Dice+CE loss — Pipeline A's convention (row "w";
|
||
pseudo-label rows down-weighted sample-by-sample)
|
||
* selection gates + head-relative consistency filter — Pipeline B's
|
||
implementation (scripts_nnu/nnu_common; absolute patient-space resampling
|
||
is unreliable across acquisitions, see README)
|
||
* DDP via torchrun for training, per-rank row sharding for inference —
|
||
Pipeline A's parallelism model
|
||
|
||
Network: monai UNet 3D, 1 in / 2 out, channels 16→128 (~1.2M params) — the
|
||
MONAI analogue of Pipeline A's Unet3D(base=16, depth=4).
|
||
"""
|
||
import os
|
||
import sys
|
||
|
||
# Data/run dirs: overridable for scratch runs. Code (src/, scripts_nnu/): next to this file.
|
||
ROOT = os.environ.get(
|
||
"LONGITUDINAL_ROOT", os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
|
||
_REPO = ROOT
|
||
sys.path.insert(0, _REPO)
|
||
sys.path.insert(0, os.path.join(_REPO, "scripts_nnu"))
|
||
|
||
import numpy as np
|
||
import torch
|
||
import torch.nn as nn
|
||
|
||
from monai.networks.nets import UNet # module is monai-specific; safe to import eagerly
|
||
from src.common import dice # noqa: E402
|
||
|
||
# Selection gates / consistency filter shared with Pipeline B (pure numpy/scipy)
|
||
from nnu_common import pos_mask, neg_frac_bg, consistency_filter, load_voxel_stats # noqa: E402,F401
|
||
|
||
WIN = 96
|
||
|
||
|
||
# ---------------- network / loss ----------------
|
||
|
||
def build_model(device, norm="batch", b_aligned=False):
|
||
from monai.networks.nets import UNet
|
||
if b_aligned:
|
||
return UNetBAligned(device, norm=norm).to(device)
|
||
net = UNet(
|
||
spatial_dims=3,
|
||
in_channels=1,
|
||
out_channels=2,
|
||
channels=(16, 32, 64, 128),
|
||
strides=(2, 2, 2),
|
||
num_res_units=2,
|
||
act=("LEAKYRELU", {"inplace": True, "negative_slope": 0.01}),
|
||
norm=("instance", {"affine": True, "eps": 1e-5}) if norm == "instance" else "batch",
|
||
)
|
||
return net.to(device)
|
||
|
||
|
||
class UNetBAligned(UNet):
|
||
"""Pipeline C's MONAI UNet + Pipeline B's deep supervision.
|
||
|
||
Adds 1x1 conv aux heads at the 1/2- and 1/4-resolution decoder outputs
|
||
(the 1/8 bottleneck is dropped, mirroring nnU-Net's zero-weight lowest
|
||
scale). In training mode forward returns (final, aux_half, aux_quarter);
|
||
in eval mode it returns just the final logits so MONAI sliding-window
|
||
inference works unchanged.
|
||
"""
|
||
|
||
def __init__(self, device, norm="instance"):
|
||
super().__init__(
|
||
spatial_dims=3, in_channels=1, out_channels=2,
|
||
channels=(16, 32, 64, 128), strides=(2, 2, 2), num_res_units=2,
|
||
act=("LEAKYRELU", {"inplace": True, "negative_slope": 0.01}),
|
||
norm=("instance", {"affine": True, "eps": 1e-5}) if norm == "instance" else "batch",
|
||
)
|
||
# locate the up-paths: model[1].submodule = [down2, SkipConn(S3), up2]
|
||
s2 = self.model[1].submodule
|
||
self._up_half = s2[2] # 16ch @ 1/2
|
||
s3 = s2[1].submodule # [down3, SkipConn(bottom), up3]
|
||
self._up_quarter = s3[2] # 32ch @ 1/4
|
||
self.aux_half = nn.Conv3d(16, 2, 1)
|
||
self.aux_quarter = nn.Conv3d(32, 2, 1)
|
||
self._aux = {}
|
||
self._up_half.register_forward_hook(lambda m, i, o: self._aux.__setitem__("half", o))
|
||
self._up_quarter.register_forward_hook(lambda m, i, o: self._aux.__setitem__("quarter", o))
|
||
from src.unet3d import init_b
|
||
self.apply(init_b)
|
||
|
||
def forward(self, x):
|
||
self._aux = {}
|
||
out = super().forward(x)
|
||
if not self.training:
|
||
return out
|
||
return out, self.aux_half(self._aux["half"]), self.aux_quarter(self._aux["quarter"])
|
||
|
||
|
||
class WeightedDiceCELoss(nn.Module):
|
||
"""Per-sample (soft-dice + CE), weighted batch mean.
|
||
|
||
Same convention as Pipeline A's src/losses.py per_sample_loss (dice over
|
||
background+foreground averaged, dice_weight 0.5), built on the MONAI
|
||
DiceLoss. Per-sample reduction is required so accepted pseudo-label rows
|
||
can be down-weighted individually via row "w"; a batch-level loss
|
||
(MONAI DiceCELoss / nnU-Net CE+Dice) cannot express that.
|
||
"""
|
||
|
||
def __init__(self, dice_weight=0.5):
|
||
super().__init__()
|
||
from monai.losses import DiceLoss
|
||
self.dice_weight = dice_weight
|
||
self.dice = DiceLoss(include_background=True, softmax=False, reduction="none")
|
||
|
||
def forward(self, logits, label, w):
|
||
logits = logits.float()
|
||
c = logits.size(1)
|
||
one = nn.functional.one_hot(label, c).permute(0, -1, *range(1, label.ndim)).float()
|
||
p = nn.functional.softmax(logits, dim=1)
|
||
dsc = self.dice(p, one) # (B, C, 1, 1, 1)
|
||
dsc = dsc.squeeze(-1).squeeze(-1).squeeze(-1).mean(dim=1) # (B,)
|
||
ce = nn.functional.cross_entropy(logits, label, reduction="none") # (B, z, y, x)
|
||
ce = ce.mean(dim=tuple(range(1, ce.ndim))) # (B,)
|
||
per = (1 - self.dice_weight) * ce + self.dice_weight * dsc
|
||
if torch.isnan(per).any() or torch.isinf(per).any():
|
||
per = torch.zeros_like(per) # zero-gradient fallback: keeps DDP collectives in sync
|
||
return (per * w).sum() / w.sum().clamp(min=1e-6)
|
||
|
||
|
||
# ---------------- data ----------------
|
||
|
||
class NNDCeLoss(nn.Module):
|
||
"""nnU-Net-style DC_and_CE (same formulation as src.losses.nn_dc_and_ce):
|
||
unweighted CE + batch-level soft dice over both classes, smooth=1.0,
|
||
loss = CE - mean(dice). Used for the A/C-vs-B loss-parity runs."""
|
||
|
||
def forward(self, logits, label):
|
||
logits = logits.float()
|
||
ce = nn.functional.cross_entropy(logits, label)
|
||
p = nn.functional.softmax(logits, dim=1)
|
||
g = nn.functional.one_hot(label, p.size(1)).permute(0, -1, *range(1, label.ndim)).float()
|
||
axes = (0, *range(2, p.dim()))
|
||
tp = (p * g).sum(axes)
|
||
fp = (p * (1.0 - g)).sum(axes)
|
||
fn = ((1.0 - p) * g).sum(axes)
|
||
dc = (2.0 * tp + 1.0) / (2.0 * tp + fp + fn + 1.0)
|
||
loss = ce - dc.mean()
|
||
if not torch.isfinite(loss):
|
||
return torch.zeros_like(loss) # zero-gradient fallback (DDP parity)
|
||
return loss
|
||
|
||
|
||
class _WorkerRNG:
|
||
"""Per-dataloader-worker numpy RandomState (forked workers reseed independently)."""
|
||
|
||
@staticmethod
|
||
def get(seed=0):
|
||
info = torch.utils.data.get_worker_info()
|
||
wid = info.id if info is not None else 0
|
||
return np.random.RandomState(seed + 104729 * wid)
|
||
|
||
|
||
class ZScoreAll:
|
||
"""Per-volume Z-score over all voxels (Pipeline B's ZScoreNormalization, no
|
||
mask). Applied to channel-first pimg."""
|
||
|
||
def __call__(self, data):
|
||
from src.augment_b import zscore_volume
|
||
d = dict(data)
|
||
img = d["pimg"]
|
||
v = img[0] if img.ndim == 4 else img
|
||
v = zscore_volume(np.asarray(v, dtype=np.float32))
|
||
d["pimg"] = v[np.newaxis] if img.ndim == 4 else v
|
||
return d
|
||
|
||
|
||
class BAugmentB:
|
||
"""Pipeline B's full augmentation (src.augment_b), applied to the sampled
|
||
96^3 patch (B's SpatialTransform deforms the sampled patch about its
|
||
center and crops back to the final size)."""
|
||
|
||
def __call__(self, data):
|
||
from src.augment_b import augment_b
|
||
d = dict(data)
|
||
img = np.asarray(d["pimg"], dtype=np.float32)
|
||
raw_lab = d.get("plabel")
|
||
lab = None if raw_lab is None else np.asarray(raw_lab)
|
||
lab_4d = lab is not None and lab.ndim == 4
|
||
if lab_4d:
|
||
lab = lab[0]
|
||
out_img, out_lab = augment_b(img, lab, _WorkerRNG.get())
|
||
d["pimg"] = out_img if img.ndim == 4 else out_img[np.newaxis]
|
||
if out_lab is not None:
|
||
d["plabel"] = out_lab[np.newaxis] if lab_4d else out_lab
|
||
return d
|
||
|
||
|
||
def train_transform(win=WIN, b_aligned=False, no_aug=False):
|
||
"""Channel-first MONAI chain: load → min-pad → augment → pos/neg crop → typed tensors.
|
||
|
||
Legacy (default): mirrors Pipeline A's augmentation set (per-axis flip p=0.5,
|
||
rot90 in the (y,x) plane p=0.1, brightness 1±0.1 p=0.3, Gaussian noise
|
||
σ=0.01 p=0.4). Patch sampling uses MONAI's foreground-aware
|
||
RandCropByPosNegLabeld (pos=neg=1).
|
||
|
||
b_aligned: Pipeline B's data protocol — per-volume Z-score, B's full
|
||
augmentation (rot/zoom/noise/blur/brightness/contrast/lowres/gamma/mirror),
|
||
and 33% foreground patch sampling (pos=neg ratio 0.5:1, like B's
|
||
oversample_foreground_percent=0.33).
|
||
"""
|
||
from monai.transforms import (
|
||
Compose, LoadImaged, EnsureChannelFirstd, EnsureTyped, SpatialPadd,
|
||
RandFlipd, RandRotate90d, RandScaleIntensityd, RandGaussianNoised,
|
||
RandCropByPosNegLabeld, Compose as _C)
|
||
if b_aligned:
|
||
crop_pos, crop_neg = 0.5, 1.0
|
||
# B's SpatialTransform augments the SAMPLED patch (rot/zoom about its
|
||
# center, then crop back to final size) — not the whole volume; that
|
||
# is also what Pipeline A's b_aug does, and whole-volume deformation
|
||
# is ~30x slower per item.
|
||
steps = [
|
||
LoadImaged(keys=["pimg", "plabel"]),
|
||
EnsureChannelFirstd(keys=["pimg", "plabel"]),
|
||
SpatialPadd(keys=["pimg", "plabel"], spatial_size=(win, win, win)),
|
||
ZScoreAll(),
|
||
RandCropByPosNegLabeld(keys=["pimg", "plabel"], label_key="plabel",
|
||
spatial_size=(win, win, win), pos=crop_pos, neg=crop_neg,
|
||
num_samples=1),
|
||
]
|
||
if not no_aug:
|
||
steps.append(BAugmentB())
|
||
steps.append(EnsureTyped(keys=["pimg", "plabel"], dtype=[torch.float32, torch.long]))
|
||
return Compose(steps)
|
||
return Compose([
|
||
LoadImaged(keys=["pimg", "plabel"]),
|
||
EnsureChannelFirstd(keys=["pimg", "plabel"]),
|
||
SpatialPadd(keys=["pimg", "plabel"], spatial_size=(win, win, win)),
|
||
RandFlipd(keys=["pimg", "plabel"], prob=0.5, spatial_axis=0),
|
||
RandFlipd(keys=["pimg", "plabel"], prob=0.5, spatial_axis=1),
|
||
RandFlipd(keys=["pimg", "plabel"], prob=0.5, spatial_axis=2),
|
||
RandRotate90d(keys=["pimg", "plabel"], prob=0.1, max_k=3, spatial_axes=(1, 2)),
|
||
RandScaleIntensityd(keys=["pimg"], factors=0.1, prob=0.3),
|
||
RandGaussianNoised(keys=["pimg"], std=0.01, prob=0.4),
|
||
RandCropByPosNegLabeld(keys=["pimg", "plabel"], label_key="plabel",
|
||
spatial_size=(win, win, win), pos=1, neg=1, num_samples=1),
|
||
EnsureTyped(keys=["pimg", "plabel"], dtype=[torch.float32, torch.long]),
|
||
])
|
||
|
||
|
||
def rows_for_dataset(rows):
|
||
out = []
|
||
for r in rows:
|
||
out.append({"key": r["key"], "pimg": r["pimg"], "plabel": r["plabel"],
|
||
"w": float(r.get("w", 1.0))})
|
||
return out
|
||
|
||
|
||
def _collate(batch):
|
||
# RandCropByPosNegLabeld(num_samples=1) yields a 1-item list per sample
|
||
items = [b[0] if isinstance(b, (list, tuple)) else b for b in batch]
|
||
return torch.utils.data.default_collate(items)
|
||
|
||
|
||
def make_dataloader(rows, win=WIN, batch=3, workers=4, b_aligned=False):
|
||
from monai.data import Dataset
|
||
ds = Dataset(data=rows_for_dataset(rows), transform=train_transform(win, b_aligned=b_aligned))
|
||
return torch.utils.data.DataLoader(
|
||
ds, batch_size=batch, shuffle=True, num_workers=workers,
|
||
collate_fn=_collate, drop_last=len(ds) > batch, pin_memory=True,
|
||
persistent_workers=workers > 0)
|
||
|
||
|
||
# ---------------- inference ----------------
|
||
|
||
@torch.no_grad()
|
||
def predict_probs(model, vol, device, win=WIN, overlap=0.5, tta=True, sw_batch=8,
|
||
zscore=False):
|
||
"""Tumor probability map (C, z, y, x → (z, y, x) numpy) for one volume.
|
||
|
||
MONAI sliding_window_inference (gaussian blend, `overlap`), plus the same
|
||
4-view flip TTA (identity + 3 axis flips) as Pipeline A. zscore=True matches
|
||
a B-aligned model's input (per-volume Z-score, no [0,1.5] clip).
|
||
"""
|
||
from monai.inferers import sliding_window_inference
|
||
a = np.nan_to_num(np.asarray(vol, dtype=np.float32), nan=0.0, posinf=1.5, neginf=0.0)
|
||
if zscore:
|
||
from src.augment_b import zscore_volume
|
||
a = zscore_volume(a)
|
||
else:
|
||
a = np.clip(a, 0.0, 1.5)
|
||
t = torch.from_numpy(a).unsqueeze(0).unsqueeze(0).to(device)
|
||
|
||
def run(vt):
|
||
with torch.autocast("cuda", dtype=torch.bfloat16):
|
||
logits = sliding_window_inference(
|
||
vt, roi_size=(win, win, win), sw_batch_size=sw_batch,
|
||
predictor=model, mode="gaussian", overlap=overlap,
|
||
sw_device=device, device=device)
|
||
return torch.softmax(logits.float(), dim=1)[0, 1].cpu().numpy()
|
||
|
||
acc = run(t)
|
||
if not tta:
|
||
return acc
|
||
for ax in (0, 1, 2):
|
||
acc += np.flip(run(t.flip(ax + 2)), axis=ax)
|
||
return acc / 4.0
|
||
|
||
|
||
def prob_dice(probs, lab, thr=0.5):
|
||
"""Dice at probability threshold 0.5 — the convention of Pipeline A's
|
||
08_eval and Pipeline B's primary metric."""
|
||
return dice((probs >= thr).astype("uint8"), (lab > 0).astype("uint8"))
|
||
|
||
|
||
def ckpt_b_aligned(sd):
|
||
"""Whether a checkpoint was trained with the B-aligned config (instance norm +
|
||
deep-sup aux heads). New checkpoints carry an explicit 'b_aligned' marker;
|
||
pre-marker ones fall back to the aux_head state-dict keys."""
|
||
if not isinstance(sd, dict):
|
||
return False
|
||
if "b_aligned" in sd:
|
||
return bool(sd["b_aligned"])
|
||
inner = sd.get("model", sd)
|
||
return isinstance(inner, dict) and any(k.startswith("aux_") for k in inner)
|
||
|
||
|
||
def load_model(ckpt, device):
|
||
sd = torch.load(ckpt, map_location=device, weights_only=True)
|
||
ba = ckpt_b_aligned(sd)
|
||
net = build_model(device, norm="instance" if ba else "batch", b_aligned=ba)
|
||
net.load_state_dict(sd.get("model", sd))
|
||
net.eval()
|
||
return net, sd
|
||
|
||
|
||
# ---------------- distributed (torchrun, as in Pipeline A) ----------------
|
||
|
||
def rank_info():
|
||
return (int(os.environ.get("RANK", 0)), int(os.environ.get("WORLD_SIZE", 1)),
|
||
int(os.environ.get("LOCAL_RANK", 0)))
|
||
|
||
|
||
def init_dist():
|
||
if int(os.environ.get("WORLD_SIZE", 1)) > 1:
|
||
from datetime import timedelta
|
||
import torch.distributed as dist
|
||
dist.init_process_group("nccl", timeout=timedelta(minutes=30))
|
||
|
||
|
||
def barrier():
|
||
if int(os.environ.get("WORLD_SIZE", 1)) > 1:
|
||
import torch.distributed as dist
|
||
dist.barrier()
|
||
|
||
|
||
def destroy_dist():
|
||
if int(os.environ.get("WORLD_SIZE", 1)) > 1:
|
||
import torch.distributed as dist
|
||
dist.destroy_process_group() |