PyTorchでMNISTをやってみる - 今日も窓辺でプログラム

はじめに

PytorchでMNISTをやってみたいと思います。
chainerに似てるという話をよく見かけますが、私はchainerを触ったことがないので、公式のCIFAR10のチュートリアルをマネする形でMNISTに挑戦してみました。

Training a classifier — PyTorch Tutorials 0.3.0.post4 documentation

はじめに
データの用意
モデルの定義
コスト関数と学習方法を定義
学習させてみる
テスト
テストデータの予測例
Jupyter Notebook

データの用意

PyTorchにはtorchvisionという、有名なデータセットや画像処理でよく使う関数などをまとめたパッケージが存在しているようです。
MNISTのデータセットをウンロード・読み込みする手段も提供されているので、今回はこれを使用してデータを準備してみます。

早速試しにダウンロードして中身を確認してみます。

import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torchvision.datasets import MNIST

# データセットをダウンロード
mnist_data = MNIST('~/tmp/mnist', train=True, download=True, transform=transforms.ToTensor())
data_loader = DataLoader(mnist_data,
                         batch_size=4,
                         shuffle=False)

torchvision.datasets.MNISTクラスでデータセットをダウンロードできます。最初の引数がデータがダウンロードされるディレクトリなので、適宜変更してください。
ダウンロードしたデータセットに変換をかけたい場合は、transform引数に関数を渡してあげるとよいみたいです。
今回は特に複雑な変換や正規化はせず、単純にtransforms.ToTensor()でtorch.Tensorクラスに変換しています。

torch.utils.data.DataLoaderクラスで、データセットをロードして指定したバッチサイズに分割して順に読み込みできる状態にしてくれます。

次のようにすると、1つ目のデータを取り出して、matplotlibで画像を表示できます。

data_iter = iter(data_loader)
images, labels = data_iter.next()

# matplotlibで1つ目のデータを可視化してみる
npimg = images[0].numpy()
npimg = npimg.reshape((28, 28))
plt.imshow(npimg, cmap='gray')
print('Label:', labels[0])

f:id:kanohk:20171218125526p:plain

ラベルは5なので、確かによく見るMNISTのデータのようです。

では、今度はMNISTやDataLoaderの引数を少しいじって、トレーニングデータとテストデータを用意しておきます。
バッチサイズは4、トレーニングデータは順番をシャッフルする設定にしてあります。

# 訓練データとテストデータを用意
train_data = MNIST('~/tmp/mnist', train=True, download=True, transform=transforms.ToTensor())
train_loader = DataLoader(mnist_data,
                         batch_size=4,
                         shuffle=True)
test_data = MNIST('~/tmp/mnist', train=False, download=True, transform=transforms.ToTensor())
test_loader = DataLoader(test_data,
                         batch_size=4,
                         shuffle=False)

モデルの定義

MNISTの画像サイズは28*28なので、1次元ベクトルにすると784次元になります。今回はPyTorchの使い方の概要を知ることが目的なので、モデル自体は単純なFeed Forwardネットワークにします。
784次元の入力層、50次元の中間層、10次元の出力層といった構成です。

モデル自体は、nn.Moduleクラスを継承して定義します。コンストラクタ内でネットワークの構成を定義してあげて、forward関数内で順伝播の処理を実装してあげます。
この2つの関数を定義するだけで、逆伝播部分は自動で面倒を見てくれます。

from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.l1 = nn.Linear(28 * 28, 50) # 入力層から隠れ層へ
        self.l2 = nn.Linear(50, 10) # 隠れ層から出力層へ
        
    def forward(self, x):
        x = x.view(-1, 28 * 28) # テンソルのリサイズ: (N, 1, 28, 28) --> (N, 784)
        x = self.l1(x)
        x = self.l2(x)
        return x
    
net = Net()

今回の方法で得た入力データはToTensor()関数で変換されているので、(1, 28, 28)という次元のテンソルになります。
上記__init__で定義したモデルは784次元のベクトルの入力を受け入れるので、torch.Tensor.view関数を使ってテンソルをリサイズする必要があります。
Nがバッチサイズの時、(N, 1, 28, 28)という次元のテンソルを(N, 784)という次元のテンソルに変換しているのがforward関数の1行目になります。

コスト関数と学習方法を定義

コスト関数と学習方法を定義します。代表的なコスト関数や最適化手法はあらかじめ提供されているので、今回はそれに乗っかります。
コスト関数にクロスエントロピー、最適化手法にSGDを選択した場合、コードは次のようになります。
optim.SGDに引数として最適化対象のパラメータ一覧を渡しています。モデルを定義する際に継承したnn.Moduleクラスがparameters()というパラメータ一覧を取得するメソッドを提供しているので、非常に簡単です。

# コスト関数と最適化手法を定義
import torch.optim as optim

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.01)

学習させてみる

今まで書いたコードを使って、早速学習させてみます。ひとまず、学習を3エポック分回してみます。
学習ループ内では次のような作業を順に行っていて、それぞれがなんと1行で書けてしまいました

入力をtorch.autograd.Variableに変換
逆伝播で使う勾配情報をリセット
順伝播
ロスの計算
逆伝播
得られた勾配を使ってパラメータを更新

for epoch in range(3):
    running_loss = 0.0
    for i, data in enumerate(train_loader):
        inputs, labels = data
        
        # Variableに変換
        inputs, labels = Variable(inputs), Variable(labels)
        
        # 勾配情報をリセット
        optimizer.zero_grad()
        
        # 順伝播
        outputs = net(inputs)
        
        # コスト関数を使ってロスを計算する
        loss = criterion(outputs, labels)
        
        # 逆伝播
        loss.backward()
        
        # パラメータの更新
        optimizer.step()
        
        running_loss += loss.data[0]
        
        if i % 5000 == 4999:
            print('%d %d loss: %.3f' % (epoch + 1, i + 1, running_loss / 1000))
            running_loss = 0.0
            
print('Finished Training')

実際にこれを走らせると、次のように損失が徐々に小さくなっていくことが確認できました。

1 5000 loss: 2.734
1 10000 loss: 1.687
1 15000 loss: 1.614
2 5000 loss: 1.553
2 10000 loss: 1.501
2 15000 loss: 1.510
3 5000 loss: 1.455
3 10000 loss: 1.429
3 15000 loss: 1.505
Finished Training

テスト

次にテストデータを使って精度を確認してみます。
テストデータ自体はtest_loaderで読み込める状態になっているので、テストデータの入力をモデルに渡して、出力から予測ラベルを得て、正解率を算出するだけです。

import torch

correct = 0
total = 0
for data in test_loader:
    inputs, labels = data
    outputs = net(Variable(inputs))
    _, predicted = torch.max(outputs.data, 1)
    total += labels.size(0)
    correct += (predicted == labels).sum()
    
print('Accuracy %d / %d = %f' % (correct, total, correct / total))

これで92%の精度が出ていることが確認できました。ちゃんと学習ができていそうです。

Accuracy 9200 / 10000 = 0.920000

テストデータの予測例

テストデータを1つ取り出して、実際にどのような画像にどのようなラベルを予測しているか見てみます。

test_iter = iter(test_loader)
inputs, labels = test_iter.next()
outputs = net(Variable(inputs))
_, predicted = torch.max(outputs.data, 1)

plt.imshow(inputs[0].numpy().reshape(28, 28), cmap='gray')
print('Label:', predicted[0])

すると、下記画像に対して7を予測していることがわかりました。確かにちゃんと予測できていそうですね。
f:id:kanohk:20171218155718p:plain