딥러닝 파이토치 강좌, PSPNet

본 강좌에서는 딥러닝을 이용한 이미지 분할의 최신 기법 중 하나인 PSPNet(Pyramid Scene Parsing Network)에 대해 알아보겠습니다. PSPNet은 이미지의 의미론적 분할을 수행하는 데 있어 특히 뛰어난 성능을 보여주며, 다양한 이미지 인식 문제에 응용될 수 있습니다.

1. PSPNet 개요

PSPNet은 2017년 Zhang 등에서 발표한 네트워크로, 이미지의 전역 컨텍스트를 포착하여 각 픽셀에 대한 클래스 확률을 예측합니다. 이 모델은 피라미드 풀링 모듈(PPM, Pyramid Pooling Module)을 통해 다양한 스케일에서 정보를 통합하는 방식으로 작동합니다. 이러한 구조는 객체 인식에 유리하여 다양한 객체 크기에 대한 인식을 가능하게 합니다.

1.1. 주요 특징

• 피라미드 풀링 모듈: 이미지의 여러 크기에서 피처를 추출하고 통합하여 보다 포괄적인 맥락 정보를 제공합니다.
• 전역 정보 통합: 네트워크의 마지막 단계에서 전역 정보를 통합하여 최종 예측을 강화합니다.
• 우수한 성능: 여러 벤치마크 데이터셋에서 뛰어난 성능을 나타내며, 다양한 응용 분야에서 활용될 수 있습니다.

2. PSPNet 구조

PSPNet의 기본 구조는 다음과 같이 나눌 수 있습니다:

1. 백본 네트워크: ResNet 등의 CNN 모델을 기반으로 사용됩니다.
2. 피라미드 풀링 모듈: 다중 스케일의 피처 맵을 통합하여 전반적인 맥락을 포착합니다.
3. 업샘플링: 최종 예측을 위해 적절한 해상도로 조정합니다.

2.1. 피라미드 풀링 모듈

PPM은 입력 이미지에 대해 여러 해상도에서 피처 맵을 생성합니다. 이 모듈은 각기 다른 크기의 풀링 연산을 수행하여 공간 정보를 수집하고, 이를 다시 원본 해상도로 통합합니다. PPM은 다음과 같은 단계로 구성됩니다:

• 입력 피처 맵에 대해 다양한 사이즈의 풀링 연산을 수행합니다 (예: 1×1, 2×2, 3×3, 6×6).
• 각 풀링 단계에서 출력된 피처 맵을 다시 원본 해상도로 업샘플링합니다.
• 최종적으로 모든 업샘플링 된 피처 맵을 concatenate하여 새로운 피처 맵을 생성합니다.

3. PyTorch로 PSPNet 구현하기

이제 PSPNet을 PyTorch로 구현해 보겠습니다. 아래의 코드는 PSPNet의 구조를 정의합니다.

3.1. 환경 설정

import torch
import torch.nn as nn
import torchvision.models as models
    

3.2. PSPNet 클래스 정의

PSPNet 클래스는 백본 네트워크와 피라미드 풀링 모듈을 통합합니다. 다음과 같은 방식으로 정의할 수 있습니다:

class PSPModule(nn.Module):
    def __init__(self, in_channels, out_channels):
        super(PSPModule, self).__init__()
        self.pool1 = nn.AvgPool2d(1, stride=1)
        self.pool2 = nn.AvgPool2d(2, stride=2)
        self.pool3 = nn.AvgPool2d(3, stride=3)
        self.pool4 = nn.AvgPool2d(6, stride=6)
        self.conv1x1 = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
        self.bn = nn.BatchNorm2d(out_channels)

    def forward(self, x):
        size = x.size()[2:]
        p1 = self.conv1x1(self.pool1(x))
        p2 = self.conv1x1(self.pool2(x))
        p3 = self.conv1x1(self.pool3(x))
        p4 = self.conv1x1(self.pool4(x))

        p1 = nn.functional.interpolate(p1, size, mode='bilinear', align_corners=True)
        p2 = nn.functional.interpolate(p2, size, mode='bilinear', align_corners=True)
        p3 = nn.functional.interpolate(p3, size, mode='bilinear', align_corners=True)
        p4 = nn.functional.interpolate(p4, size, mode='bilinear', align_corners=True)

        return torch.cat((x, p1, p2, p3, p4), dim=1)

class PSPNet(nn.Module):
    def __init__(self, num_classes):
        super(PSPNet, self).__init__()
        self.backbone = models.resnet101(pretrained=True)
        self.ppm = PSPModule(2048, 512)
        self.final_convolution = nn.Conv2d(2048 + 512 * 4, num_classes, kernel_size=1)

    def forward(self, x):
        x = self.backbone(x)
        x = self.ppm(x)
        x = self.final_convolution(x)
        return x

3.3. 모델 학습하기

모델을 학습하기 위해 데이터셋 준비, 옵티마이저 설정, 그리고 학습 루프를 작성해야 합니다. torchvision의 Cityscapes 데이터셋을 예로 들어봅시다.

from torchvision import datasets, transforms
from torch.utils.data import DataLoader

# 데이터셋 준비
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
])

train_dataset = datasets.Cityscapes(root='path/to/cityscapes/', split='train', mode='fine', target_type='semantic', transform=transform)
train_loader = DataLoader(train_dataset, batch_size=4, shuffle=True)

# 모델과 옵티마이저 설정
model = PSPNet(num_classes=19).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)
criterion = nn.CrossEntropyLoss()

# 학습 루프
for epoch in range(num_epochs):
    model.train()
    for images, masks in train_loader:
        images, masks = images.to(device), masks.to(device)

        optimizer.zero_grad()
        outputs = model(images)
        loss = criterion(outputs, masks)
        loss.backward()
        optimizer.step()

    print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')

4. 실험 및 평가

학습이 완료된 후, 모델을 평가하기 위해 validation 데이터셋에 대한 성능을 측정합니다. 평가 지표로는 주로 IoU(Intersection over Union) 또는 Pixel Accuracy가 사용됩니다. 다음의 코드는 모델의 성능을 평가하는 방법을 보여줍니다.

def evaluate(model, val_loader):
    model.eval()
    total_loss = 0
    total_correct = 0
    total_pixels = 0

    with torch.no_grad():
        for images, masks in val_loader:
            images, masks = images.to(device), masks.to(device)
            outputs = model(images)
            loss = criterion(outputs, masks)

            total_loss += loss.item()
            preds = outputs.argmax(dim=1)
            total_correct += (preds == masks).sum().item()
            total_pixels += masks.numel()

    print(f'Validation Loss: {total_loss / len(val_loader):.4f}, Pixel Accuracy: {total_correct / total_pixels:.4f}')

# 평가 수행하기
evaluate(model, val_loader)

5. 결론

이번 강좌에서는 PSPNet의 구조와 동작 원리를 살펴보았습니다. PyTorch로 모델을 구현하고 학습하는 과정을 통해 의미론적 분할 문제를 해결하는 방법을 이해하셨길 바랍니다. PSPNet은 뛰어난 성능을 보여주는 네트워크로, 실제 이미지 처리 문제와 다양한 응용 분야에서 활용될 수 있습니다.