파이토치(PyTorch) 이해하기 with MNIST

지난 글에서 텐서플로우와 파이토치의 차이를 살펴보고 파이토치의 특징과 장단점을 간단하게 알아봤다.

이번 포스팅에서는 파이토치의 MNIST 예제를 통해 파이토치의 작동 방식을 보다 자세히 공부하려고 한다.

코드는 지난 글에서 그대로 가져왔다.

 

라이브러리 및 config

import torch
import torchvision
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader

import numpy as np
import matplotlib.pyplot as plt

 

- torchvision은 이미지/영상 처리를 위한 파이토치의 라이브러리다.

- torch.nn 이 무엇인지에 대해 파이토치에서 도큐멘트를 제공하고 있다. (torch.nn 없이 모델 쌓는 법부터 안내해줌)

 

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

torch.device를 통해 디바이스 정보를 담은 객체를 생성한다. 가용한 GPU가 있을 경우 'cuda'로 디바이스 정보 저장.

 

num_classes = 10
input_size = 784
batch_size = 64
lr = 0.0001
epochs = 3

 

하이퍼파라미터 설정

 

데이터셋 셋팅

T = torchvision.transform.Compose([torchvision.transforms.ToTensor()])

X_train = torchvision.datasets.MNIST(root='/datasets', \
                                     train=True, \
                                     download=True, \
                                     transform=T)
train_loader = DataLoader(dataset=X_train, batch_size=batch_size, shuffle=True)

X_test = torchvision.datasets.MNIST(root='/datasets', \
                                     train=False, \
                                     download=True, \
                                     transform=T)
test_loader = DataLoader(dataset=X_test, batch_size=batch_size, shuffle=True)

 

개인적으로 파이토치에 대해서 가장 진입장벽 느낀 게 데이터를 파이토치의 Dataset 객체로 불러오는 부분이었다. 차근차근 하나씩 살펴보겠다.

 

위 코드에서는 torchvision에서 제공되는 기본 데이터셋인 MNIST를 불러와 tensor 객체로 변환하고 미니배치로 쪼개고 있다.

 

먼저 torchvision.transform.Compose

torchvision은 데이터를 변환할 수 있는 기능을 제공하고 있는데, 만약 여러 단계에 걸쳐 변환을 적용한다면 Compose 를 통해 이 단계들을 묶어줄 수 있다. 이때 transform 기능들은 리스트에 담아 넘겨주어야 한다.

 

예를 들어,

import torchvision.transforms as transforms

mnist_transform = transforms.Compose([
    transforms.Resize((300, 300))
    transforms.CenterCrop(10),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2),
    transforms.RandomHorizontalFlip(p = 1)
    transforms.PILToTensor(),
    transforms.ConvertImageDtype(torch.float),
])

리사이징, 크롭, 플립 등 다양한 처리 기능들을 Compose로 묶어주고 있다.

 

이 Compose 클래스의 객체를 torchvision.datasets.MNIST의 transform 인자에 넘겨주면, 정의된 이미지 처리 기능들이 데이터에 적용된다.

torchvision.datasets.MNIST의 객체는 튜플 형태로, 즉 (image, target) 으로 MNIST 데이터를 저장하고 있다.

소스코드: https://pytorch.org/vision/stable/_modules/torchvision/datasets/mnist.html#MNIST

 

그 다음 DataLoader

DataLoader(dataset, batch_size=1, shuffle=False, sampler=None,
           batch_sampler=None, num_workers=0, collate_fn=None,
           pin_memory=False, drop_last=False, timeout=0,
           worker_init_fn=None, *, prefetch_factor=2,
           persistent_workers=False)

파이토치의 공식 도큐멘트를 기반으로 살펴보자면, DataLoader에서 가장 중요한 인자는 'dataset' (Dataset 객체)이다.

이때 파이토치는 두 가지 타입의 dataset을 지원한다.

 

1) map-stype datsets

__getitem__() 과 __len__() 을 구현하는 데이터셋이다. 인덱스 및 키로 접근할 수 있다.

예를 들어, dataset[idx]로 idx번째 이미지와 그에 해당하는 라벨을 확인할 수 있다.

 

2) iterable-style datsets

__iter()__를 구현하는 데이터셋으로 IterableDataset 하위 클래스의 객체다.

데이터를 랜덤으로 읽어오는 것이 비싸거나 어려울 때, 그리고 메모리에 학습 데이터를 모두 올리기 어려울 때 사용한다. (ex. 실시간성 데이터)

iter(dataset)을 통해 데이터 스트림을 읽어올 수 있다.

 

이번 예제에서는 MNIST 데이터셋 객체를 가져와 바로 DataLoader로 넘겨줄 수 있었다.

 

커스텀 데이터셋

하지만 나만의 커스텀 이미지 데이터셋을 활용하려면?

Dataset 클래스를 상속받아서, 나만의 Dataset 객체를 생성해주는 클래스를 구현해야 한다.

 

[1] 사용할 이미지 파일들의 경로를 리스트에 저장해준다.

filelist = glob.glob('/path/to/datset/*.jpg')

 

[2] 이미지 전처리를 묶어준다

my_transform = transforms.Compose([
    transforms.Resize((300, 300))
    transforms.CenterCrop(10),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2),
    transforms.RandomHorizontalFlip(p = 1)
    transforms.PILToTensor(),
    transforms.ConvertImageDtype(torch.float),
])

 

 

[3] Dataset 클래스를 구현한다

class MyDataset(torch.utils.data.Dataset):
     def __init__(self, filelist, transform):
          self.filelist = filelist
          self.transform = transform
          
     def __len__(self):
          return len(self.filelist)
          
     def __getitem__(self, index):
          img_path = self.filelist[index]
          img = Image.open(img_path)
          img_transformed = self.transform(img)
          return img_transformed

 

