[Pytorch] MNIST Classification (2020/12/02 수정)

(pytorch v1.7.0+cu101 / Colab을 사용했습니다.)

 

(2020/12/12 수정내용)

model의 마지막에 log_softmax는 빼야합니다. 아래에서 loss function으로 CrossEntropyLoss를 사용하는데, CrossEntropyLoss내에서 log_softmax 연산을 수행하고 있습니다. 따라서, model의 output은 log_softmax를 처리하지 않은 raw output이 되어야 합니다. log_softmax를 한번 더 수행했을뿐, 결과는 거의 유사합니다.

 

tensorflow를 사용하다가 pytorch는 어떤가 궁금해서 pytorch도 사용을 해보려고 합니다. 

기본적인 사용법은 pytorch 공식 홈페이지 튜토리얼을 참조하였고, 그 내용을 토대로 MNIST Classification을 구현해보려고 합니다.

1. Dataset 준비

torchvision이라는 모듈을 통해서 dataset을 받아올 수 있습니다.

(tensorflow나 sklearn을 사용해도 무관할 것 같네요)

 

import torch
import torchvision

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

train_dataset = torchvision.datasets.MNIST('./', train=True, download=True)
test_dataset = torchvision.datasets.MNIST('./', train=False)

이런식으로 저장하고, 받아오게 됩니다. 한번 다운로드 받으면, 두 번째는 다운받지 않아도 자동으로 읽어옵니다. 또한, train이라는 매개변수가 있는데, True/False를 통해서 train/test data를 각각 받아올 수 있습니다. 이 부분은 굳이 두 번씩 호출에서 읽어오는게 조금 불편하긴 하네요.

 

x_train_orig = train_dataset.train_data
y_train_orig = train_dataset.train_labels
x_test_orig = test_dataset.test_data
y_test_orig = test_dataset.test_labels
print('x_train shape : ', x_train_orig.shape)
print('y_train shape : ', y_train_orig.shape)
print('x_test  shape : ', x_test_orig.shape)
print('y_test  shape : ', y_test_orig.shape)

그리고 train/val/test set으로 나누어주고 normalization을 적용하였습니다.(학습에 val data는 사용하지는 않았습니다 - 어떤 방식으로 사용하는지 알아보고 다음 예제를 통해서 적용해보도록 하겠습니다.)

# normalization
x_train = (x_train_orig[10000:] / 255.).to(device=device)
y_train = y_train_orig[10000:].to(device=device)
x_val = (x_train_orig[:10000] / 255.).to(device=device)
y_val = y_train_orig[:10000].to(device=device)
x_test = x_test_orig / 255.
y_test = y_test_orig.to

print('x_train shape : ', x_train.shape)
print('y_train shape : ', y_train.shape)
print('x_train shape : ', x_val.shape)
print('y_train shape : ', y_val.shape)
print('x_test  shape : ', x_test.shape)
print('y_test  shape : ', y_test.shape)

그리고 tensorflow의 DataGenerator와 유사한 torch.utils.data.DataLoader를 사용해보려고 합니다. 

이 메소드는 사용할 데이터를 batch size 크기로 읽을 수 있는 iterator 역할을 합니다.

trainingset = torch.utils.data.TensorDataset(x_train, y_train)
train_loader = torch.utils.data.DataLoader(trainingset, batch_size=32, shuffle=True)

그전에 torch.utils.data.TensorDataset을 통해서 input과 target으로 나누어진 dataset을 합쳐주고, DataLoader의 인자로 합쳐진 dataset을 입력해줍니다. batch size는 32로 했으며, 랜덤하게 섞여서 iteration되도록 하였습니다.

 

이제 데이터 준비는 끝이 났습니다.

2. Model 구현 및 학습

두 가지 모델을 구현해서 진행해보려고 합니다.

 

첫 번째는 단순히 fc layer - relu를 반복적으로 쌓아서 구성해보겠습니다.

