[실전 예제/인스턴스 분할/PyTorch] Mask R-CNN 모델 구성과 COCO 학습

인스턴스 분할(Instance Segmentation)이란?

  인스턴스 분할(Instance Segmentation)은 객체 검출(Object Detection)과 의미 분할(Semantic Segmentation)을 결합한 컴퓨터 비전 태스크입니다. 이미지 안의 각 객체를 구분하면서, 객체마다 픽셀 단위의 마스크를 예측하는 것이 핵심입니다.

  Mask R-CNN은 Faster R-CNN 구조를 기반으로, ROI 단위에서 Bounding Box + Class + Mask를 동시에 예측하는 대표적인 인스턴스 분할 모델입니다. 이번 글에서는 이전 글에서 구성한 COCO 인스턴스 분할 데이터셋을 바탕으로, PyTorch에서 Mask R-CNN 모델을 구성하고 실제 학습까지 연결하는 과정을 정리합니다.

• 예: 사람·차량·사물 영역을 개별 객체 단위로 분할
• 예: 객체 개수 계산, 정밀 영역 분석

 

목표

1. COCO 인스턴스 분할 데이터셋을 입력으로 하는 학습 파이프라인 구성
2. torchvision Mask R-CNN 모델 구조 이해
3. Box / Mask loss를 포함한 학습 코드 작성

 

학습을 위한 데이터 구조

  이전 글에서 구성한 COCO 인스턴스 분할 데이터셋은 COCO annotation 포맷(JSON)을 그대로 사용하며, 이미지와 마스크 정보를 함께 포함합니다.

coco/
  train/
    images/
      xxx.jpg
    annotations.json
  val/
    images/
      xxx.jpg
    annotations.json

  Mask R-CNN 학습에서는 각 샘플마다 다음 정보가 필요합니다.

• boxes: (N,4) 형태의 bounding box
• labels: 각 객체의 class id
• masks: (N,H,W) 형태의 binary mask

 

Mask R-CNN 모델 구성

모델 개요

• Backbone: ResNet50 + FPN
• Detector: Faster R-CNN 기반
• Segmentation Head: ROIAlign + FCN Mask Head

import torch
from torchvision.models.detection import maskrcnn_resnet50_fpn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.mask_rcnn import MaskRCNNPredictor

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

    # Box predictor 교체
    in_features = model.roi_heads.box_predictor.cls_score.in_features
    model.roi_heads.box_predictor = FastRCNNPredictor(
        in_features,
        num_classes
    )

    # Mask predictor 교체
    in_features_mask = model.roi_heads.mask_predictor.conv5_mask.in_channels
    hidden_layer = 256
    model.roi_heads.mask_predictor = MaskRCNNPredictor(
        in_features_mask,
        hidden_layer,
        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}")

  Mask R-CNN은 box loss 외에도 mask loss가 추가되기 때문에, 객체 검출보다 메모리 사용량이 큽니다. batch size를 무리하게 키우지 않는 것이 중요합니다.

 

마무리

  이번 글에서는 COCO 데이터셋을 기반으로 Mask R-CNN 모델을 구성하고 PyTorch에서 실제 학습까지 연결하는 과정을 살펴보았습니다. 인스턴스 분할은 객체 검출보다 정보량이 훨씬 풍부한 태스크이지만, 그만큼 데이터 구조와 메모리 관리가 중요합니다.

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

 

관련 내용

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