PyTorchでMNIST

はじめに

PyTorchを触ってみました．
試しに，単純なパーセプトロンのみでMNIST分類をやってみます．

すべてGoogle Colabの環境で実行しました．
Python: 3.6.9
PyTorch: 1.6.0+cu101

色々読み込み

import torch, torch.nn as nn, torchvision
from torchvision import models, transforms
import matplotlib.pyplot as plt

データローダ作成

transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, ), (0.5, ))])
trainset = torchvision.datasets.MNIST(root='./', train=True, download=True, transform=transform)
testset = torchvision.datasets.MNIST(root='./', train=False, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=32, shuffle=True)
testloader = torch.utils.data.DataLoader(testset, batch_size=32, shuffle=False)

transforms.Compose()めちゃ便利

transforms.Compose([処理1, 処理2, ...])と書くだけで，上から順に加工処理を適用してくれる．
適当に抜粋↓

transforms.Resize(size) # 画像の短辺がこのサイズにフィットするようリサイズ
transforms.CenterCrop() # 画像の中央を(おそらく)短辺の長さで正方形に切り取る
transforms.RandomHorizontalFlip() # ランダムで反転
transforms.ToTensor() # Tensor型に変換 ※必須級
transforms.Normalize(mean, std) # 引数はTensor型に限る(故に最後に配置すると良さそう)

ちなみに，transforms.Normalizeは引数をタプルで渡しますが，
transforms.Normalize((0.1, 0.2, 0.3), (1.1, 1.2, 1.3) )のように書くと，左からR, G, Bの各チャネルに対応するそうです．

モデル定義

class Net(nn.Module):
  def __init__(self):
    super(Net, self).__init__()

    self.layer1 = nn.Sequential(
      nn.Linear(784, 1000),
      nn.Sigmoid()
    )
    self.classifier = nn.Sequential(
      nn.Linear(1000, 10),
      nn.Softmax(dim=1)
    )

  def forward(self, x):
    x = self.layer1(x)
    return self.classifier(x)

学習モデル．中間層のユニット数は1,000としました．
forward()に順伝播の処理順を示します．

モデルのインスタンス化

device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
net = Net()to(device)

PyTorchはTensorごとにリソース(CPU or GPU)を選択できます．
CUDAが利用できる環境であればtorch.cuda.is_available()はTrueを返すので，deviceの中身はtorch.device('cuda:0')となります．
そしてインスタンス生成時にnet = Net().to(device)としてやれば，モデルに含まれるパラメータをすべてGPU上で処理できるようになります．

損失関数と最適化手法の決定

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(net.parameters(), lr=1e-3)

損失関数はクロスエントロピー．
最適化手法はAdamで，学習率は0.001です．

学習

train_loss_history = []
val_loss_history = []
acc_history = []
epochs = 10
for epoch in range(epochs):
    train_loss_sum = 0.0
    val_loss_sum = 0.0
    
    net.train() # 学習モード
    for x, y in trainloader:
        # 勾配を初期化
        optimizer.zero_grad()

        # データ整形 & GPU処理用に変換
        x = x.view(-1, 28*28)
        x = x.to(device)
        y = y.to(device)
        
        # 順伝播
        outputs = net(x)
        loss = criterion(outputs, y)

        # 逆伝播(バックプロパゲーション)
        loss.backward()

        # 重みの更新
        optimizer.step()

        # lossの累積を加算
        train_loss_sum += loss.item()
      
    net.eval() # 評価モード
    correct = 0.0
    with torch.no_grad():
      for x, y in testloader:
        # データ整形 & GPU処理用に変換
        x = x.view(-1, 28*28)
        x = x.to(device)
        y = y.to(device)
        
        # 順伝播
        outputs = net(x)
        loss = criterion(outputs, y)

        # lossの累積を加算
        val_loss_sum += loss.item()

        # ラベルと一致していたら正答としてカウント
        pred = outputs.argmax(dim=1, keepdim=True)
        correct += pred.eq(y.view_as(pred)).sum().item()

    # 色々表示
    train_loss = train_loss_sum / len(trainloader)
    val_loss = val_loss_sum / len(testloader)
    accuracy = correct / len(testset)
    print('Epoch:{}, train_loss: {:.3f}, val_loss: {:.3f}, val_accuracy: {}'
          .format(epoch + 1, train_loss, val_loss, accuracy))
    
    # 記録
    train_loss_history.append(train_loss)
    val_loss_history.append(val_loss)
    acc_history.append(accuracy)

