| """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) |
| 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() |
|
|