[4] DataLoader로 불러온다

train_dataset = MyDataset(filelist = filelist, 
                          transform = my_transform)
                          
train_dataloader = DataLoader(train_dataset,
                              batch_size = 64,
                              shuffle = True)
               
# 확인해봅시다
batch_iterator = iter(train_dataloader)
images = next(batch_iterator)
print(images.size())

 

만약에 학습할 데이터를 계층적인 폴더에 저장해두었다면, pytorch의 ImageFolder 라이브러리를 통해 편리하게 Dataset 객체로 불러올 수 있다.

예를 들어 아래와 같이 각 이미지가 클래스명으로 된 폴더 아래 들어가 있다고 해보자.

 

 

mydata/

          0/

                    0.jpg

                    1.jpg

                    ...

          1/

                    0.jpg

                    1.jpg

                    ...

...

          9/

                    0.jpg

                    1.jpg

                    ...

 

이 경우 아래와 같이 데이터셋을 준비할 수 있다.

from torchvision.datasets import ImageFolder

train_dataset = ImageFolder(root='./mydata', transform=my_transform)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=128, shuffle=True, num_workers=4)

 

이는 이곳에서 다룬 텐서플로우tensorflow의 ImageDataGenerator와 매우 유사해보인다!

 

 

모델 정의 및 학습

class MyNetwork(nn.Module):
     def __init__(self, input_size, num_classes):
          super(MyNetwork, self).__init__()
          self.fc1 = nn.Linear(in_features=input_size, out_features=50)
          self.fc2 = nn.Linear(in_features=50, out_features=num_classes)
          
     def forward(self, x):
          x = self.fc1(x)
          x = F.relu(x)
          x = self.fc2(x)
          return x
          
model = MyNetwork()

기본적으로 torch.nn.Module 을 상속해온다.

만약 위 모델구조를 텐서플로우 케라스로 구현한다면 아래와 같지 않을까 싶다.

model = tf.keras.models.Sequential()
model.add(tf.keras.layers.Dense(out_features, input_shape=(input_size,), activation='relu'))
model.add(tf.keras.layers.Dense(num_classes))

 

그 다음

# 손실함수와 옵티마이저
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=lr)

# 훈련
model.train()
for epoch in range(epochs):
     for batch, (data, target) in enumerate(train_loader):
          data = data.to(device=device)
          target = target.to(device=device)
          
          data = data.reshape(data.shape[0], -1)
          
          # 순전파
          score = model(data)
          loss = criterion(score, target)
          
          # 역전파
          optimizer.zero_grad()
          loss.backward()
          optimizer.step()

생소한 것이 많은데 하나씩 짚어보겠다.

-  loss function을 criterion이라는 객체명으로 불러온다. 

- 옵티마이저를 불러온다. 이때, 업데이트가 필요한 모델의 파라미터를 넘겨준다 (model.parameters())

 

- model.train() : 모델을 훈련모드로 진입시킨다.

- train_loader 를 enumerate 했으므로 batch 에는 인덱스 번호가 불러와진다

- 앞서 언급했다시피 (image, label) 튜플로 저장된다고 했으니 튜플 형식으로 데이터를 받아온다

- 데이터를 지정된 디바이스(CPU/GPU)에 할당해준다

 

- 네트워크 객체인 model 에 데이터를 입력하면 feed forward 되어 출력한다

- 출력된 값과 타겟값을 손실함수에 넣어 loss 를 계산한다

(nn.CrossEntropyLoss()의 소스코드를 보면 입력된 값과 타겟값의 크로스엔트로피를 구하게 됨)

 

- optimizer.zero_grad() : 한번 iteration이 끝날 때마다 gradient를 0으로 초기화해준다

- loss.backward() : loss 는 크로스엔트로피를 계산한 tensor값이다. tensor 객체의 backward 함수를 불러 그래디언트를 계산한다.

backward() 함수에 대해서 공식 도큐멘트에서는 

Tensor.backward : computes the gradient of current tensor with respect to graph leaves

라고 설명하고 있는데, leave 노드는 연산 그래프에서의 데이터를 말한다.

즉, 그래프의 노드들에 대해 현재 텐서의 그래디언트를 구하는 것이다 (현재 텐서값이 도출되는 데 노드들이 얼마나 기여했는지 계산한다는 뜻 같다)

 

- 마지막으로, optimizer.step() : step 메소드를 통해 한 번의 최적화 과정(파라미터 업데이트)을 수행한다

 

 

모델 평가(추론)

num_correct = 0
num_samples = 0
model.eval()

with torch.no_grad():
    for x, y in loader:
        x = x.to(device=device)
        y = y.to(device=device)
        x = x.reshape(x.shape[0], -1)

        scores = model(x)
        _, predictions = scores.max(1)
        num_correct += (predictions == y).sum()
        num_samples += predictions.size(0)

 

- model.eval() : evaluation 모드로 모델을 전환. dropout이나 batchnorm과 같이 훈련에만 사용되는 레이어를 off하는 역할을 한다.

- torch.no_grad() : 그래디언트 연산을 수행하지 않을 때 사용한다. 이 컨텍스트 내부에서 발생하는 텐서 객체들은 requires_grad=False 옵션이 저절로 세팅되어 메모리 사용량을 아낄 수 있다.

 

 

이로써 pytorch에서 데이터를 준비, 전처리하고 모델을 훈련, 추론하는 과정을 확인해봤다.

다음 글에서는 requires_grad, 즉 파이토치에서의 그래디언트 연산에 대해서 좀더 자세히 이야기해보려 한다.