[Python] 간단한 Pytorch 코드 예제(MNIST)

Pytorch를 처음 접했을 때 tensorflow, keras와는 코드 생김새(?)가 달라서 접근하기 어려웠다. 하지만 계속 쓰다 보니 유사한 코드 작성 패턴이 있어서 기록해 두려고 한다. 아래는 유명한 MNIST 데이터 셋을 이용한 기본적인 Pytorch 예제이고 최소한의 코드만 작성했다.

 

1. 필요한 모듈 로드

import pandas as pd
import numpy as np
import random
import os

from sklearn.model_selection import train_test_split

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader
from torchvision import transforms
import torchvision.models as models

 

2. 데이터 읽어오기

# cvs 파일 읽어오기
test = pd.read_csv('mnist_test.csv').iloc[:, 1:]
train = pd.read_csv('mnist_train.csv').iloc[:, 1:]
sub_df = pd.read_csv('submission.csv')

# 훈련용, 검증용으로 데이터 나누기
train_images, val_images, train_labels, val_labels = train_test_split(train.iloc[:, :-1], train.iloc[:, -1], test_size=0.05)

# 훈련/검증/테스트
train_images.reset_index(drop=True, inplace=True)
val_images.reset_index(drop=True, inplace=True)
train_labels.reset_index(drop=True, inplace=True)
val_labels.reset_index(drop=True, inplace=True)
test_images = test

 

3. 데이터셋과 데이터로더 작성하기

Pytorch를 처음 접했을 때 데이터 셋과 데이터 로더의 개념이 어려웠고 왜 이걸 사용하는지 이해가 되지 않았다. 그러나 Pytorch 코드를 계속해서 접해보니 주변에서 많이 쓰는 이유가 있었다. 정말 '편리'하기 때문이다. 데이터 셋을 사용하는 이유는 csv와 같은 파일로 읽어온 데이터를 feature와 label로 나누어 정리해 둘 수 있다. 그리고 데이터 로더를 사용해 자신이 정한 배치 사이즈만큼 가져와 데이터를 확인할 수 있고 모델을 훈련시킬 수 있다. 비유를 하자면 데이터 셋은 케이크, 데이터 로더는 그 케이크를 조각낸 조각 케이크로 생각할 수 있다. 케이크를 한 번에 먹기보다는 조각내서 먹는 게 더 깔끔하고 먹기 편하다. 내가 느낀 데이터 셋과 데이터 로더도 그러하다.

 

• transforms 정의
  transforms 모듈은 Pytorch에서 이미지 데이터를 다룰 때 많이 쓰인다. 아래의 function 말고도 다양한 function이 있으니 직접 확인해서 사용하면 좋을 것 같다.
  

  train_tf = transforms.Compose([
      transforms.ToPILImage(),
      transforms.ToTensor(),
      transforms.Normalize((0.1307,), (0.3081,))
  ])
  
  val_tf = transforms.Compose([
      transforms.ToPILImage(),
      transforms.ToTensor(),    
      transforms.Normalize((0.1307,), (0.3081,))
  ])

 

• dataset 정의
  커스텀 데이터셋은 'torch.utils.data.Dataset' 을 상속 받아서 만들 수 있다.
  - __init__: feature, label, transforms를 정의
  - __len__: 데이터셋의 길이를 반환
  -__getitem__: 해당 index의 데이터를 가져올 수 있음
  

  class MnistDataSet(torch.utils.data.Dataset):
      
      def __init__(self, images, labels, transforms=None):
          self.X = images
          self.y = labels
          self.transforms = transforms
          
      def __len__(self):
          return len(self.X)
      
      def __getitem__(self, i):
          data = self.X.iloc[i, :]
          data = np.array(data).astype(np.uint8).reshape(28, 28, 1)
          
          if self.transforms:
              data = self.transforms(data)
              
          if self.y is not None:
              return (data, self.y[i])
          else:
              return data
             
  # train/val dataset 정의
  train_dataset = MnistDataSet(train_images, train_labels, train_tf)
  val_dataset = MnistDataSet(val_images, val_labels, val_tf)

• dataloader 정의
  데이터 로더는 데이터 셋에 비해 그나마 간단한 편이다.
  

  train_loader = DataLoader(train_dataset, batch_size=50)
  valid_loader = DataLoader(val_dataset, batch_size=50)

4. 모델 정의하기

훈련하고자 하는 모델을 정의한다.

class MnistResNet(nn.Module):
    def __init__(self, in_channels=1):
        super(MnistResNet, self).__init__()

        # Load a pretrained resnet model from torchvision.models in Pytorch
        self.model = models.resnet152(pretrained=True)

        # Change the input layer to take Grayscale image, instead of RGB images. 
        # Hence in_channels is set as 1 or 3 respectively
        # original definition of the first layer on the ResNet class
        # self.conv1 = nn.Conv2d(3, 64, kernel_size=7, stride=2, padding=3, bias=False)
        self.model.conv1 = nn.Conv2d(in_channels, 64, kernel_size=7, stride=2, padding=3, bias=False)

        # Change the output layer to output 10 classes instead of 1000 classes
        num_ftrs = self.model.fc.in_features
        self.model.fc = nn.Linear(num_ftrs, 10)

    def forward(self, x):
        return self.model(x)

 

5. loss & optimizer 정의

# 모델 생성
my_resnet = MnistResNet().to(device)

# cuda가 있으면 gpu 사용하고 없으면 cpu 사용
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# loss 정의
loss_fn = nn.CrossEntropyLoss()
# optimizer 정의
optimizer = optim.Adam(my_resnet.parameters(), lr=3e-4)

 

6. 훈련 & 검증

훈련과 검증을 반복한다. feature(여기선 images)와 label 데이터를 device에 넣어준다. loss 계산하고 optimizer 업데이트도 진행한다. 

for epoch in range(epochs):
    print('Epoch: ', epoch+1)
    
    epoch_loss = 0
    epoch_correct = 0
    
    my_resnet.train() # train mode
    
    for images, labels in train_loader:
        x, y = images.to(device), labels.long().to(device) # 앞서 정의한 device에 넣어줌
        
        pred = my_resnet(x)
        
        optimizer.zero_grad()
        loss = loss_fn(pred, y)
        loss.backward()
        optimizer.step()
        
        epoch_loss += loss.item() * batch_size
        epoch_correct += pred.argmax(dim=1).eq(y).sum().item()
        
    val_loss = 0
    val_correct = 0
    
    my_resnet.eval() # eval mode
    
    with torch.no_grad():
        for images, labels in valid_loader:
            x, y = images.to(device), labels.long().to(device)
            
            pred = my_resnet(x)
            
            loss = loss_fn(pred, y)
            
            val_loss += loss.item() * batch_size
            val_correct += pred.argmax(dim=1).eq(y).sum().item()
            
    print(" Val Loss: ", val_loss)
    print(" Val Acc: ", (val_correct/len(val_images))*100)

 

7. 테스트

# test 데이터에 대한 데이터 셋과 데이터 로더 정의
test_dataset = MnistDataSet(test_images, None, val_tf)
test_dataloader = DataLoader(test_dataset, batch_size=50, shuffle=False)

# 테스트 진행
my_resnet.eval() # 평가 모드

predictions = torch.LongTensor().to(device)

for images in test_dataloader:
    pred = my_resnet(images.to(device))
    predictions = torch.cat((predictions, pred.argmax(dim=1)), dim=0)