[Pytorch] MNIST 데이터셋으로 간단한 CNN 구현하기

이번 글은 MNIST 데이터셋으로 간단한 CNN을 구현해볼 것이다. MLP를 구현하였을 때와 같이 관련 패키지를 불러들이고, parameter 설정을 하고, MNIST 데이터셋을 불러들어와 로딩까지 한 번에 진행할 것이다.

import torch
import torchvision.datasets as dsets
import torchvision.transforms as transforms
import torch.nn.init
device = 'cuda' if torch.cuda.is_available() else 'cpu'

# for reproducibility
torch.manual_seed(777)
if device == 'cuda':
    torch.cuda.manual_seed_all(777)

# parameters
learning_rate = 0.001
training_epochs = 15
batch_size = 100

# MNIST dataset
mnist_train = dsets.MNIST(root='MNIST_data/',
                          train=True,
                          transform=transforms.ToTensor(),
                          download=True)

mnist_test = dsets.MNIST(root='MNIST_data/',
                         train=False,
                         transform=transforms.ToTensor(),
                         download=True)

# dataset loader
data_loader = torch.utils.data.DataLoader(dataset=mnist_train,
                                          batch_size=batch_size,
                                          shuffle=True,
                                          drop_last=True)

이번에는 CNN(Convolutional Neural Network) 모델을 구성할 것이다. CNN은 MLP와 달리 클래스 설정을 해준 다음 각각의 layer들을 생성한다. 여담으로 kernel size를 filter size라고도 한다. 이를 인지하면 아래 주석을 이해하는 데에 더욱 편리할 것이다. 

class CNN(torch.nn.Module):
    
    def __init__(self):
        super(CNN, self).__init__()
        #첫번째 layer 생성
        self.layer1 = torch.nn.Sequential(
            '''
            아래는 convolutional layer을 만드는 코드이다. 코드 구성은
            torch.nn.Conv2d(input channel 개수, output channel 개수, filter의 크기, stride, padding)
            로 이루어진다.
            '''
            torch.nn.Conv2d(1, 32, kernel_size = 3, stride =1, padding = 1),
            #ReLU함수에 통과시킨다.
            torch.nn.ReLU(),
            '''
            Max pooling을 진행한다. 코드 구성은
            torch.nn.MaxPool2d(filter 크기, stride)로 이루어진다.
            '''
            torch.nn.MaxPool2d(kernel_size = 2, stride = 2))
        #두 번째 layer 생성, 이번에는 output filter이 64개이다.
        self.layer2 = torch.nn.Sequential(
            torch.nn.Conv2d(32, 64, kernel_size=3, stride =1, padding = 1),
            torch.nn.ReLU(),
            torch.nn.MaxPool2d(kernel_size = 2, stride = 2))
        #fully connected layer 생성
        self.fc = torch.nn.Linear(7 * 7 * 64, 10, bias = True) #보충설명1
        #가중치 초기화
        torch.nn.init.xavier_uniform_(self.fc.weight)
   
   #모델 설계 후 데이터셋이 layer들을 통과할 수 있게 함.
    def forward(self, x):
        out = self.layer1(x)
        out = self.layer2(out)
        out = out.view(out.size(0), -1) #보충설명2
        out = self.fc(out)
        return out

#보충설명1 으로 표기된 부분을 살펴보자. fully connected layer들은 convolutional layer이 아닌 MLP와 같은 형태로 torch.nn.Linear() 함수를 사용한다. 다만 이 줄을 유심히 살펴보면 7 * 7 * 64개의 input node로 10개의 output node(0~9까지의 숫자)를 출력하라는 것을 확인할 수 있다. 그렇다면 input node가 7 * 7 * 64 개인 이유는 무엇일까?

 

그 해답은 convolution layer의 구조를 보면 알 수 있다. 

먼저 임의의 변수가 이미지에서 tensor 형태로 변형되었다고 치자. 그 임의의 변수를 inputs이라고 하겠다

inputs = torch.Tensor(1,1,28,28) #(batch size, channel, height, width)

위와 같은 2개의 convolution layer이 모델링 되었다고 치자. 

conv1 = torch.nn.Conv2d(1,32,kernel_size = 3, stride = 1, padding = 1) #input channel, output channel, filter크기, padding
pool = torch.nn.MaxPool2d(kernel_size = 2, stride = 2)
conv2 = torch.nn.Conv2d(32, 64, 3, 1, padding =1)

