train.py

"""
Faster R-CNN 밑바닥 구현 — [5/5] 학습 + 평가
================================================
전체 파이프라인:
  학습:  이미지 → 모델(train) → RPN 손실 + RoI 손실 → 역전파
  평가:  이미지 → 모델(eval) → 탐지 결과 → mAP@0.5 계산

실행 예:
  python train.py --voc_root /path/VOCdevkit/VOC2007 --epochs 12
  python train.py --voc_root /path/VOCdevkit/VOC2007 --eval_only --ckpt frcnn.pth

주의:
  - batch_size=1 로 설계(이미지 크기가 제각각이라 단순화).
  - GPU 권장. CPU로도 돌지만 매우 느리다.
"""

import argparse
import torch
from torch.utils.data import DataLoader

from dataset import VOCDataset, collate_fn, NUM_CLASSES, VOC_CLASSES
from model import FasterRCNN
from losses import rpn_loss, roi_loss
from box_utils import box_iou


# ---------------------------------------------------------------
# 학습 한 에폭
# ---------------------------------------------------------------
def train_one_epoch(model, loader, optimizer, device, epoch):
    model.train()
    running = 0.0
    for i, (imgs, targets) in enumerate(loader):
        img = imgs[0].to(device).unsqueeze(0)          # [1,3,H,W]
        gt_boxes = targets[0]["boxes"].to(device)
        gt_labels = targets[0]["labels"].to(device)
        if gt_boxes.numel() == 0:
            continue

        out = model(img)  # training=True → 중간 산출물 반환
        img_hw = img.shape[-2:]

        l_rpn = rpn_loss(out["rpn_logits"], out["rpn_deltas"],
                         out["anchors"], gt_boxes, img_hw)
        l_roi = roi_loss(model.head, out["feat"], out["stride"],
                         out["proposals"], gt_boxes, gt_labels)
        loss = l_rpn + l_roi

        optimizer.zero_grad()
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 10.0)  # 폭주 방지
        optimizer.step()

        running += loss.item()
        if (i + 1) % 100 == 0:
            print(f"[epoch {epoch}] iter {i+1}/{len(loader)} "
                  f"loss {running/(i+1):.4f} (rpn {l_rpn.item():.3f} roi {l_roi.item():.3f})")
    return running / max(1, len(loader))


# ---------------------------------------------------------------
# 평가: VOC 스타일 mAP@0.5
# ---------------------------------------------------------------
@torch.no_grad()
def evaluate(model, loader, device, iou_thresh=0.5):
    model.eval()
    # 클래스별 (점수, 맞음여부) 수집 + 정답 개수
    preds = {c: [] for c in range(1, NUM_CLASSES)}
    n_gt = {c: 0 for c in range(1, NUM_CLASSES)}

    for imgs, targets in loader:
        img = imgs[0].to(device).unsqueeze(0)
        det = model(img)  # eval → {boxes, labels, scores}
        gt_boxes = targets[0]["boxes"].to(device)
        gt_labels = targets[0]["labels"].to(device)

        for c in range(1, NUM_CLASSES):
            gmask = gt_labels == c
            gboxes = gt_boxes[gmask]
            n_gt[c] += gboxes.shape[0]

            pmask = det["labels"] == c
            pboxes = det["boxes"][pmask]
            pscores = det["scores"][pmask]
            if pboxes.numel() == 0:
                continue

            order = pscores.argsort(descending=True)
            pboxes, pscores = pboxes[order], pscores[order]

            matched = torch.zeros(gboxes.shape[0], dtype=torch.bool)
            for k in range(pboxes.shape[0]):
                if gboxes.numel() == 0:
                    preds[c].append((pscores[k].item(), 0))
                    continue
                ious = box_iou(pboxes[k:k+1], gboxes).squeeze(0)
                best_iou, best_j = ious.max(0)
                if best_iou >= iou_thresh and not matched[best_j]:
                    preds[c].append((pscores[k].item(), 1))  # TP
                    matched[best_j] = True
                else:
                    preds[c].append((pscores[k].item(), 0))  # FP

    # 클래스별 AP → mAP
    aps = []
    for c in range(1, NUM_CLASSES):
        ap = _voc_ap(preds[c], n_gt[c])
        aps.append(ap)
        print(f"  {VOC_CLASSES[c-1]:12s} AP = {ap:.4f}")
    mAP = sum(aps) / len(aps)
    print(f"  {'mAP@0.5':12s} = {mAP:.4f}")
    return mAP


def _voc_ap(pred_list, n_gt):
    """(점수, TP여부) 목록으로 precision-recall 곡선 아래 넓이(AP) 계산."""
    if n_gt == 0 or len(pred_list) == 0:
        return 0.0
    pred_list.sort(key=lambda x: x[0], reverse=True)
    tp = torch.tensor([p[1] for p in pred_list], dtype=torch.float32)
    fp = 1 - tp
    tp_cum = torch.cumsum(tp, 0)
    fp_cum = torch.cumsum(fp, 0)
    recall = tp_cum / n_gt
    precision = tp_cum / (tp_cum + fp_cum).clamp(min=1e-6)

    # 11-point 보간 (VOC2007 방식)
    ap = 0.0
    for t in torch.linspace(0, 1, 11):
        mask = recall >= t
        p = precision[mask].max().item() if mask.any() else 0.0
        ap += p / 11.0
    return ap


# ---------------------------------------------------------------
# 메인
# ---------------------------------------------------------------
def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--voc_root", required=True, help="VOCdevkit/VOC2007 경로")
    ap.add_argument("--epochs", type=int, default=12)
    ap.add_argument("--lr", type=float, default=1e-3)
    ap.add_argument("--ckpt", default="frcnn.pth")
    ap.add_argument("--eval_only", action="store_true")
    args = ap.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print("device:", device)

    model = FasterRCNN(NUM_CLASSES).to(device)

    if args.eval_only:
        model.load_state_dict(torch.load(args.ckpt, map_location=device))
        test_ds = VOCDataset(args.voc_root, split="test")
        test_loader = DataLoader(test_ds, batch_size=1, shuffle=False,
                                 collate_fn=collate_fn, num_workers=4)
        evaluate(model, test_loader, device)
        return

    # 학습
    train_ds = VOCDataset(args.voc_root, split="trainval")
    train_loader = DataLoader(train_ds, batch_size=1, shuffle=True,
                              collate_fn=collate_fn, num_workers=4)

    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.SGD(params, lr=args.lr, momentum=0.9,
                                weight_decay=5e-4)
    scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=8, gamma=0.1)

    for epoch in range(1, args.epochs + 1):
        avg = train_one_epoch(model, train_loader, optimizer, device, epoch)
        scheduler.step()
        print(f"[epoch {epoch}] avg loss = {avg:.4f}")
        torch.save(model.state_dict(), args.ckpt)
        print(f"  checkpoint saved → {args.ckpt}")

    print("학습 완료. 평가하려면 --eval_only 로 실행하세요.")


if __name__ == "__main__":
    main()