[Python] PyTorch로 합성곱신경망(CNN) 모델 구축

지난 포스팅에서는 PyTorch를 활용한 심층신경망(DNN) 모델 학습에 대한 코드를 알아보았다.

이번 포스팅에서는 이미지 분석을 위한 합성곱신경망(CNN) 모델을 생성하고 학습하는 코드를 알아보도록 하겠다.

 

 

CNN 모델 구축과 학습 과정은 아래 순서대로 진행한다.

1. 필요한 라이브러리/패키지 불러오기
2. 연산을 수행할 장치(CPU or GPU) 설정
3. 데이터 불러오기
4. CNN 모델 생성 및 파라미터 정의
5. 모델 학습 및 평가

 

1. 라이브러리/패키지 불러오기

import numpy as np # 넘파이 배열
import matplotlib.pyplot as plt # 그래프 그리기

import torch
import torch.nn as nn  # 딥러닝 모듈
from torch.autograd import Variable # 자동 미분
import torch.nn.functional as F # 활성화 함수

import torchvision # FashionMNIST 데이터셋 사용
import torchvision.transforms as transforms # 데이터셋 변형
from torch.utils.data import Dataset, DataLoader # 데이터셋 로딩, 전처리, 순회

 

 

2. 연산 수행 장치 설정 (CPU/GPU)

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# CUDA 사용가능하면 GPU, 아니면 CPU

 

 

3. 데이터 불러오기

3.1 데이터 불러오기

# torchvision의 FashionMNIST 데이터셋 텐서 형태로 불러오기
train_dataset = datasets.FashionMNIST(
    root='data', train=True, download=True, transform=transforms.ToTensor()
)
test_dataset = datasets.FashionMNIST(
    root='data', train=False, download=True, transform=transforms.ToTensor()
)

 

3.2 데이터를 DataLoader에 전달하여 원하는 크기의 batch 단위로 학습

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=100)
test_loader = torch.utils.data.DataLoader(test_dataset, batch_size=100)

 

 

4. CNN 모델 생성 및 파라미터 정의

4.1 CNN 모델 생성

class FashionCNN(nn.Module):
  def __init__(self):
    super(FashionCNN, self).__init__()

    # 합성곱층 : 특징 추출, 차원 축소
    self.layer1 = nn.Sequential( # nn.sequential로 layer를 차례로 쌓아줌
        # convolution layer
        nn.Conv2d(in_channels=1, out_channels=32, kernel_size=3, padding=1),
        # 각 배치별로 다양한 분포의 데이터를 정규화                         
        nn.BatchNorm2d(32),
        nn.ReLU(),
        # Pooling layer 차원 축소                         
        nn.MaxPool2d(kernel_size=2, stride=2)
    )
    self.layer2 = nn.Sequential(
        nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3),        
        nn.BatchNorm2d(64),
        nn.ReLU(),             
        nn.MaxPool2d(2)
    )
    
    # 완전연결층(FCL) : 분류 
    self.fc1 = nn.Linear(in_features=64*6*6, out_features=600)
    self.drop = nn.Dropout2d(0.25)
    self.fc2 = nn.Linear(in_features=600, out_features=120)
    self.fc3 = nn.Linear(in_features=120, out_features=10)

  def forward(self, x):
      out = self.layer1(x)
      out = self.layer2(out)
      # 합성곱층에서 완전연결층으로 변경, 데이터를 1차원으로 변경
      out = out.view(out.size(0), -1)
      out = self.fc1(out)
      out = self.drop(out)
      out = self.fc2(out)
      out = self.fc3(out)       
      return out

 

4.2 모델 파라미터 정의

# 파라미터 정의
learning_rate = 0.001;
model = FashionCNN();
model.to(device);

criterion = nn.CrossEntropyLoss();
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate);
print(model)

 

5. 모델 학습 및 평가 

# 모델 학습
num_epochs = 5
count = 0
loss_list = []
iteration_list = []
accuracy_list = []

predictions_list = []
labels_list = []

for epoch in range(num_epochs):
    for images, labels in train_loader:
        images, labels = images.to(device), labels.to(device)
    
        train = Variable(images.view(100, 1, 28, 28))
        labels = Variable(labels)
        
        outputs = model(train)
        loss = criterion(outputs, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        count += 1
    
        if not (count % 50):    
            total = 0
            correct = 0        
            for images, labels in test_loader:
                images, labels = images.to(device), labels.to(device)
                labels_list.append(labels)            
                test = Variable(images.view(100, 1, 28, 28))            
                outputs = model(test)            
                predictions = torch.max(outputs, 1)[1].to(device)
                predictions_list.append(predictions)
                correct += (predictions == labels).sum()            
                total += len(labels)
            
            accuracy = correct * 100 / total
            loss_list.append(loss.data)
            iteration_list.append(count)
            accuracy_list.append(accuracy)
        
        if not (count % 500):
            print("Iteration: {}, Loss: {}, Accuracy: {}%".format(count, loss.data, accuracy))