x = x.to(device)は，先程と同じGPU処理用の命令です．
モデルではnet.to(device)とすることで，その内部に含まれるパラメータすべてに適用しました．
データローダから新たに取り出してきたTensor達は未だCPU用なので，同様に適用させる必要があります．

net.eval()を実行することで，推論時にDropout層やBatchNormalization層をオフにできます．
（今回の例では上記レイヤは使っていないので，なくても問題ありません）

with torch.no_grad():をつけると，勾配計算が省略されるため，メモリ節約になるそうです．

ラベルと一致していたら…の部分の挙動の詳細

outputsの形状は(32, 10)です．(32: バッチサイズ, 10: クラス数)
まず，argmax(指定した軸(ここでは1)に沿って，最大値のindexを返す)を適用します．
ここで，keepdim=Trueをつけることで，2階のテンソルという情報が維持され，predの形状は(32, 1)となります．

続いて，pred.eq(y.view_as(pred)).sum().item()についてです．
a.view_as(b)は，aの形状をbの形状に合わせて変換します．
（view()という似た名前の形状変換用のメソッドがありますが，その仲間だと思います）

a = torch.arange(0, 12) # shape=(12, )
b = torch.Tensor(12).reshape(3, 4) # shape=(3, 4)

a.shape # 変換前
# torch.Size([12])

a = a.view_as(b) # bの形状に合わせてaを変換

a.shape # 変換後
# torch.Size([3, 4])

勿論，aとbの要素数は一致している必要があります．

eqは，2つのTensorを要素ごとに比較します．
要素ごとに，一致していればTrue, 違えばFalseが格納されたTensorを返します．

a = torch.tensor([1,2,3,4,5])
b = torch.tensor([0,2,0,4,5])
a.eq(b)

# tensor([False,  True, False,  True,  True])

勿論，比較対象同士の形状が同じでなければエラーが出るので注意しましょう．

a = torch.tensor([0,1,2,3,4,5,6,7,8,9])
b = torch.tensor([0,1,2])
a.eq(b)

# The size of tensor a (10) must match the size of tensor b (3) at non-singleton dimension 0

つまり，pred.eq(y.view_as(pred))は，『yをpredの形状に合わせ，yとpredの各要素を比較し，TrueかFalseが計32(=batch_size)個入ったTensorを求める』という処理です．

そして，PythonではTrueは1，Falseは0として扱われるため，sum()を適用することでTrueの数の累積を取ります．
また，sum()を適用したままではTensor型なので，実際の値を取得するためにitem()を適用します．

こうして，32個中いくつ正解できたか，という正答数が得られるので，これをすべて合算して全テストデータ数で割ることで正答率を算出しています．

学習結果

出力

Epoch:1, train_loss: 1.615, val_loss: 1.549, val_accuracy: 0.9172
Epoch:2, train_loss: 1.528, val_loss: 1.529, val_accuracy: 0.9356
Epoch:3, train_loss: 1.513, val_loss: 1.509, val_accuracy: 0.9561
Epoch:4, train_loss: 1.503, val_loss: 1.501, val_accuracy: 0.9632
Epoch:5, train_loss: 1.499, val_loss: 1.498, val_accuracy: 0.9658
Epoch:6, train_loss: 1.495, val_loss: 1.496, val_accuracy: 0.9686
Epoch:7, train_loss: 1.492, val_loss: 1.494, val_accuracy: 0.97
Epoch:8, train_loss: 1.490, val_loss: 1.493, val_accuracy: 0.971
Epoch:9, train_loss: 1.488, val_loss: 1.491, val_accuracy: 0.9732
Epoch:10, train_loss: 1.486, val_loss: 1.494, val_accuracy: 0.9692

lossグラフ

精度グラフ