"""Fine-tune FlexiCT-2D with LoRA on CT-RATE for 18-label abnormality classification. Recipe used for all three released models. They differ in --pool, --z-spacing, --max-slices and the head width (512 for the 1.5 mm models, --head-dim 256 --pool-d-attn 256 for the 2.0 mm model): * backbone frozen; LoRA r=8, alpha=16, dropout 0.05 on attention projections and the conv branch; LayerNorm and layer-scale parameters also trained * head, slice position code and conv branch at lr 4e-4 (weight decay 1e-4), everything else trainable at lr 1.4e-4 (weight decay 1e-5); fused AdamW * 400 warm-up steps, then cosine towards 5 % over a nominal 30,000 steps, stopped after 3 epochs of 2,880 optimizer steps (8,640 steps) * micro-batch 1, gradient accumulation to an effective batch of 16 across GPUs; each epoch draws 46,080 training volumes at random * asymmetric loss (gamma_neg 4, gamma_pos 0, clip 0.05); bf16 autocast; gradient checkpointing; no data augmentation torchrun --standalone --nproc_per_node=4 train.py --volumes ctrate/train \ --labels train_predicted_labels.csv --out runs/z1p5_mhead --pool mhead \ --z-spacing 1.5 --max-slices 240 """ from __future__ import annotations import argparse import csv import json import math import os from pathlib import Path import nibabel as nib import numpy as np import torch import torch.distributed as dist from model import (HEAD_PREFIXES, HEAD_SUBSTRINGS, CTClassifier, FlexiCT2DEncoder, head_kwargs, preprocess_ct) LABELS = ["Medical material", "Arterial wall calcification", "Cardiomegaly", "Pericardial effusion", "Coronary artery wall calcification", "Hiatal hernia", "Lymphadenopathy", "Emphysema", "Atelectasis", "Lung nodule", "Lung opacity", "Pulmonary fibrotic sequela", "Pleural effusion", "Mosaic attenuation pattern", "Peribronchial thickening", "Consolidation", "Bronchiectasis", "Interlobular septal thickening"] EFFECTIVE_BATCH, STEPS_PER_EPOCH, EPOCHS = 16, 2880, 3 WARMUP, NOMINAL_STEPS, MIN_LR_FRAC = 400, 30_000, 0.05 class AsymmetricLoss(torch.nn.Module): def __init__(self, gamma_neg=4.0, gamma_pos=0.0, clip=0.05, eps=1e-8): super().__init__() self.gn, self.gp, self.clip, self.eps = gamma_neg, gamma_pos, clip, eps def forward(self, logits, y): y = y.float() p = torch.sigmoid(logits) q = (1.0 - p + self.clip).clamp(max=1.0) if self.clip > 0 else 1.0 - p loss = y * torch.log(p.clamp(min=self.eps)) + (1 - y) * torch.log(q.clamp(min=self.eps)) pt = p * y + q * (1 - y) loss = loss * torch.pow(1.0 - pt.detach(), self.gp * y + self.gn * (1 - y)) return -loss.mean() def lr_factor(step: int) -> float: if step < WARMUP: return step / max(1, WARMUP) progress = min(1.0, (step - WARMUP) / (NOMINAL_STEPS - WARMUP)) return MIN_LR_FRAC + (1 - MIN_LR_FRAC) * 0.5 * (1 + math.cos(math.pi * progress)) class CTRateVolumes(torch.utils.data.Dataset): """CT-RATE training volumes (*.nii.gz under --volumes) with labels from the CSV (VolumeName + 18 columns).""" def __init__(self, root, labels_csv, geo, cache_dir=None): paths = {p.name: p for p in Path(root).rglob("*.nii.gz")} self.items = [] for row in csv.DictReader(open(labels_csv)): name = row["VolumeName"] if row["VolumeName"].endswith(".nii.gz") else row["VolumeName"] + ".nii.gz" if name in paths: self.items.append((paths[name], np.array([float(row[l]) for l in LABELS], dtype=np.float32))) self.geo = geo self.cache = Path(cache_dir) if cache_dir else None def __len__(self): return len(self.items) def __getitem__(self, i): path, y = self.items[i] cached = self.cache / f"{path.name[:-7]}.npy" if self.cache else None if cached is not None and cached.is_file(): x = torch.from_numpy(np.load(cached)) else: img = nib.load(str(path)) x = preprocess_ct(np.asanyarray(img.dataobj).astype(np.float32), np.asarray(img.affine, dtype=np.float32), **self.geo) if cached is not None: tmp = cached.with_name(f".{cached.name}.{os.getpid()}.npy") np.save(tmp, x.numpy()) os.replace(tmp, cached) return x, torch.from_numpy(y) def build(args): from peft import LoraConfig, get_peft_model enc = FlexiCT2DEncoder(img_size=480, patch_size=16, max_slices=args.max_slices, slice_microbatch=args.max_slices, conv_branch=True, grad_checkpointing=True) dim = enc.embed_dim enc = get_peft_model(enc, LoraConfig(r=8, lora_alpha=16, lora_dropout=0.05, bias="none", target_modules=list(enc.lora_target_modules))) for name, p in enc.named_parameters(): n = name.lower() p.requires_grad_(any(t in n for t in ("lora_", "norm", "ls1", "ls2", "conv_branch", "cb_gamma"))) cfg = dict(model="flexict2d", head="radiology", pool=args.pool, head_heads=4, aggregation="classwise", class_head="mlp", use_lora=True, lora_r=8, lora_alpha=16, lora_dropout=0.05, img_size=480, patch_size=16, resolution=None, max_slices=args.max_slices, slice_microbatch=args.max_slices, pool_d_attn=args.pool_d_attn, head_dropout=0.1, token_mode="all", slice_posenc=True, head_dim=args.head_dim, head_layerscale=1e-4, head_identity_proj=True, xy_mode="spacing", xy_spacing=0.75, z_spacing=args.z_spacing, preproc_dtype="float16", conv_branch=True) model = CTClassifier(enc, dim, len(LABELS), slice_posenc=True, **head_kwargs(cfg)) return model, cfg def main(): ap = argparse.ArgumentParser() ap.add_argument("--volumes", required=True) ap.add_argument("--labels", required=True) ap.add_argument("--out", type=Path, required=True) ap.add_argument("--pool", default="mhead", choices=("mhead", "hybrid")) ap.add_argument("--z-spacing", type=float, default=1.5) ap.add_argument("--max-slices", type=int, default=240) ap.add_argument("--head-dim", type=int, default=512) ap.add_argument("--pool-d-attn", type=int, default=512) ap.add_argument("--cache-dir", help="optional folder for preprocessed volumes (.npy)") ap.add_argument("--seed", type=int, default=0) ap.add_argument("--workers", type=int, default=16) ap.add_argument("--save-every", type=int, default=250, help="optimizer steps between resume checkpoints") ap.add_argument("--max-steps", type=int, default=0, help="stop early (smoke tests only)") args = ap.parse_args() world, rank, local = int(os.environ.get("WORLD_SIZE", 1)), int(os.environ.get("RANK", 0)), int(os.environ.get("LOCAL_RANK", 0)) if world > 1: dist.init_process_group("nccl") torch.cuda.set_device(local) device = torch.device("cuda", local) torch.manual_seed(args.seed) np.random.seed(args.seed) accum = EFFECTIVE_BATCH // world assert accum * world == EFFECTIVE_BATCH, "world size must divide 16" model, cfg = build(args) slow, fast = [], [] for name, p in model.named_parameters(): if p.requires_grad: (fast if name.startswith(HEAD_PREFIXES) or any(s in name for s in HEAD_SUBSTRINGS) else slow).append(p) model = model.to(device) ddp = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local], broadcast_buffers=False) if world > 1 else model torch.manual_seed(args.seed + 1_000_003 * rank) # distinct dropout masks per rank opt = torch.optim.AdamW([{"params": slow, "lr": 1.4e-4, "weight_decay": 1e-5}, {"params": fast, "lr": 4e-4, "weight_decay": 1e-4}], fused=True) sched = torch.optim.lr_scheduler.LambdaLR(opt, [lr_factor, lr_factor]) crit = AsymmetricLoss() geo = dict(img_size=480, xy_spacing=0.75, z_spacing=args.z_spacing, max_slices=args.max_slices) data = CTRateVolumes(args.volumes, args.labels, geo, args.cache_dir) per_epoch = STEPS_PER_EPOCH * EFFECTIVE_BATCH step, epoch, cursor = 0, 0, 0 last = args.out / "last.pt" if last.is_file(): ck = torch.load(last, map_location="cpu", weights_only=False) model.load_state_dict(ck["trainable"], strict=False) opt.load_state_dict(ck["opt"]) sched.load_state_dict(ck["sched"]) step, epoch, cursor = ck["step"], ck["epoch"], ck["cursor"] def save(path, full=False): if rank: return args.out.mkdir(parents=True, exist_ok=True) if full: payload = {"model": model.state_dict(), "config": cfg, "epoch": epoch, "step": step} else: names = {n for n, p in model.named_parameters() if p.requires_grad} trainable = {k: v for k, v in model.state_dict().items() if k in names} payload = {"trainable": trainable, "opt": opt.state_dict(), "sched": sched.state_dict(), "step": step, "epoch": epoch, "cursor": cursor} tmp = path.with_name(f".{path.name}.part") torch.save(payload, tmp) os.replace(tmp, path) ddp.train() while epoch < EPOCHS: order = np.random.default_rng(args.seed + 1_000_003 * epoch).permutation(len(data))[:per_epoch] mine = order[rank::world] sub = torch.utils.data.Subset(data, mine[cursor // world:].tolist()) loader = torch.utils.data.DataLoader(sub, batch_size=1, num_workers=args.workers, pin_memory=True, prefetch_factor=2, persistent_workers=False) running = 0.0 for i, (x, y) in enumerate(loader, 1): x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True) sync = i % accum == 0 ctx = ddp.no_sync() if (world > 1 and not sync) else torch.enable_grad() with ctx, torch.autocast("cuda", dtype=torch.bfloat16): loss = crit(ddp(x).float(), y) / accum loss.backward() running += float(loss) if sync: opt.step() sched.step() opt.zero_grad(set_to_none=True) step += 1 cursor += EFFECTIVE_BATCH if rank == 0 and step % 50 == 0: print(f"epoch {epoch} step {step} loss {running / 50:.4f} lr {sched.get_last_lr()[0]:.2e}", flush=True) running = 0.0 if step % args.save_every == 0: save(last) if args.max_steps and step >= args.max_steps: save(last) return save(args.out / f"epoch{epoch}.pt", full=True) epoch, cursor = epoch + 1, 0 save(last) if rank == 0: (args.out / "config.json").write_text(json.dumps(cfg, indent=1)) if world > 1: dist.destroy_process_group() if __name__ == "__main__": main()