이제  inputs에 convolution layer 한 층을 통과시키고 그 크기를 알아보자. 

out = conv1(inputs)
out.size()

결과값으로 torch.Size([1, 32, 28, 28])가 나온다. 

그 말인 즉슨 output volume size가 28 * 28 * 32 (witdth * height * depth) 라는 뜻이다.

이를 도출하는 방법은 아래와 같다. Stanford University의 cs321n강의 Lecture 5에서 퍼온 피피티 슬라이드이다.

 

 

그렇게 두 번째 convolution layer 및 pooling을 통과한 후의 volume size를 확인해보자

out = pool(out)
out = conv2(out)
out = pool(out)
out.size()

결과값으로 torch.Size([1, 64, 7, 7])이 나왔으며, output volume size가 7 * 7 * 64 임을 의미한다. 그래서 fully connected layer에서의 input size가 7 * 7 * 64로 나오게 되는 것이다.

 

이번에는 #보충설명2로 표기된 부분을 살펴보자.

우리는 이 7 * 7 * 64 volume size를 batch size마다 하나의 벡터로 나열하여 fully connected layer을 통과시킨다. 

앞서 batch size를 1이라고 설정하였고, 그 1 값은 torch.Size의 첫 번째 인덱스에서 out.size(0)으로 확인할 수 있다. 그래서 out = out.view(out.size(0), -1)를 통하여 batch size를 기준으로 한 row를 만들고 나머지 값을 그 열에다 입력하라고 지시한다. 

 

 

 

이제 CNN model을 실행시켜보자. 

# instantiate CNN model
model = CNN().to(device)

 

cost/loss & optimizer 설정을 해주자.

# define cost/loss & optimizer
criterion = torch.nn.CrossEntropyLoss().to(device)    # Softmax is internally computed.
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

실제 MNIST training set으로 CNN을 학습시켜 보자.

total_batch = len(data_loader)
print('Learning started. It takes sometime.')
for epoch in range(training_epochs):
    avg_cost = 0

    for X, Y in data_loader:
        # image is already size of (28x28), no reshape
        # label is not one-hot encoded
        X = X.to(device)
        Y = Y.to(device)

        optimizer.zero_grad()
        '''
        우리가 만든 CNN을 model로 instantiate 해주었는데,
        model에 우리의 training set을 대입해 본다.
        '''
        hypothesis = model(X)
        cost = criterion(hypothesis, Y)
        cost.backward()
        optimizer.step()

        avg_cost += cost / total_batch

    print('[Epoch: {:>4}] cost = {:>.9}'.format(epoch + 1, avg_cost))

print('Learning Finished!')

그 결과값은 아래와 같다.

Learning started. It takes sometime.

[Epoch: 1] cost = 0.225655958

[Epoch: 2] cost = 0.0630160421

[Epoch: 3] cost = 0.0462756902

[Epoch: 4] cost = 0.0374853872

[Epoch: 5] cost = 0.0314579383

[Epoch: 6] cost = 0.026188707

[Epoch: 7] cost = 0.0217339192

[Epoch: 8] cost = 0.0183844231

[Epoch: 9] cost = 0.016300844

[Epoch: 10] cost = 0.0133926449

[Epoch: 11] cost = 0.0100590801

[Epoch: 12] cost = 0.00998833869

[Epoch: 13] cost = 0.00781068346

[Epoch: 14] cost = 0.00675791921

[Epoch: 15] cost = 0.00778320711

Learning Finished!

 

epoch가 진행될수록 cost가 줄어드는 것을 확인할 수 있다. 그러므로 성공적인 모델이다.

 

마지막으로 모델을 testing set에다가 학습시킬 것이다. 

 

# Test model and check accuracy
with torch.no_grad():
    X_test = mnist_test.test_data.view(len(mnist_test), 1, 28, 28).float().to(device)
    Y_test = mnist_test.test_labels.to(device)

    prediction = model(X_test)
    correct_prediction = torch.argmax(prediction, 1) == Y_test
    accuracy = correct_prediction.float().mean()
    print('Accuracy:', accuracy.item())

Accuracy: 0.986299991607666

 

출처: GitHub - deeplearningzerotoall/PyTorch: Deep Learning Zero to All - Pytorch

GitHub - deeplearningzerotoall/PyTorch: Deep Learning Zero to All - Pytorch