딥러닝 파이토치 강좌, 완전 합성곱 네트워크

완전 합성곱 네트워크(Fully Convolutional Networks, FCN)는 주로 이미지 세분화, 즉 이미지 내에서 특정 객체를 픽셀 단위로 분리하는 문제에 적합한 신경망 구조입니다. 기존의 CNN(합성곱 신경망)은 주로 분류 작업에 사용되며, 최종 출력으로 고정된 크기의 결과를 생성합니다. 하지만 FCN은 시각적인 정보를 유지하면서 변환된 출력을 생성할 수 있는 구조로, 이미지의 모든 픽셀에 의미를 부여할 수 있습니다.

1. 완전 합성곱 네트워크의 기본 구조

FCN은 기본적으로 CNN의 구조를 계승합니다. 하지만 중요한 점은 CNN의 마지막 부분에서 전결합층을 제거하고, 대신에 컨볼루션 층과 업샘플링 층을 사용하여 원하는 출력 크기를 갖도록 한다는 것입니다.

FCN의 주요 구성 요소는 다음과 같습니다:

  • 컨볼루션 층 (Convolutional Layer): 입력 이미지에서 특성을 추출하는 층.
  • 비선형 활성화 함수: 주로 ReLU(Rectified Linear Unit) 함수가 사용됨.
  • 업샘플링: 다운샘플링된 데이터를 원본 이미지의 크기로 복원.
  • 스킵 연결 (Skip Connection): 원본 해상도의 특성을 유지하면서 통합하기 위해 사용됨.

2. 파이토치를 이용한 FCN 구현

이제 FCN을 파이토치를 사용해 구현해보겠습니다. 아래는 간단한 FCN의 파이썬 코드 예제입니다.

import torch
import torch.nn as nn
import torch.nn.functional as F

class FCN(nn.Module):
    def __init__(self, num_classes):
        super(FCN, self).__init__()
        self.conv1 = nn.Conv2d(3, 64, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(128, 256, kernel_size=3, padding=1)
        self.conv4 = nn.Conv2d(256, 256, kernel_size=3, padding=1)
        
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        
        self.upconv1 = nn.ConvTranspose2d(256, 128, kernel_size=2, stride=2)
        self.upconv2 = nn.ConvTranspose2d(128, 64, kernel_size=2, stride=2)
        
        self.final_conv = nn.Conv2d(64, num_classes, kernel_size=1)

    def forward(self, x):
        x1 = F.relu(self.conv1(x))
        x2 = self.pool(x1)
        x2 = F.relu(self.conv2(x2))
        x3 = self.pool(x2)
        x3 = F.relu(self.conv3(x3))
        x4 = self.pool(x3)
        x4 = F.relu(self.conv4(x4))

        x = self.upconv1(x4)
        x = self.upconv2(x)
        x = self.final_conv(x)
        
        return x

2.1 모델 설명

위의 코드에서 우리 FCN 모델은 다음과 같은 단계를 거칩니다:

  1. 3채널(일반적인 RGB 이미지) 입력을 받아 64개의 특성 맵을 생성하는 첫 번째 컨볼루션 층을 통과합니다.
  2. 다음 두 개의 컨볼루션 층을 지나며 점차적으로 더 많은 특성 맵을 생성하고, 최대 풀링을 통해 이미지 크기를 절반으로 줄입니다.
  3. 업샘플링을 통해 이미지 크기를 원래 크기로 늘립니다.
  4. 최종 출력은 클래스 개수만큼의 채널을 가진 컨볼루션 층을 거칩니다.

3. 데이터셋 준비

FCN 모델을 학습시키기 위해선 적절한 데이터셋이 필요합니다. 일반적으로 이미지 세분화에 사용되는 데이터셋으로는 Pascal VOC, COCO 등이 있으며, 여기서는 간단한 이미지와 마스크의 예제를 사용합니다.

3.1 예제 데이터셋 생성

import numpy as np
import cv2
import matplotlib.pyplot as plt

def generate_example_data():
    h, w = 128, 128
    image = np.random.randint(0, 255, (h, w, 3), dtype=np.uint8)
    mask = np.zeros((h, w), dtype=np.uint8)
    mask[30:70, 30:70] = 1  # 사각형 객체
    
    return image, mask

image, mask = generate_example_data()

# 이미지와 마스크 시각화
plt.subplot(1, 2, 1)
plt.title('Image')
plt.imshow(image)
plt.subplot(1, 2, 2)
plt.title('Mask')
plt.imshow(mask, cmap='gray')
plt.show()

4. 모델 학습

데이터셋을 만들었다면, 이제 FCN 모델을 학습시킬 준비가 되었습니다. 학습 과정에서 PyTorch의 손실 함수와 옵티마이저를 설정해야 합니다.

import torch.optim as optim

# 모델, 손실 함수 및 옵티마이저 초기화
num_classes = 2  # 객체와 배경
model = FCN(num_classes)
criterion = nn.CrossEntropyLoss()  # 크로스 엔트로피 손실 함수
optimizer = optim.Adam(model.parameters(), lr=0.001)

# Dummy training loop
for epoch in range(5):  # 5번의 epoch 동안 학습
    model.train()
    optimizer.zero_grad()
    
    # Forward pass
    inputs = torch.Tensor(image).permute(2, 0, 1).unsqueeze(0)  # (1, 3, 128, 128)
    targets = torch.Tensor(mask).long().unsqueeze(0)  # (1, 128, 128)
    
    outputs = model(inputs)
    loss = criterion(outputs, targets)
    
    # Backward pass
    loss.backward()
    optimizer.step()
    
    print(f'Epoch [{epoch+1}/5], Loss: {loss.item():.4f}')

5. 결론

이 글에서는 완전 합성곱 네트워크(FCN)의 기본 개념부터, 파이토치를 이용한 간단한 FCN 모델 구현 과정, 데이터셋 준비, 학습 방법에 대해 살펴보았습니다. FCN은 이미지 세분화에 매우 유용한 모델로, 다양한 응용 분야에서 사용될 수 있습니다.

더욱 발전된 FCN 모델에 대한 연구와 추가적인 데이터셋을 통해 더 나은 성능을 목표로 할 수 있습니다. FCN의 응용이 궁금하시다면, 더 많은 내용을 탐색해보시길 추천드립니다!