[실전 예제/객체 탐지/PyTorch] Faster R-CNN 모델 구성과 COCO 학습

객체 검출(Object Detection)이란?

  객체 검출(Object Detection)은 이미지 내에서 객체의 위치(Bounding Box)와 클래스(Class)를 동시에 예측하는 컴퓨터 비전 태스크입니다. 단순히 무엇이 있는지를 넘어서, 어디에 있는지까지 함께 추론하는 것이 핵심입니다.

  Faster R-CNN은 Region Proposal Network(RPN)를 통해 후보 영역을 직접 학습으로 생성하는 2-stage 객체 검출 모델로, 정확도가 중요한 환경에서 널리 사용됩니다. 이번 글에서는 이전 글에서 구성한 COCO 데이터셋을 기반으로, PyTorch에서 Faster R-CNN 모델을 구성하고 실제 학습까지 연결하는 과정을 정리합니다.

• 예: 이미지 속 사람·차량·동물 검출
• 예: CCTV, 자율주행, 산업용 비전

 

목표

1. COCO 객체 검출 데이터셋을 이용한 학습 파이프라인 구성
2. torchvision Faster R-CNN 모델 구조 이해
3. Bounding Box / Classification loss 기반 학습 코드 작성

 

학습을 위한 데이터 구조

  이전 글에서 구성한 COCO 객체 검출 데이터셋은 COCO annotation 포맷(JSON)을 사용하며, 이미지와 bounding box 정보를 함께 포함합니다.

coco/
  train2017/
  val2017/
  annotations/
    instances_train2017.json
    instances_val2017.json

  Faster R-CNN 학습 시, 각 샘플은 다음과 같은 형태로 전달됩니다.

• boxes: (N,4) bounding box 좌표
• labels: 객체 클래스 ID
• image_id / area / iscrowd

 

Faster R-CNN 모델 구성

모델 개요

• Backbone: ResNet50 + FPN
• Region Proposal Network(RPN)
• ROIAlign 기반 Box Head

import torch
from torchvision.models.detection import fasterrcnn_resnet50_fpn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

def build_model(num_classes):
    model = fasterrcnn_resnet50_fpn(weights="DEFAULT")

    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(
        in_features,
        num_classes
    )

    return model

  COCO 데이터셋을 사용할 경우, 클래스 수는 background 포함으로 설정해야 합니다.

 

학습 코드

실행 예시

python train_fasterrcnn_coco_official.py --coco_root ./coco --epochs 10 --batch 2 --num_workers 4

# train_fasterrcnn_coco_official.py
import os
import json
import argparse
from typing import Any, Dict, List, Tuple

import torch
from torch.utils.data import Dataset, DataLoader
from PIL import Image

import torchvision
from torchvision.transforms import functional as F
from torchvision.models.detection import fasterrcnn_resnet50_fpn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

# 실행 예시
# python train_fasterrcnn_coco_official.py --coco_root ./coco --epochs 10 --batch 2 --num_workers 4

# -------------------------
# Utils
# -------------------------
def collate_fn(batch):
    # detection task: variable number of boxes per image
    return tuple(zip(*batch))


def set_seed(seed: int):
    import random
    import numpy as np
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)


