실습: CNN으로 MNIST 분류 구현하기
이번에는 앞서 배운 CNN을 활용하여 기존 MNIST 분류 프로젝트를 고도화하고자 합니다. 이미 우리는 기존에 잘 동작하는 학습 및 추론 코드를 가지고 있기 때문에, 최소한으로 수정하여 구현하고자 합니다.
본격적인 시작에 앞서 파일 구조를 먼저 고도화하려 합니다. 다음은 기존의 파일 구성 형태입니다. 각 파일들의 역할은 이전 챕터의 설명을 참고해주세요.
- 기존 파일 구성
- predict.ipynb
- model.py
- utils.py
- train.py
- trainer.py
기존에는 평면적인 디렉터리 구조를 가지고 있었다면, 이번에는 좀 더 계층화하여 파일 구조를 구성하고자 합니다.
- 새로운 파일 구성
- predict.ipynb
- mnist_classifier
- __init__.py
- models
- fc.py
- __init__.py
- cnn.py
- utils.py
- trainer.py
- train.py
mnist_classifier 라는 패키지를 구성하였으며, 패키지 내부의 models 디렉터리에 기존 선형 계층 기반 모델을 fc.py 라는 파일로 넣고, 이번에 구현할 CNN 기반의 모델을 cnn.py 라는 파일로 구성하였습니다. 따라서 추후 순환 신경망RNN을 활용하여 모델을 추가하고자 할 때, models 디렉터리 내에 파일을 추가하게 될 것입니다.
모델 구조 설계
우리는 MNIST 이미지를 분류하고자 합니다. MNIST 데이터셋은 $28\times28$ 크기의 흑백 이미지가 학습셋 60,000장, 테스트셋 10,000장으로 구성되어 있습니다. 예전에는 선형 계층linear layer 기반으로 모델을 구성하였기 때문에 행렬 형태의 이미지를 784차원의 벡터로 만들어 모델에 넣어주었지만, CNN 기반의 모델에서는 행렬 형태 그대로 넣어줄 수 있게 되었습니다.
앞서 이야기한대로 CNN 블록을 재활용하는 형태로 모델 구조를 설계하면 다음의 그림과 같이 진행될 것입니다.
마지막 CNN 블록을 통과하면 미니배치 텐서는 $(N,512,1,1)$ 의 형태가 될 것이고, 마지막 두 차원을 squeeze 함수를 통해 제거하면 $(N,512)$ 형태의 텐서를 얻을 수 있게 됩니다. 이 텐서를 ‘선형 계층 + 활성 함수 + 선형 계층 + 소프트맥스 함수’에 차례대로 통과시키면 $(N, |\mathcal{C}|)$ 크기의 텐서를 얻을 수 있습니다. 이 $(N,|\mathcal{C}|)$ 크기의 출력 텐서는 미니배치의 각 샘플별로 클래스에 대한 확률 값을 담고 있을 것입니다.
모델 클래스 구현
./mnist_classifier/cnn.py 파일에 CNN 기반 분류 모델을 구현할 것입니다. 먼저 CNN 블록을 구현하도록 합니다. 다음의 코드를 보면 역시 nn.Module을 상속받아 클래스를 정의하고 있는 것을 볼 수 있습니다. 그리고 init 함수와 forward 함수를 오버라이드하여 모델을 구현하고 있습니다.
class ConvolutionBlock(nn.Module):
def __init__(self, in_channels, out_channels):
self.in_channels = in_channels
self.out_channels = out_channels
super().__init__()
self.layers = nn.Sequential(
nn.Conv2d(in_channels, out_channels, (3, 3), padding=1),
nn.ReLU(),
nn.BatchNorm2d(out_channels),
nn.Conv2d(out_channels, out_channels, (3, 3), stride=2, padding=1),
nn.ReLU(),
nn.BatchNorm2d(out_channels),
)
def forward(self, x):
# |x| = (batch_size, in_channels, h, w)
y = self.layers(x)
# |y| = (batch_size, out_channels, h, w)
return y
init 함수 내부에 nn.Sequential을 활용하여 여러 계층을 구성하고 있고, 그 내부에 합성곱 계층과 활성 함수, 그리고 배치 정규화batch normalization이 구현되어 있습니다. 합성곱 계층은 nn.Conv2d 클래스를 통해 구현하고 있는데요. 첫 번째 인자는 입력 채널의 갯수, 두 번째 인자는 출력 채널의 갯수[1]를 넣어줍니다. 그리고 $3\times3$ 크기의 커널을 사용하도록 하고, 필요에 따라 패딩 크기와 스트라이드 크기를 넣어줍니다. 이 CNN 블록을 통과하는 입출력 텐서의 모양은 주석으로 forward 메서드 내부에 정리되어 있습니다.
이제 이렇게 정의한 CNN 블록을 활용하여 심층신경망 전체를 구성하려 합니다. 다음의 코드는 심층신경망 전체를 구성하는 코드입니다.
class ConvolutionalClassifier(nn.Module):
def __init__(self, output_size):
self.output_size = output_size
super().__init__()
self.blocks = nn.Sequential( # |x| = (n, 1, 28, 28)
ConvolutionBlock(1, 32), # (n, 32, 14, 14)
ConvolutionBlock(32, 64), # (n, 64, 7, 7)
ConvolutionBlock(64, 128), # (n, 128, 4, 4)
ConvolutionBlock(128, 256), # (n, 256, 2, 2)
ConvolutionBlock(256, 512), # (n, 512, 1, 1)
)
self.layers = nn.Sequential(
nn.Linear(512, 50),
nn.ReLU(),
nn.BatchNorm1d(50),
nn.Linear(50, output_size),
nn.LogSoftmax(dim=-1),
)
def forward(self, x):
assert x.dim() > 2
if x.dim() == 3:
# |x| = (batch_size, h, w)
x = x.view(-1, 1, x.size(-2), x.size(-1))
# |x| = (batch_size, 1, h, w)
z = self.blocks(x)
# |z| = (batch_size, 512, 1, 1)
y = self.layers(z.squeeze())
# |y| = (batch_size, output_size)
return y
두 개의 nn.Sequential 객체가 self.blocks 와 self.layers 에 할당되는 것을 볼 수 있는데요. self.blocks 에는 CNN 블록들을 넣어주고, self.layers 에는 선형 계층과 활성함수, 배치 정규화, 로그소프트맥스log-softmax 계층을 넣어준 것을 볼 수 있습니다. 마지막 CNN 블록의 출력 채널 갯수와 첫 번째 선형 계층의 입력 크기가 같음에 주목하세요.
forward 함수에는 self.blocks 와 self.layers 에 차례로 텐서를 통과시키는 코드가 구현되어 있는데요. self.layers 에 텐서를 넣기 전에 squeeze 함수를 활용해서 마지막 두 차원을 없애주는 것을 확인할 수 있습니다.
[1]: 앞서 다루었던 대로 이것은 커널의 갯수와 같습니다.
train.py 수정하기
앞서 우리는 이미 잘 짜여진 MNIST 분류 프로젝트를 가지고 있기 때문에, 아주 약간의 수정만으로도 학습과 추론이 원활하게 동작하는 것을 볼 수 있을것입니다. 먼저 파일 구조가 바뀌었기 때문에, import 를 수행하는 부분의 코드들이 다음과 같이 바뀐 것을 볼 수 있습니다.
import argparse
import torch
import torch.nn as nn
import torch.optim as optim
from mnist_classifier.trainer import Trainer
from mnist_classifier.utils import load_mnist
from mnist_classifier.utils import split_data
from mnist_classifier.utils import get_model
그리고 이제 train.py 를 실행할 때 어떤 모델을 활용하여 분류기를 학습할지 선택하는 옵션을 추가해야 합니다. define_argparser 함수 내부에 다음의 코드를 추가합니다.
p.add_argument("--model", default="fc", choices=["fc", "cnn"])
그럼 사용자는 다음과 같이 train.py 를 실행하면 CNN 기반의 MNIST 분류기를 얻을 수 있습니다.
$ train.py --model_fn ./model.pth --model cnn
그리고 앞서 우리는 선형 계층 기반으로 분류 문제를 풀었기 때문에 $28\times28$ 크기의 MNIST 이미지를 784 차원의 벡터로 변환하여 불러왔습니다만, 이제 CNN 기반의 모델에서는 기존 행렬 형태 그대로 활용할 수 있습니다. 그래서 load_mnist 함수에는 미리 flatten 이라는 인자를 만들어두었는데요. 항상 True로 넣어주던 인자를 config.model 의 값이 “fc” 일때만 벡터로 변환하도록 다음과 같이 수정해줍니다.
x, y = load_mnist(is_train=True, flatten=(config.model == "fc"))
그리고 이전에는 train.py 에서 직접 ImageClassifier 클래스를 생성하여 객체를 할당하였는데, utils.py 에 함수를 만들어 재사용하도록 합니다. 다음의 코드는 utils.py 에 정의된 get_model 함수입니다.
from mnist_classifier.models.fc import ImageClassifier
from mnist_classifier.models.cnn import ConvolutionalClassifier
def get_model(input_size, output_size, config, device):
if config.model == "fc":
model = ImageClassifier(
input_size=input_size,
output_size=output_size,
hidden_sizes=get_hidden_sizes(
input_size,
output_size,
config.n_layers
),
use_batch_norm=not config.use_dropout,
dropout_p=config.dropout_p,
).to(device)
elif config.model == "cnn":
model = ConvolutionalClassifier(output_size)
else:
raise NotImplementedError
return model
config.model 의 값에 따라 필요한 클래스를 생성하여 반환하는 것을 볼 수 있습니다. 이것을 train.py 에서 받아 model 이라는 변수에 할당합니다.
model = get_model(
input_size,
output_size,
config,
device,
)
그럼 이제 다음과 같이 코드를 실행해보면 동작하는 것을 볼 수 있습니다.
$ python train.py --model_fn ./model.pth --n_epochs 20 --model cnn
Train: torch.Size([48000, 28, 28]) torch.Size([48000])
Valid: torch.Size([12000, 28, 28]) torch.Size([12000])
ConvolutionalClassifier(
(blocks): Sequential(
(0): ConvolutionBlock(
(layers): Sequential(
(0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
(1): ReLU()
(2): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(3): Conv2d(32, 32, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
(4): ReLU()
(5): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)
...
)
(layers): Sequential(
(0): Linear(in_features=512, out_features=50, bias=True)
(1): ReLU()
(2): BatchNorm1d(50, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(3): Linear(in_features=50, out_features=10, bias=True)
(4): LogSoftmax(dim=-1)
)
)
Adam (
Parameter Group 0
amsgrad: False
betas: (0.9, 0.999)
eps: 1e-08
lr: 0.001
weight_decay: 0
)
NLLLoss()
Epoch(1/20): train_loss=1.9247e-01 valid_loss=9.3223e-02 lowest_loss=9.3223e-02
Epoch(2/20): train_loss=5.1895e-02 valid_loss=6.6272e-02 lowest_loss=6.6272e-02
Epoch(3/20): train_loss=3.8363e-02 valid_loss=3.8659e-02 lowest_loss=3.8659e-02
Epoch(4/20): train_loss=2.7718e-02 valid_loss=3.6389e-02 lowest_loss=3.6389e-02
Epoch(5/20): train_loss=2.2694e-02 valid_loss=3.4458e-02 lowest_loss=3.4458e-02
Epoch(6/20): train_loss=1.7611e-02 valid_loss=3.7527e-02 lowest_loss=3.4458e-02
Epoch(7/20): train_loss=1.7853e-02 valid_loss=4.3038e-02 lowest_loss=3.4458e-02
Epoch(8/20): train_loss=1.4361e-02 valid_loss=5.5245e-02 lowest_loss=3.4458e-02
Epoch(9/20): train_loss=1.4388e-02 valid_loss=4.4018e-02 lowest_loss=3.4458e-02
Epoch(10/20): train_loss=1.1478e-02 valid_loss=3.7882e-02 lowest_loss=3.4458e-02
Epoch(11/20): train_loss=1.0984e-02 valid_loss=4.4866e-02 lowest_loss=3.4458e-02
Epoch(12/20): train_loss=9.7303e-03 valid_loss=2.9497e-02 lowest_loss=2.9497e-02
Epoch(13/20): train_loss=8.5712e-03 valid_loss=2.9458e-02 lowest_loss=2.9458e-02
Epoch(14/20): train_loss=7.2910e-03 valid_loss=3.9303e-02 lowest_loss=2.9458e-02
Epoch(15/20): train_loss=7.4639e-03 valid_loss=3.8095e-02 lowest_loss=2.9458e-02
Epoch(16/20): train_loss=9.1744e-03 valid_loss=4.2406e-02 lowest_loss=2.9458e-02
Epoch(17/20): train_loss=8.3403e-03 valid_loss=3.1964e-02 lowest_loss=2.9458e-02
Epoch(18/20): train_loss=6.2705e-03 valid_loss=3.6520e-02 lowest_loss=2.9458e-02
Epoch(19/20): train_loss=5.7355e-03 valid_loss=3.3537e-02 lowest_loss=2.9458e-02
Epoch(20/20): train_loss=5.2022e-03 valid_loss=4.2961e-02 lowest_loss=2.9458e-02
애초에 우리는 매우 잘 정돈된 MNIST 분류 코드를 가지고 있었기 때문에, 매우 적은 수정만으로도 훌륭하게 동작하는 CNN 분류기를 구현할 수 있습니다.
predict.ipynb 수정하기
테스트셋에 대한 추론을 수행하고 성능을 구하는 코드인 predict.ipynb 파일의 경우에도 매우 적은 수정만으로도 여전히 잘 동작하게 만들 수 있습니다. 앞서와 마찬가지로 import 를 수행하는 코드들을 알맞게 수정해줍니다. utils.py 에 정의된 get_model 함수를 import 하는 것을 볼 수 있습니다.
import sys
import numpy as np
import matplotlib.pyplot as plt
from mnist_classifier.utils import load_mnist
from mnist_classifier.utils import get_model
그리고 train.py 에서와 마찬가지로 load_mnist 함수의 flatten 인자를 항상 True 에서 train_config.model 의 값이 “fc” 일 때만 True 가 되도록 합니다. 또한 get_model 함수를 통해 모델 생성을 대체합니다.
# Load MNIST test set.
x, y = load_mnist(is_train=False, flatten=(train_config.model == "fc"))
x, y = x.to(device), y.to(device)
print(x.shape, y.shape)
input_size = int(x.shape[-1])
output_size = int(max(y)) + 1
model = get_model(
input_size,
output_size,
train_config,
device,
)
model.load_state_dict(model_dict)
test(model, x, y, to_be_shown=False)
이 코드를 실행하면 다음과 같이 실행 결과가 나타날텐데요.
... torch.Size([10000, 28, 28]) torch.Size([10000])
Accuracy: 0.9916
테스트셋에 대해서 무려 99.16%의 정확도를 보여주는 것을 볼 수 있습니다. 앞서 선형 계층 기반 모델은 98.37%의 점수를 보였던 것에 비하면 매우 큰 성과라고 볼 수 있습니다. 오류가 약 1.6%에서 약 0.8%로 줄어들었다는 것은 무려 절반이나 오류를 개선한 것이기 때문입니다. 또한 저자가 딱히 모델 구조를 튜닝한 것은 아니기 때문에 만약 독자분들이 실험을 거치며 튜닝을 수행한다면 더 높은 성능을 얻을 수도 있겠지요. 다만 테스트셋이 아닌 검증셋validation set에 대해서 튜닝을 수행해야함을 잊지 마세요.[2]
[2]: 제일 좋은 방법은 비공개된 테스트셋이 존재하여, 도전자들이 공개된 데이터셋을 통해 학습과 튜닝을 수행하고 비공개 테스트셋을 통해 진검 승부를 펼치는 것입니다.
마무리하며
이처럼 우리는 매우 적은 코드의 수정만으로도 기존 선형 계층 기반의 모델보다 훨씬 뛰어난 CNN 기반의 모델을 만들 수 있었습니다. 이를 비춰볼 때, 만약 추후에도 우리가 새로운 모델을 추가하거나 하이퍼 파라미터가 변경되고 학습 알고리즘이 변하더라도, 현재의 코드를 수정하면 훨씬 더 쉽게 구현할 수 있을 것입니다. 독자분들도 만약 실제로 머신러닝 프로젝트를 수행한다면 이와 같이 주피터 노트북을 벗어나서 이와 같이 CLI 환경에서 연구개발을 진행할 수 있는 환경을 구축해야 할 것입니다.