프로젝트 개요
이 튜토리얼에서는 PyTorch를 활용해 MNIST 데이터셋으로 손글씨 숫자를 인식하는 CNN(합성곱 신경망) 분류기를 학습하는 방법을 단계별로 살펴봅니다.
MNIST는 손글씨 숫자 분류 작업에 가장 널리 사용되는 데이터셋으로, 라벨이 붙은 28×28 픽셀 크기의 흑백 이미지 7만 장 이상으로 구성되어 있습니다. 학습용 이미지 약 6만 장과 테스트용 이미지 1만 장으로 나뉘며, 우리가 할 일은 6만 장의 학습 이미지로 모델을 훈련한 뒤, 1만 장의 테스트 이미지에서 분류 정확도를 평가하는 것입니다.
라이브러리 설치
먼저 최신 버전의 PyTorch와 torchvision이 필요합니다. 아직 설치하지 않았다면 터미널에서 pip 명령어로 간단히 설치할 수 있습니다.
$pip install torch torchvision
라이브러리 임포트
import torch import torchvision
MNIST 데이터셋 불러오기
본격적인 코딩에 앞서 MNIST 데이터셋이 필요합니다. 이미지와 라벨을 메모리에 불러오고, 실험에 사용할 하이퍼파라미터를 정의하겠습니다.
# n_epochs: 전체 학습 데이터셋을 반복하는 횟수 n_epochs = 3 batch_size_train = 64 batch_size_test = 1000 # learning_rate와 momentum은 옵티마이저 설정값 learning_rate = 0.01 momentum = 0.5 log_interval = 10 random_seed = 1 torch.backends.cudnn.enabled = False torch.manual_seed(random_seed)
이제 TorchVision을 사용해 MNIST 데이터셋을 불러옵니다. 학습에는 배치 크기 64, 테스트에는 배치 크기 1000을 사용하고, 정규화에는 MNIST 데이터셋의 평균값 0.1307과 표준편차 0.3081을 적용합니다.
train_loader = torch.utils.data.DataLoader(
torchvision.datasets.MNIST('/files/', train=True, download=True,
transform=torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(
(0.1307,), (0.3081,))
])
),
batch_size=batch_size_train, shuffle=True)
test_loader = torch.utils.data.DataLoader(
torchvision.datasets.MNIST('/files/', train=False, download=True,
transform=torchvision.transforms.Compose([
torchvision.transforms.ToTensor(),
torchvision.transforms.Normalize(
(0.1307,), (0.3081,))
])
),
batch_size=batch_size_test, shuffle=True)출력 결과
Downloading https://yann.lecun.com/exdb/mnist/train-images-idx3-ubyte.gz Downloading https://yann.lecun.com/exdb/mnist/train-labels-idx1-ubyte.gz Downloading https://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz Downloading https://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz Processing... Done!
test_loader로 테스트 데이터를 불러옵니다.
examples = enumerate(test_loader) batch_idx, (example_data, example_targets) = next(examples) example_data.shape
출력 결과
torch.Size([1000, 1, 28, 28])
출력 결과에서 확인할 수 있듯이, 하나의 테스트 배치는 [1000, 1, 28, 28] 형태의 텐서입니다. 즉, 28×28 픽셀 흑백 이미지 1,000장이 담겨 있다는 의미입니다.
matplotlib으로 데이터셋 일부를 시각화해 보겠습니다.
import matplotlib.pyplot as plt
fig = plt.figure()
for i in range(5):
plt.subplot(2,3,i+1)
plt.tight_layout()
plt.imshow(example_data[i][0], cmap='gray', interpolation='none')
plt.title("Ground Truth: {}".format(example_targets[i]))
plt.xticks([])
plt.yticks([])
print(fig)네트워크 구축
이번에는 2차원 합성곱(convolutional) 레이어 두 개와 완전연결(fully-connected) 레이어 두 개로 구성된 네트워크를 만들어 보겠습니다. 먼저 필요한 모듈을 임포트한 후, 네트워크 클래스를 새로 정의합니다.
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
class Net(nn.Module):
def __init__(self):
super(Net, self).__init__()
self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
self.conv2_drop = nn.Dropout2d()
self.fc1 = nn.Linear(320, 50)
self.fc2 = nn.Linear(50, 10)
def forward(self, x):
x = F.relu(F.max_pool2d(self.conv1(x), 2))
x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
x = x.view(-1, 320)
x = F.relu(self.fc1(x))
x = F.dropout(x, training=self.training)
x = self.fc2(x)
return F.log_softmax(x)네트워크와 옵티마이저를 초기화합니다.
network = Net() optimizer = optim.SGD(network.parameters(), lr=learning_rate, momentum=momentum)
모델 학습
학습 로직을 구현해 보겠습니다. 먼저 네트워크를 학습 모드로 전환한 뒤, 에포크(epoch)마다 전체 학습 데이터를 한 번씩 순회합니다. DataLoader가 개별 배치를 불러오고, optimizer.zero_grad()로 그래디언트를 초기화합니다.
train_losses = [] train_counter = [] test_losses = [] test_counter = [i*len(train_loader.dataset) for i in range(n_epochs + 1)]
깔끔한 학습 곡선을 그리기 위해 학습 손실과 테스트 손실을 저장할 리스트 두 개를 준비했습니다. x축에는 학습에 사용된 예시 데이터의 수를 표시할 계획입니다.
backward() 호출로 새로운 그래디언트 집합을 수집하면, optimizer.step()을 통해 해당 그래디언트를 네트워크의 각 파라미터에 역전파합니다.
def train(epoch):
network.train()
for batch_idx, (data, target) in enumerate(train_loader):
optimizer.zero_grad()
output = network(data)
loss = F.nll_loss(output, target)
loss.backward()
optimizer.step()
if batch_idx % log_interval == 0:
print('Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
epoch, batch_idx * len(data), len(train_loader.dataset),
100. * batch_idx / len(train_loader), loss.item()))
train_losses.append(loss.item())
train_counter.append(
(batch_idx*64) + ((epoch-1)*len(train_loader.dataset)))
torch.save(network.state_dict(), '/results/model.pth')
torch.save(optimizer.state_dict(), '/results/optimizer.pth')신경망 모듈과 옵티마이저는 .state_dict()를 통해 내부 상태를 저장하고 불러올 수 있습니다.
다음은 테스트 루프입니다. 테스트 손실을 합산하고 올바르게 분류된 숫자 개수를 추적하여 네트워크의 정확도를 계산합니다.
def test():
network.eval()
test_loss = 0
correct = 0
with torch.no_grad():
for data, target in test_loader:
output = network(data)
test_loss += F.nll_loss(output, target, size_average=False).item()
pred = output.data.max(1, keepdim=True)[1]
correct += pred.eq(target.data.view_as(pred)).sum()
test_loss /= len(test_loader.dataset)
test_losses.append(test_loss)
print('\nTest set: Avg. loss: {:.4f}, Accuracy: {}/{} ({:.0f}%)\n'.format(
test_loss, correct, len(test_loader.dataset),
100. * correct / len(test_loader.dataset)))학습을 시작하기 전에 test()를 한 번 호출하여, 무작위로 초기화된 파라미터 상태에서 모델 성능을 먼저 평가합니다.
test()
for epoch in range(1, n_epochs + 1):
train(epoch)
test()출력 결과
Test set: Avg. loss: 2.3048, Accuracy: 1063/10000 (10%) Train Epoch: 1 [0/60000 (0%)]Loss: 2.294911 Train Epoch: 1 [640/60000 (1%)]Loss: 2.314225 Train Epoch: 1 [1280/60000 (2%)]Loss: 2.290719 Train Epoch: 1 [1920/60000 (3%)]Loss: 2.294191 Train Epoch: 1 [2560/60000 (4%)]Loss: 2.246799 ... Train Epoch: 1 [36480/60000 (61%)]Loss: 0.496494 ...... ...... ....... Train Epoch: 3 [49920/60000 (83%)]Loss: 0.253500 Train Epoch: 3 [50560/60000 (84%)]Loss: 0.364354 Train Epoch: 3 [51200/60000 (85%)]Loss: 0.333843 Train Epoch: 3 [51840/60000 (86%)]Loss: 0.096922 Train Epoch: 3 [52480/60000 (87%)]Loss: 0.282102 Train Epoch: 3 [53120/60000 (88%)]Loss: 0.236428 Train Epoch: 3 [53760/60000 (90%)]Loss: 0.610584 Train Epoch: 3 [54400/60000 (91%)]Loss: 0.198840 Train Epoch: 3 [55040/60000 (92%)]Loss: 0.344225 Train Epoch: 3 [55680/60000 (93%)]Loss: 0.158644 Train Epoch: 3 [56320/60000 (94%)]Loss: 0.216912 Train Epoch: 3 [56960/60000 (95%)]Loss: 0.309554 Train Epoch: 3 [57600/60000 (96%)]Loss: 0.243239 Train Epoch: 3 [58240/60000 (97%)]Loss: 0.176541 Train Epoch: 3 [58880/60000 (98%)]Loss: 0.456749 Train Epoch: 3 [59520/60000 (99%)]Loss: 0.318569 Test set: Avg. loss: 0.0912, Accuracy: 9716/10000 (97%)
모델 성능 평가
단 3 에포크의 학습만으로 테스트 세트에서 97%의 정확도를 달성했습니다. 학습 시작 전 무작위로 초기화된 파라미터 상태에서는 테스트 세트 정확도가 10%에 불과했다는 점을 고려하면 놀라운 향상입니다.
학습 곡선 시각화
fig = plt.figure()
plt.plot(train_counter, train_losses, color='blue')
plt.scatter(test_counter, test_losses, color='red')
plt.legend(['Train Loss', 'Test Loss'], loc='upper right')
plt.xlabel('number of training examples seen')
plt.ylabel('negative log likelihood loss')
fig위 학습 곡선을 보면 에포크 수를 늘릴수록 정확도가 더 향상될 여지가 충분히 있음을 알 수 있습니다. 3 에포크 기준으로도 손실이 꾸준히 감소하는 추세를 보였습니다.
그 전에 몇 가지 예시를 더 실행하여 모델의 예측 결과와 실제 값을 비교해 보겠습니다.
with torch.no_grad():
output = network(example_data)
fig = plt.figure()
for i in range(6):
plt.subplot(2,3,i+1)
plt.tight_layout()
plt.imshow(example_data[i][0], cmap='gray', interpolation='none')
plt.title("Prediction: {}".format(
output.data.max(1, keepdim=True)[1][i].item()))
plt.xticks([])
plt.yticks([])
fig모델의 예측 결과를 확인해 보면, 해당 예시들에 대해 거의 정확하게 숫자를 맞추고 있는 것을 볼 수 있습니다.