파이토치(PyTorch)로 텐서플로우 튜토리얼에 있는 MNIST 예제를 재현해 보았습니다. 이 코드는 파이토치의 MNIST 예제를 참고했으며 주피터 노트북으로 작성되어 깃허브에 올려져 있습니다. 당연하지만 분류 결과는 텐서플로우로 만든 예제와 큰 차이가 없습니다.

파이토치는 버전 0.1.6을 사용하였습니다. 파이토치 설치는 이 포스트를 참고하세요. 파이토치는 코어 C 라이브러리를 토치(Torch)와 공유하고 있습니다. 그 외에는 모두 파이썬으로 구현되어 있으며 별도로 C API 가 없습니다. 오직 파이썬을 위해서 구현된 라이브러리이고 또 다이나믹 컴퓨테이션 그래프(Dynamic Computation Graph)를 지원하기 때문에 코드 모습이 일반 프로그래밍처럼 자연스러웠습니다.

파이토치의 예제와 유사하게 torch.nn.Module 클래스를 상속받아 모델 클래스를 만들었습니다.

class MnistModel(nn.Module):
    def __init__(self):
        super(MnistModel, self).__init__()
        # input is 28x28
        # padding=2 for same padding
        self.conv1 = nn.Conv2d(1, 32, 5, padding=2)
        # feature map size is 14*14 by pooling
        # padding=2 for same padding
        self.conv2 = nn.Conv2d(32, 64, 5, padding=2)
        # feature map size is 7*7 by pooling
        self.fc1 = nn.Linear(64*7*7, 1024)
        self.fc2 = nn.Linear(1024, 10)
        
    def forward(self, x):
        x = F.max_pool2d(F.relu(self.conv1(x)), 2)
        x = F.max_pool2d(F.relu(self.conv2(x)), 2)
        x = x.view(-1, 64*7*7)   # reshape Variable
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x)
    
model = MnistModel()

Module 를 상속받은 클래스는 뉴럴 네트워크의 정방향 계산을 수행하는 forward() 메소드를 구현해야만 합니다. 이 forward() 메소드는 model 오브젝트를 데이터와 함께 호출하면 자동으로 실행이됩니다.

풀링과 활성화 함수들은 torch.nn.functional 아래에 있고 함수처럼 불러서 사용할 수 있습니다. 이런 함수들은 Conv2d 클래스처럼 유사하게 forward(), backward() 메소드를 가진 클래스를 감싸고 있는 형태입니다. 최종 loss 에서 backward() 함수가 호출되면 거꾸로 짚어가면서 역방향 계산이 연쇄적으로 일어납니다. Conv2d 함수의 파라미터에서 볼 수 있듯이 패딩과 스트라이드를 명시적으로 지정해 주어야 합니다. 콘볼루션의 패딩 계산은 이 포스트를 참고하세요.

Linear 함수는 일반적인 완전연결(fully connected) 레이어를 의미하는 것으로 이것뿐만 아니라 많은 요소들이 토치의 네이밍을 따르고 있는 것 같습니다. 아마도 토치를 썼던 사람은 쉽게 사용할 수 있을 것 같습니다.

model.train()
train_loss = []
train_accu = []
i = 0
for epoch in range(15):
    for data, target in train_loader:
        data, target = Variable(data), Variable(target)
        optimizer.zero_grad()
        output = model(data)
        loss = F.nll_loss(output, target)
        loss.backward()    # calc gradients
        train_loss.append(loss.data[0])
        optimizer.step()   # update gradients
        prediction = output.data.max(1)[1]   # first column has actual prob.
        accuracy = prediction.eq(target.data).sum()/batch_size*100
        train_accu.append(accuracy)
        if i % 1000 == 0:
            print('Train Step: {}\tLoss: {:.3f}\tAccuracy: {:.3f}'.format(i, loss.data[0], accuracy))
        i += 1

모델을 학습시키는 부분도 직관적입니다. Negative Log Likelihood를 의미하는 nll_loss 함수에 출력과 타겟값을 넣고 loss.backward() 를 호출한 다음 optimizer.step() 으로 파라미터를 업데이트 합니다. 여기서 target 은 원-핫-인코딩이 아닌 타겟 레이블을 그대로 넣습니다. loss 와 accuracy 를 리스트에 저장하여 맷플롯립으로 그래프를 그렸습니다.

loss

accuracy

모델을 훈련할 땐 model.train() , 테스트할 땐 model.eval() 함수를 호출하면 Module 클래스의 훈련 상태 여부를 바꾸게 되며 이는 dropout 같은 함수가 드롭아웃을 적용할지 안할지를 결정하게 만듭니다.

model.eval()
correct = 0
for data, target in test_loader:
    data, target = Variable(data, volatile=True), Variable(target)
    output = model(data)
    prediction = output.data.max(1)[1]
    correct += prediction.eq(target.data).sum()

print('Test set: Accuracy: {:.2f}%'.format(100. * correct / len(test_loader.dataset)))

최종결과는 99.29%를 얻었습니다. 전체 코드는 깃허브에서 확인할 수 있습니다.