class CocoDetectionDataset(Dataset):
    """
    COCO 객체 검출(Instances) 어노테이션을 사용한 Dataset
      - instances_train2017.json
      - instances_val2017.json
    
    Reads:
      - images_dir: .../train2017 or .../val2017
      - ann_file  : .../annotations/instances_train2017.json (or val)

    Returns:
      image: FloatTensor (3,H,W) in [0,1]
      target: dict with boxes (xyxy), labels (contiguous), image_id, area, iscrowd
    """
    def __init__(self, images_dir: str, ann_file: str, cat_id_to_contiguous: Dict[int, int]):
        self.images_dir = images_dir
        self.ann_file = ann_file
        self.cat_map = cat_id_to_contiguous

        if not os.path.isdir(self.images_dir):
            raise FileNotFoundError(f"[ERR] images_dir not found: {self.images_dir}")
        if not os.path.isfile(self.ann_file):
            raise FileNotFoundError(f"[ERR] ann_file not found: {self.ann_file}")

        with open(self.ann_file, "r", encoding="utf-8") as f:
            coco = json.load(f)

        # image_id -> image_info
        self.id_to_img = {img["id"]: img for img in coco.get("images", [])}
        self.img_ids = sorted(self.id_to_img.keys())

        # group annotations by image_id
        self.anns_by_img: Dict[int, List[dict]] = {}
        for ann in coco.get("annotations", []):
            img_id = ann["image_id"]
            self.anns_by_img.setdefault(img_id, []).append(ann)

    def __len__(self) -> int:
        return len(self.img_ids)

    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, Dict[str, Any]]:
        img_id = self.img_ids[idx]
        info = self.id_to_img[img_id]
        file_name = info["file_name"]

        img_path = os.path.join(self.images_dir, file_name)
        img = Image.open(img_path).convert("RGB")
        image = F.to_tensor(img)  # float32, [0,1]

        anns = self.anns_by_img.get(img_id, [])

        boxes = []
        labels = []
        areas = []
        iscrowd = []

        for ann in anns:
            # COCO bbox: [x,y,w,h] -> xyxy
            x, y, w, h = ann["bbox"]
            if w <= 1 or h <= 1:
                continue

            x1, y1 = float(x), float(y)
            x2, y2 = float(x + w), float(y + h)

            cat_id = int(ann["category_id"])
            if cat_id not in self.cat_map:
                continue

            boxes.append([x1, y1, x2, y2])
            labels.append(self.cat_map[cat_id])
            areas.append(float(ann.get("area", w * h)))
            iscrowd.append(int(ann.get("iscrowd", 0)))

        if len(boxes) == 0:
            boxes_t = torch.zeros((0, 4), dtype=torch.float32)
            labels_t = torch.zeros((0,), dtype=torch.int64)
            areas_t = torch.zeros((0,), dtype=torch.float32)
            iscrowd_t = torch.zeros((0,), dtype=torch.int64)
        else:
            boxes_t = torch.tensor(boxes, dtype=torch.float32)
            labels_t = torch.tensor(labels, dtype=torch.int64)
            areas_t = torch.tensor(areas, dtype=torch.float32)
            iscrowd_t = torch.tensor(iscrowd, dtype=torch.int64)

        target = {
            "boxes": boxes_t,
            "labels": labels_t,
            "image_id": torch.tensor([img_id], dtype=torch.int64),
            "area": areas_t,
            "iscrowd": iscrowd_t,
        }
        return image, target


def build_category_mapping(instances_ann_file: str) -> Tuple[Dict[int, int], int]:
    """
    COCO category_id는 연속이 아니어서, torchvision 학습용으론 contiguous label(1..K)로 재매핑 권장.
    background는 0으로 예약.
    """
    with open(instances_ann_file, "r", encoding="utf-8") as f:
        coco = json.load(f)

    cat_ids = sorted([int(c["id"]) for c in coco.get("categories", [])])
    cat_map = {cid: i + 1 for i, cid in enumerate(cat_ids)}  # 1..K
    num_classes = len(cat_ids) + 1  # + background(0)
    return cat_map, num_classes


# -------------------------
# Model
# -------------------------
def build_model(num_classes: int):
    # torchvision 버전 호환
    try:
        model = fasterrcnn_resnet50_fpn(weights="DEFAULT")
    except TypeError:
        model = fasterrcnn_resnet50_fpn(pretrained=True)

    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
    return model


# -------------------------
# Train / "Val loss"
# -------------------------
def train_one_epoch(model, loader, optimizer, device, epoch: int, print_every: int = 50):
    model.train()
    loss_sum = 0.0

    for it, (images, targets) in enumerate(loader, start=1):
        images = [img.to(device) for img in images]
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

        loss_dict = model(images, targets)
        losses = sum(v for v in loss_dict.values())

        optimizer.zero_grad(set_to_none=True)
        losses.backward()
        optimizer.step()

        loss_sum += float(losses.item())

        if it % print_every == 0:
            avg = loss_sum / it
            detail = {k: float(v.item()) for k, v in loss_dict.items()}
            print(f"[Epoch {epoch}] iter={it}/{len(loader)} loss={avg:.4f} detail={detail}")

    return loss_sum / max(1, len(loader))


@torch.no_grad()
def compute_loss_on_loader(model, loader, device):
    """
    torchvision detection은 eval()에서 loss를 주지 않으므로,
    train() 모드에서 no_grad로 loss만 계산(드랍아웃 영향이 싫으면 backbone 등 고정 필요).
    튜토리얼/체크 용도로만 간단히 사용.
    """
    model.train()
    loss_sum = 0.0
    n = 0

    for images, targets in loader:
        images = [img.to(device) for img in images]
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]

        loss_dict = model(images, targets)
        losses = sum(v for v in loss_dict.values()).item()
        loss_sum += float(losses)
        n += 1

    return loss_sum / max(1, n)