class Net(torch.nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.fc1 = torch.nn.Linear(784, 256)
        self.fc2 = torch.nn.Linear(256, 256)
        self.fc3 = torch.nn.Linear(256, 256)
        self.fc4 = torch.nn.Linear(256, 128)
        self.fc5 = torch.nn.Linear(128, 128)
        self.fc6 = torch.nn.Linear(128, 10)
    
    def forward(self, x):
        x = x.view(-1, 784)
        x = torch.nn.functional.relu(self.fc1(x))
        x = torch.nn.functional.relu(self.fc2(x))
        x = torch.nn.functional.relu(self.fc3(x))
        x = torch.nn.functional.relu(self.fc4(x))
        x = torch.nn.functional.relu(self.fc5(x))
        x = torch.nn.functional.dropout(x, training=self.training)
        x = self.fc6(x)
        return x

m1 = Net()

tf.keras.models.Model을 상속받아서 모델 Class를 구현하는 것와 유사합니다.

torch.nn.Module은 신경망 모듈의 기본이 되는 클래스로, 클래스 안에 각 layer와 구조를 정의하게 됩니다.

fc layer를 쌓는다고 하였는데, tf.keras.layers.Dense에 해당하는 layer가 torch에서는 torch.nn.Linear입니다. 단순히 선형 layer이며 Aw + b의 output을 출력합니다.

 

조금 코드에 대해서 설명하자면 __init__을 통해서 모델을 선언할 때 사용될 layer들을 생성합니다. 

그리고, forward 함수를 정의하는데 forward 함수는 모델이 불릴 때, forward propagation을 진행하는데 자동으로 forward 함수가 호출되게 됩니다.

저는 forward 함수 내에 relu와 log_softmax를 정의했는데, __init__ 함수에 layer로 정의해서 forward에서 호출하는 형식으로 구현할 수도 있습니다.

 

이렇게 생성한 모델을 선언해주고, 저는 GPU 환경에서 학습을 진행하기 때문에 아래와 같이 모델을 GPU 메모리에 로드해줍니다. 

m1.cuda()

그리고 loss function과 optimizer를 정의해줍니다. tensorflow의 'categorical_crossentropy'와 'adam' optimizer가 pytorch에서는 다음과 같이 설정될 수 있습니다.

criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(m1.parameters(), lr=0.001)

 

그리고 학습을 진행하기 위한 코드를 작성합니다.

for epoch in range(10):
    start = time.time()
    total_loss = 0

    for xb, yb in train_loader:
        #xb, yb = torch.autograd.Variable(xb), torch.autograd.Variable(yb)

        pred = m1(xb)
        loss = criterion(pred, yb)

        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        total_loss += loss.item()
    
    with torch.no_grad():
        pred = m1(x_train)
        acc = pred.data.max(1)[1].eq(y_train.data).sum()/len(x_train) * 100
        loss = criterion(pred, y_train)
    print(f"{time.time() - start} sec - loss : {loss} / acc : {acc}")

tensorflow에서는 단순히 model.fit 을 통해서 간단하게 학습할 수 있었지만, pytorch에서는 이렇게 for 문을 통해서 학습을 진행해주어야 하네요. tensorflow에서 GradientTape를 통해서 학습하는 것와 유사해보입니다.

(-> skorch를 사용하면 model.fit과 같이 학습이 가능합니다.. !)

 

코드를 간단히 설명하자면, 10 epoch로 학습을 진행하며 Data 전처리 과정에서 DataLoader를 사용해서 32개의 샘플 씩 스텝을 진행합니다. 

m1(xb)를 통해서 output을 얻고, 얻은 output으로 loss를 계산하고, loss.backward()를 통해서 Gradient를 계산합니다. 

pytorch가 forward 계산을 추적하고 있기 때문에, 자동으로 미분이 계산됩니다.

그리고 optimizer.step()을 통해서 parameter를 업데이트하며, 업데이트가 끝나면 optimizer.zero_grad()를 통해서 계산한 Gradient를 0으로 초기화합니다.(backward()를 통해서 구한 Gradient가 누적되므로, 업데이트가 끝나면 0으로 초기화해주는 것입니다.)

 

그리고 각 epoch마다 loss와 acc를 구하기 위해서 마지막 코드를 작성했습니다.

10번의 반복 학습으로 99%의 train acc를 얻었습니다.

pred = m1(x_test.to(device=device))
acc = pred.data.max(1)[1].eq(y_test.to(device=device).data).sum()/len(x_test) * 100
loss = criterion(pred, y_test.to(device=device))
print(f"loss : {loss} / acc : {acc}")

그리고 97.5%의 test acc를 얻었습니다.

 

 

다음으로 CNN 구조의 Model을 구현해보도록 하겠습니다. 이 구조는 tensorflow에서 구현한 모델과 동일합니다.

2020/11/22 - [ML & DL/tensorflow] - MNIST dataset 예제(ConvNet, VGG16)

MNIST dataset 예제(ConvNet, VGG16)

class CNN_Model(torch.nn.Module):
    def __init__(self):
        super(CNN_Model, self).__init__()
        self.conv1 = torch.nn.Conv2d(1, 32, 3, 1, padding=1)
        self.conv2 = torch.nn.Conv2d(32, 16, 2, 1, padding=1)
        self.fc1 = torch.nn.Linear(784, 10)
    
    def forward(self, x):
        x = x.view(-1, 1, 28, 28)
        x = torch.nn.functional.relu(self.conv1(x))
        x = torch.nn.functional.max_pool2d(x, 2)
        x = torch.nn.functional.relu(self.conv2(x))
        x = torch.nn.functional.max_pool2d(x, 2)
        x = torch.nn.Flatten()(x)
        x = self.fc1(x)
        return x

m2 = CNN_Model()

마찬가지로 torch.nn.Module을 상속받아서 CNN 구조의 모델을 정의하였습니다.

tf.keras.layers.Conv2D에 해당하는 torch.nn.Conv2d를 사용했습니다.

tensorflow와 약간의 차이가 있는데, 우선 ConvNet에서 사용되는 input shape의 형태가 (m, c, h, w) 입니다. 즉, channel이 2번째 차원에 있습니다. 그리고 padding을 'same', 'valid'로 설정할 수 있는 것과 달리 pytorch에서는 integer로 입력해주어야 합니다... ! 즉... 같은 형태로 만들고 싶다면 same shape를 얻기 위해서 계산을 해야합니다. tensorflow에 익숙해져있어서 그런지 모르겠지만.. 이 부분은 tensorflow가 훨씬 편리하네요.

 

모델 구성을 살펴보면, Conv2d - relu - max_pool2d가 두 번 반복되고 마지막에 fc층으로 10개의 output을 가지고, softmax layer를 거쳐서 출력이 나오게 됩니다. 마찬가지로 relu나 max_pool2d를 __init__에서 정의하고 호출하는 형식으로 구현할 수 있습니다.

 

m2.cuda()
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(m2.parameters(), lr=0.001)

마찬가지로 GPU에 로드시켜주고, loss와 optimizer를 정의한 후에 학습을 진행합니다.

이번에는 5 epoch로 진행했습니다.

for epoch in range(5):
    start = time.time()
    total_loss = 0

    for xb, yb in train_loader:
        #xb, yb = torch.autograd.Variable(xb), torch.autograd.Variable(yb)

        pred = m2(xb)
        loss = criterion(pred, yb.long())

        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

        total_loss += loss.item()
    
    with torch.no_grad():
        pred = m2(x_train)
        acc = pred.data.max(1)[1].eq(y_train.data).sum()/len(x_train) * 100
        loss = criterion(pred, y_train)
    print(f"{time.time() - start} sec - loss : {loss} / acc : {acc}")

Data가 많아서 그런지 5번의 epoch에도 98.7%의 정확도를 얻었습니다.

 

pred = m2(x_test.to(device=device))
acc = pred.data.max(1)[1].eq(y_test.to(device=device).data).sum()/len(x_test) * 100
loss = criterion(pred, y_test.to(device=device))
print(f"loss : {loss} / acc : {acc}")

Test Set에 대해서는 98%의 정확도를 나타내고 있습니다.

 

 

tensorflow만 사용하다가 pytorch를 사용해보니 새로웠는데, keras Model을 통해서 조금 편하게 구현할 수 있었다면, pytorch의 경우에는 조금 번거러울 수는 있지만, 그래도 custom하기에 조금 더 편하지 않을까하는 생각을 갖게되는 것 같습니다.

 

앞으로 새로운 dataset으로 구현할 때, 둘 다 사용해봐야겠네요.