[실전 예제/이미지 분류/PyTorch] ResNet 기반 이미지 분류 모델 구성과 학습

이미지 분류(Image Classification)이란?

  이미지 분류(Image Classification)는 입력 이미지 전체를 하나의 클래스로 분류하는 가장 기본적인 컴퓨터 비전 태스크입니다. 객체의 위치나 개수보다는 이미지가 무엇을 나타내는지를 판단하는 데 초점을 둡니다.

  ResNet(Residual Network)은 깊은 신경망에서 발생하는 기울기 소실 문제를 Residual Connection으로 해결한 대표적인 CNN 모델입니다. 이번 글에서는 이전 글에서 구성한 CIFAR-10 / ImageNet 형식의 데이터셋을 기반으로, PyTorch에서 ResNet 모델을 구성하고 실제 학습까지 연결하는 과정을 정리합니다.

• 예: 동물·사물·풍경 이미지 분류
• 예: 제품 카테고리 자동 분류

 

목표

1. CIFAR-10 / ImageNet 구조의 데이터셋을 이용한 학습 파이프라인 구성
2. torchvision ResNet 모델 구조 이해
3. Cross Entropy Loss 기반 이미지 분류 학습 코드 작성

 

학습을 위한 데이터 구조

  이전 글에서 구성한 이미지 분류 데이터셋은 클래스별 디렉토리 구조를 사용하는 표준적인 형태입니다.

dataset/
  train/
    cat/
    dog/
    car/
  val/
    cat/
    dog/
    car/

  torchvision의 ImageFolder를 사용하면 별도의 라벨 파일 없이도 바로 데이터 로딩이 가능합니다.

 

ResNet 모델 구성

모델 개요

• Backbone: ResNet18 / ResNet50
• 입력: RGB 이미지
• 출력: 클래스 확률

import torch
import torch.nn as nn
from torchvision.models import resnet18

def build_model(num_classes):
    model = resnet18(weights="DEFAULT")
    in_features = model.fc.in_features
    model.fc = nn.Linear(in_features, num_classes)
    return model

  ImageNet으로 사전 학습된 가중치를 사용하고, 마지막 분류기(fc)만 데이터셋 클래스 수에 맞게 교체합니다.

 

학습 코드

import os
import torch
import torch.nn as nn
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision.models import resnet18

# -----------------------------
# 1) Model
# -----------------------------
def build_model(num_classes: int):
    model = resnet18(weights="DEFAULT")  # 구버전 torchvision이면 pretrained=True
    in_features = model.fc.in_features
    model.fc = nn.Linear(in_features, num_classes)
    return model

# -----------------------------
# 2) Transforms
# - CIFAR-10: 32x32이지만 ResNet18 입력을 맞추기 위해 224로 업샘플
# - ImageNet: 원래 224 기반
# -----------------------------
transform_train = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])

transform_eval = transforms.Compose([
    transforms.Resize(256),
    transforms.CenterCrop(224),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406],
                         std=[0.229, 0.224, 0.225]),
])

# -----------------------------
# 3) DataLoader (CIFAR-10 / ImageNet 스위치)
# -----------------------------
def get_dataloaders(dataset_name: str,
                    batch_size: int = 64,
                    num_workers: int = 4):
    dataset_name = dataset_name.lower()

    if dataset_name == "cifar10":
        train_ds = datasets.CIFAR10(root="./data", train=True, download=True, transform=transform_train)
        val_ds   = datasets.CIFAR10(root="./data", train=False, download=True, transform=transform_eval)
        num_classes = 10
        class_names = train_ds.classes

    elif dataset_name == "imagenet":
        # ✅ ImageNet은 torchvision에서 제공하지만 자동 다운로드는 안 됩니다.
        # root 경로 아래에 ImageNet 폴더가 준비되어 있어야 합니다.
        train_ds = datasets.ImageNet(root="./data", split="train", transform=transform_train)
        val_ds   = datasets.ImageNet(root="./data", split="val", transform=transform_eval)
        num_classes = 1000
        class_names = None

    else:
        raise ValueError("dataset_name은 'cifar10' 또는 'imagenet'만 가능합니다.")

    train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,
                              num_workers=num_workers, pin_memory=True)
    val_loader = DataLoader(val_ds, batch_size=batch_size*2, shuffle=False,
                            num_workers=num_workers, pin_memory=True)

    return train_loader, val_loader, num_classes, class_names

# -----------------------------
# 4) Train / Eval
# -----------------------------
def train_one_epoch(model, loader, optimizer, device):
    model.train()
    criterion = nn.CrossEntropyLoss()
    loss_sum = 0.0

    for images, labels in loader:
        images = images.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)

        optimizer.zero_grad(set_to_none=True)
        outputs = model(images)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()

        loss_sum += loss.item()

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

@torch.no_grad()
def evaluate(model, loader, device):
    model.eval()
    criterion = nn.CrossEntropyLoss()
    loss_sum = 0.0
    correct = 0
    total = 0

    for images, labels in loader:
        images = images.to(device, non_blocking=True)
        labels = labels.to(device, non_blocking=True)

        outputs = model(images)
        loss = criterion(outputs, labels)

        pred = outputs.argmax(dim=1)
        correct += (pred == labels).sum().item()
        total += labels.size(0)
        loss_sum += loss.item()

    avg_loss = loss_sum / max(1, len(loader))
    acc = correct / max(1, total)
    return avg_loss, acc

# -----------------------------
# 5) Main
# -----------------------------
def main():
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

    # ✅ 글에서는 여기만 바꾸면 된다고 설명하면 깔끔합니다.
    dataset_name = "cifar10"   # "imagenet" 로 변경 가능

    train_loader, val_loader, num_classes, class_names = get_dataloaders(
        dataset_name=dataset_name,
        batch_size=64,
        num_workers=4
    )

    model = build_model(num_classes=num_classes).to(device)
    optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

    for epoch in range(10):
        train_loss = train_one_epoch(model, train_loader, optimizer, device)
        val_loss, val_acc = evaluate(model, val_loader, device)
        print(f"[{dataset_name}] Epoch {epoch+1:02d} | train_loss={train_loss:.4f} | val_loss={val_loss:.4f} | val_acc={val_acc*100:.2f}%")

    if class_names is not None:
        print("classes:", class_names)

if __name__ == "__main__":
    main()

 

• CIFAR-10은 PyTorch가 자동 다운로드/로딩을 지원합니다.
• ImageNet은 torchvision에서 Dataset 클래스를 제공하지만, 데이터는 별도로 준비되어 있어야 합니다.
• 학습 파이프라인은 동일하고, 데이터 로더만 교체해도 재사용할 수 있도록 구성합니다.

 

마무리

  이번 글에서는 이미지 분류 데이터셋을 기반으로 ResNet 모델을 구성하고 PyTorch에서 실제 학습까지 연결하는 과정을 살펴보았습니다. 이미지 분류는 이후 객체 검출·분할 모델을 이해하는 데에도 중요한 기초가 됩니다.

  다음 글에서는 학습된 ResNet 모델을 이용해 검증 데이터에 대한 정확도를 측정하고, 예측 결과를 시각화하는 방법을 정리해보겠습니다.

 

관련 내용

• [실전 예제/이미지 분류/PyTorch] 이미지 분류 튜토리얼: CIFAR-10과 ImageNet으로 PyTorch 데이터셋 만들기
• [PyTorch] CNN 모델의 기초: torch.nn.Conv2d() 사용 가이드
• [PyTorch] 다중 클래스 분류에서 필수: torch.nn.CrossEntropyLoss() 사용 가이드