def save_ckpt(model, optimizer, epoch: int, out_path: str):
    os.makedirs(os.path.dirname(out_path), exist_ok=True)
    torch.save(
        {"epoch": epoch, "model": model.state_dict(), "optimizer": optimizer.state_dict()},
        out_path
    )


# -------------------------
# Main
# -------------------------
def parse_args():
    p = argparse.ArgumentParser()
    p.add_argument("--coco_root", default="coco", help="folder that contains train2017/val2017/annotations")
    p.add_argument("--epochs", type=int, default=10)
    p.add_argument("--batch", type=int, default=2)
    p.add_argument("--lr", type=float, default=0.005)
    p.add_argument("--momentum", type=float, default=0.9)
    p.add_argument("--weight_decay", type=float, default=0.0005)
    p.add_argument("--num_workers", type=int, default=4)
    p.add_argument("--seed", type=int, default=42)
    p.add_argument("--device", default="cuda")
    p.add_argument("--out_dir", default="runs/fasterrcnn_coco_official")
    p.add_argument("--print_every", type=int, default=50)
    return p.parse_args()


def main():
    args = parse_args()
    set_seed(args.seed)

    device = torch.device(args.device if torch.cuda.is_available() and args.device.startswith("cuda") else "cpu")

    train_images = os.path.join(args.coco_root, "train2017")
    val_images = os.path.join(args.coco_root, "val2017")
    ann_dir = os.path.join(args.coco_root, "annotations")
    train_ann = os.path.join(ann_dir, "instances_train2017.json")
    val_ann = os.path.join(ann_dir, "instances_val2017.json")

    cat_map, num_classes = build_category_mapping(train_ann)
    print(f"[INFO] COCO root      : {args.coco_root}")
    print(f"[INFO] num_classes   : {num_classes} (incl background)")
    print(f"[INFO] train_images  : {train_images}")
    print(f"[INFO] train_ann     : {train_ann}")
    print(f"[INFO] val_images    : {val_images}")
    print(f"[INFO] val_ann       : {val_ann}")

    train_ds = CocoDetectionDataset(train_images, train_ann, cat_map)
    val_ds = CocoDetectionDataset(val_images, val_ann, cat_map)

    train_loader = DataLoader(
        train_ds,
        batch_size=args.batch,
        shuffle=True,
        num_workers=args.num_workers,
        pin_memory=True,
        collate_fn=collate_fn,
    )
    val_loader = DataLoader(
        val_ds,
        batch_size=args.batch,
        shuffle=False,
        num_workers=args.num_workers,
        pin_memory=True,
        collate_fn=collate_fn,
    )

    model = build_model(num_classes)
    model.to(device)

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

    print(f"[INFO] device={device} train={len(train_ds)} val={len(val_ds)}")

    for epoch in range(1, args.epochs + 1):
        train_loss = train_one_epoch(
            model, train_loader, optimizer, device,
            epoch=epoch, print_every=args.print_every
        )
        val_loss = compute_loss_on_loader(model, val_loader, device)

        lr_scheduler.step()

        lr_now = optimizer.param_groups[0]["lr"]
        print(f"[EPOCH {epoch}] train_loss={train_loss:.4f} val_loss={val_loss:.4f} lr={lr_now:.6f}")

        ckpt_path = os.path.join(args.out_dir, f"epoch_{epoch:03d}.pth")
        save_ckpt(model, optimizer, epoch, ckpt_path)

    print(f"[DONE] checkpoints saved to: {args.out_dir}")


if __name__ == "__main__":
    main()

  Faster R-CNN은 2-stage 구조이기 때문에 YOLO 계열 모델보다 학습 속도는 느리지만, 작은 객체나 복잡한 장면에서 더 안정적인 성능을 보입니다.

 

마무리

  이번 글에서는 COCO 데이터셋을 기반으로 Faster R-CNN 모델을 구성하고 PyTorch에서 실제 학습까지 연결하는 과정을 살펴보았습니다. 정확도가 중요한 객체 검출 문제에서는 Faster R-CNN이 여전히 강력한 베이스라인이 됩니다.

  다음 글에서는 학습된 Faster R-CNN 모델을 이용해 실제 이미지에 대한 추론 결과를 시각화하고, IOU 기준 성능을 확인하는 방법을 정리해보겠습니다.

 

관련 내용

• [실전 예제/객체 탐지/PyTorch] 객체 검출 튜토리얼: COCO 데이터셋으로 PyTorch 데이터셋 만들기
• [실전 예제/객체 탐지/PyTorch] 객체 검출 튜토리얼: DOTA 데이터셋으로 PyTorch 데이터셋 만들기
• [실전 예제/PyTorch/Detection] YOLO 모델을 이용한 이미지 객체 탐지