""" ResNet50 multi-label classification for chest X-ray (13 diseases). 数据集 jsonl 格式 (每行): {"file_name": "...png", "attn_list": [["Consolidation", "...mask.png"], ...]} - 输入: 512x512 灰度 PNG (0-255) - 标签: 13 维 multi-hot,从 attn_list[*][0] 提取;Normal 样本所有维度为 0 - Loss: BCEWithLogitsLoss - Metric: 每个类的 AUROC + mAP (macro 平均) 用法: python train_resnet50_cls.py \\ --exp_name resnet50_real \\ --data_dir /home/jovyan \\ --train_prompt /home/jovyan/AURAD/detection/json/train_real.json \\ --val_prompt /home/jovyan/AURAD/detection/json/val_real.json \\ --batch_size 32 --lr 1e-4 --epochs 50 """ import argparse import json import os from pathlib import Path import numpy as np import torch import torch.nn as nn import torch.optim as optim import torchvision import torchvision.transforms as T import wandb from PIL import Image from sklearn.metrics import roc_auc_score, average_precision_score from torch.utils.data import DataLoader, Dataset DISEASES = [ "Atelectasis", "Calcification", "Cardiomegaly", "Consolidation", "Diffuse Nodule", "Effusion", "Emphysema", "Fibrosis", "Fracture", "Mass", "Nodule", "Pleural Thickening", "Pneumothorax", ] DISEASE2IDX = {d: i for i, d in enumerate(DISEASES)} NUM_CLASSES = len(DISEASES) # ---------------- Dataset ---------------- class ChestXRayClsDataset(Dataset): """读取 jsonl, 从 attn_list 提取 multi-hot label. Normal (或任何不在 DISEASES 里的 label) 一律视为全 0 label.""" def __init__(self, jsonl_path, data_dir, transform=None): self.data_dir = Path(data_dir) self.items = [] with open(jsonl_path) as f: for line in f: line = line.strip() if line: self.items.append(json.loads(line)) self.transform = transform def __len__(self): return len(self.items) def __getitem__(self, idx): item = self.items[idx] img_path = self.data_dir / item["file_name"] img = Image.open(img_path).convert("L") # 灰度 label = torch.zeros(NUM_CLASSES, dtype=torch.float32) for entry in item.get("attn_list", []): name = entry[0] if name in DISEASE2IDX: label[DISEASE2IDX[name]] = 1.0 # Normal 或其它名字 -> 不动 label,保持全 0 if self.transform: img = self.transform(img) return img, label # ---------------- Model ---------------- def get_resnet50(num_classes=NUM_CLASSES, in_channels=1): """torchvision resnet50, 改单通道输入 + 改分类头. 不加载 ImageNet 权重 (从头训).""" model = torchvision.models.resnet50(weights=None) # 改第一层为单通道 if in_channels != 3: model.conv1 = nn.Conv2d( in_channels, 64, kernel_size=7, stride=2, padding=3, bias=False ) # 改分类头 model.fc = nn.Linear(model.fc.in_features, num_classes) return model # ---------------- Train / Val ---------------- def train_one_epoch(model, dataloader, optimizer, criterion, device, epoch): model.train() total_loss = 0.0 n_batches = 0 for images, labels in dataloader: images = images.to(device, non_blocking=True) labels = labels.to(device, non_blocking=True) logits = model(images) loss = criterion(logits, labels) optimizer.zero_grad() loss.backward() optimizer.step() total_loss += loss.item() n_batches += 1 avg = total_loss / max(n_batches, 1) return avg @torch.no_grad() def validate(model, dataloader, criterion, device): model.eval() total_loss = 0.0 n_batches = 0 all_probs, all_labels = [], [] for images, labels in dataloader: images = images.to(device, non_blocking=True) labels = labels.to(device, non_blocking=True) logits = model(images) loss = criterion(logits, labels) total_loss += loss.item() n_batches += 1 all_probs.append(torch.sigmoid(logits).cpu().numpy()) all_labels.append(labels.cpu().numpy()) probs = np.concatenate(all_probs, 0) labels = np.concatenate(all_labels, 0) # 每类 AUROC / AP, 只在有正样本的类上计算 per_class = {} aurocs, aps = [], [] for i, dis in enumerate(DISEASES): y, p = labels[:, i], probs[:, i] if y.sum() == 0 or y.sum() == len(y): # 没正样本或全是正样本, 跳过 per_class[dis] = {"auroc": np.nan, "ap": np.nan, "n_pos": int(y.sum())} continue auc = roc_auc_score(y, p) ap = average_precision_score(y, p) per_class[dis] = {"auroc": auc, "ap": ap, "n_pos": int(y.sum())} aurocs.append(auc) aps.append(ap) macro_auroc = float(np.mean(aurocs)) if aurocs else 0.0 macro_ap = float(np.mean(aps)) if aps else 0.0 avg_loss = total_loss / max(n_batches, 1) return avg_loss, macro_auroc, macro_ap, per_class def save_full_ckpt(path, model, optimizer, scheduler, epoch, best_auroc): """保存可续训的完整状态。epoch 存的是“已完成的 epoch 数”,即下次的起点。""" torch.save({ "epoch": epoch, "model": model.state_dict(), "optimizer": optimizer.state_dict(), "scheduler": scheduler.state_dict(), "best_auroc": best_auroc, }, path) # ---------------- Main ---------------- def parse_args(): ap = argparse.ArgumentParser() ap.add_argument("--exp_name", default="resnet50_cls") ap.add_argument("--data_dir", required=True) ap.add_argument("--train_prompt", required=True) ap.add_argument("--val_prompt", required=True) ap.add_argument("--img_resolution", type=int, default=512) ap.add_argument("--batch_size", type=int, default=32) ap.add_argument("--lr", type=float, default=1e-4) ap.add_argument("--epochs", type=int, default=50) ap.add_argument("--val_epoch", type=int, default=1) ap.add_argument("--checkpoint_save_epoch", type=int, default=50) ap.add_argument("--num_workers", type=int, default=8) ap.add_argument("--weight_decay", type=float, default=1e-4) ap.add_argument("--optim", choices=["adamw", "sgd"], default="adamw") ap.add_argument("--resume", default="", help="要续训的 checkpoint 路径,留空则从头训") return ap.parse_args() def main(): args = parse_args() device = torch.device("cuda" if torch.cuda.is_available() else "cpu") wandb.init(project="classification", name=args.exp_name) ckpt_dir = Path(f"./checkpoint/{args.exp_name}") ckpt_dir.mkdir(parents=True, exist_ok=True) # ------- Transforms ------- # 训练带轻量增强, 验证只 resize + ToTensor train_tf = T.Compose([ T.Resize((args.img_resolution, args.img_resolution)), T.RandomHorizontalFlip(p=0.5), T.RandomAffine(degrees=10, translate=(0.05, 0.05), scale=(0.95, 1.05)), T.ToTensor(), # [0,1], shape (1, H, W) T.Normalize(mean=[0.5], std=[0.5]), # 灰度图标准化 ]) val_tf = T.Compose([ T.Resize((args.img_resolution, args.img_resolution)), T.ToTensor(), T.Normalize(mean=[0.5], std=[0.5]), ]) # ------- Data ------- train_set = ChestXRayClsDataset(args.train_prompt, args.data_dir, transform=train_tf) val_set = ChestXRayClsDataset(args.val_prompt, args.data_dir, transform=val_tf) print(f"Train: {len(train_set)}, Val: {len(val_set)}") # 算每类正样本数, 打印 + 算 pos_weight (类别极不平衡时有用) label_sum = np.zeros(NUM_CLASSES, dtype=np.int64) for it in train_set.items: for entry in it.get("attn_list", []): if entry[0] in DISEASE2IDX: label_sum[DISEASE2IDX[entry[0]]] += 1 print("Train pos counts per class:") for i, d in enumerate(DISEASES): print(f" {d:20s}: {label_sum[i]}") pos_weight = torch.tensor( [(len(train_set) - c) / max(c, 1) for c in label_sum], dtype=torch.float32, device=device, ) train_loader = DataLoader( train_set, batch_size=args.batch_size, shuffle=True, num_workers=args.num_workers, pin_memory=True, drop_last=True, ) val_loader = DataLoader( val_set, batch_size=args.batch_size, shuffle=False, num_workers=args.num_workers, pin_memory=True, ) # ------- Model / Loss / Optim ------- model = get_resnet50(NUM_CLASSES, in_channels=1).to(device) criterion = nn.BCEWithLogitsLoss(pos_weight=pos_weight) if args.optim == "adamw": optimizer = optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) else: optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=args.weight_decay) scheduler = optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=args.epochs, eta_min=1e-6, ) # ------- Resume ------- start_epoch = 0 best_auroc = 0.0 if args.resume: print(f"Resuming from {args.resume}") ckpt = torch.load(args.resume, map_location=device) if isinstance(ckpt, dict) and "model" in ckpt: # 完整 checkpoint model.load_state_dict(ckpt["model"]) optimizer.load_state_dict(ckpt["optimizer"]) scheduler.load_state_dict(ckpt["scheduler"]) start_epoch = ckpt["epoch"] best_auroc = ckpt.get("best_auroc", 0.0) print(f" -> resumed at epoch {start_epoch}, best_AUROC={best_auroc:.4f}") else: # 只含权重的老 checkpoint(比如 best.pth),只能恢复模型 model.load_state_dict(ckpt) print(" -> 只加载了模型权重(无 optimizer/epoch 状态,LR 调度会从头算)") # ------- Train loop ------- # best_auroc = 0.0 for epoch in range(start_epoch, args.epochs): train_loss = train_one_epoch(model, train_loader, optimizer, criterion, device, epoch) scheduler.step() cur_lr = optimizer.param_groups[0]["lr"] print(f"[Epoch {epoch+1}/{args.epochs}] train_loss={train_loss:.4f} lr={cur_lr:.2e}") wandb.log({"epoch": epoch + 1, "train_loss": train_loss, "lr": cur_lr}) if (epoch + 1) % args.val_epoch == 0: val_loss, macro_auc, macro_ap, per_cls = validate(model, val_loader, criterion, device) print(f" val_loss={val_loss:.4f} macro_AUROC={macro_auc:.4f} macro_AP={macro_ap:.4f}") log = { "val_loss": val_loss, "macro_AUROC": macro_auc, "macro_AP": macro_ap, } for d, m in per_cls.items(): if not np.isnan(m["auroc"]): log[f"AUROC/{d}"] = m["auroc"] log[f"AP/{d}"] = m["ap"] wandb.log(log) if macro_auc > best_auroc: best_auroc = macro_auc torch.save(model.state_dict(), ckpt_dir / "best.pth") print(f" -> saved best.pth (macro_AUROC={best_auroc:.4f})") if (epoch + 1) % args.checkpoint_save_epoch == 0: save_full_ckpt( ckpt_dir / f"checkpoint_epoch_{epoch+1}.pth", model, optimizer, scheduler, epoch + 1, best_auroc, ) wandb.finish() if __name__ == "__main__": main()