[실전 예제/객체 탐지/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/
  train/
    images/
      xxx.jpg
    annotations.json
  val/
    images/
      xxx.jpg
    annotations.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 포함으로 설정해야 합니다.

 

학습 코드

from torch.utils.data import DataLoader

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

model = build_model(num_classes=81)  # COCO 기준
model.to(device)

optimizer = torch.optim.SGD(
    model.parameters(),
    lr=0.005,
    momentum=0.9,
    weight_decay=0.0005
)

for epoch in range(10):
    model.train()
    loss_sum = 0.0

    for images, targets in train_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(loss for loss in loss_dict.values())

        optimizer.zero_grad()
        losses.backward()
        optimizer.step()

        loss_sum += losses.item()

    print(f"[Epoch {epoch+1}] loss={loss_sum/len(train_loader):.4f}")

  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 모델을 이용한 이미지 객체 탐지