classification/train_resnet.py

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