PyTorchでMNISTする

(2019-01-19)

PyTorchはFacebookによるOSSの機械学習フレームワーク。 TensorFlow(v1)よりも簡単に使うことができる。 TensorFlow 2.0ではPyTorchのようにDefine-by-runなeager executionがデフォルトになるのに加え、パッケージも整理されるようなのでいくらか近くなると思われる。

使い方

インストール

Colabで動かす。まずpipでインストール

!pip install torch torchvision

autograd(自動微分)

Tensorは自身が作成された関数の参照.grad_fnを持ち、backward()が呼ばれるとbackpropしてrequires_grad=TrueなTensorの勾配を自動で計算し.gradに入れてくれる。

MLPと誤差逆伝搬法(Backpropagation) - sambaiz-net

import torch
x = torch.randn(4, 4)
y = torch.randn(4, 1)
w = torch.randn(4, 1, requires_grad=True)
b = torch.randn(1, requires_grad=True)
y_pred = torch.matmul(x, w) + b
loss = (y_pred - y).pow(2).sum()

print(x.grad, w.grad) # None None 
loss.backward()
print(x.grad, w.grad) # None tensor([...])

with torch.no_grad():
    y_eval = torch.matmul(x, w) + b
print(y_eval.requires_grad) # False

Module

nnパッケージにLinearやConv2dといったModuleが実装されていて、次のように呼び出すとforward()が 呼ばれ順伝播する。

model = nn.Sequential(
    nn.Conv2d(1,20,5),
    nn.ReLU(),
    nn.Conv2d(20,64,5),
    nn.ReLU()
)
y = model(x)

nn.Moduleを継承すれば自分でModuleを作ることもできる。

import torch.nn as nn
import torch.nn.functional as F

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(1, 20, 5)
        self.conv2 = nn.Conv2d(20, 20, 5)

    def forward(self, x):
       x = F.relu(self.conv1(x))
       return F.relu(self.conv2(x))

train()とeval()で学習/推論のモードが切り替わり Dropoutなどの挙動に影響する。

model.train() 
model.eval()

Optimizer

optimパッケージにOptimizerが実装されていて、step()で勾配をもとにパラメータが更新される。

optimizer = optim.SGD(model.parameters(), lr = 0.01, momentum=0.9)
for input, target in dataset:
    optimizer.zero_grad()
    output = model(input)
    loss = loss_fn(output, target)
    loss.backward()
    optimizer.step()

Save/Load

学習したモデルのパラメータをSave/Loadする。

torch.save(model.state_dict(), PATH)
model.load_state_dict(torch.load(PATH))

torchvision

Datasetやモデルが含まれるパッケージ。

import torchvision
dataset_train = torchvision.datasets.MNIST('~/mnist', train=True, download=True)
!ls ~/mnist/processed # => test.pt  training.pt

%matplotlib inline
from PIL import Image
import matplotlib.pyplot as plt
import numpy as np
import tensorflow as tf
plt.imshow(np.asarray(dataset_train[0][0]))
print(dataset_train[0][1]) # tensor(5)

MNIST Datasetの内容

DatasetをDataLoader に渡してやるとシャッフルなどしてイテレーションしてくれるが、MNISTのデータはPIL Imageで入っているのでそのまま渡すと TypeError: batch must contain tensors, numbers, dicts or lists; found <class 'PIL.Image.Image'>のエラーになる。 そこでtransforms.ToTesnor()でTensorに変換してやる必要がある。

import torch
import torchvision
dataset_train = torchvision.datasets.MNIST('~/mnist', train=True, download=True, transform=torchvision.transforms.ToTensor())
dataloader_train = torch.utils.data.DataLoader(dataset_train,
                                          batch_size=4,
                                          shuffle=True,
                                          num_workers=0)
for x, t in dataloader_train:
    print(x.shape, t.shape) # torch.Size([4, 1, 28, 28]) torch.Size([4])
    break

GPUを使う

GPUで計算させるにはto()でdeviceを明示する必要がある。

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

MNISTのモデル作成

MNISTのモデルを作りtorchvisionのDatasetで学習させてみる。

Dataset

まずDatasetを準備する。

from torchvision import datasets, transforms, models
from torch import nn, optim, utils, device as device_, cuda
import torch.nn.functional as F
import numpy as np
from sklearn import metrics

dataset_train = datasets.MNIST(
    '~/mnist', 
    train=True, 
    download=True, 
    transform=transforms.ToTensor())
dataset_valid = datasets.MNIST(
    '~/mnist', 
    train=False, 
    download=True, 
    transform=transforms.ToTensor())

dataloader_train = utils.data.DataLoader(dataset_train,
                                          batch_size=1000,
                                          shuffle=True,
                                          num_workers=4)
dataloader_valid = utils.data.DataLoader(dataset_valid,
                                          batch_size=1000,
                                          shuffle=True,
                                          num_workers=4)

モデル

モデルとOptimizer、損失関数。

lr = 0.01

device = device_("cuda" if cuda.is_available() else "cpu")

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.conv1 = nn.Conv2d(1, 64, 5) # -> 24x24
        self.pool1 = nn.MaxPool2d(2) # -> 12x12
        self.conv2 = nn.Conv2d(64, 128, 5) # -> 8x8
        self.dropout = nn.Dropout(p=0.4)
        self.dense = nn.Linear(128 * 8 * 8, 10)

    def forward(self, x):
       x = F.relu(self.conv1(x))
       x = self.pool1(x)
       x = F.relu(self.conv2(x))
       x = self.dropout(x)
       x = x.view(x.size(0), -1) # Flatten
       return F.relu(self.dense(x))
    
model = Model().to(device)
optimizer = optim.SGD(model.parameters(), lr=lr)
criterion = nn.CrossEntropyLoss()

学習

trainモードにして学習していく。

model.train()
for i in range(20):
  print(i)
  for x, t in dataloader_train:
      x = x.to(device)
      t = t.to(device)
      model.zero_grad()
      y = model(x)
      loss = criterion(y, t)
      loss.backward()
      optimizer.step()

評価

evalモードにしてテストする。

model.eval()
labels = []
preds = []
losses = []
for x, t in dataloader_valid:
    x = x.to(device)
    t = t.to(device)
    labels.extend(t.tolist())
    y = model(x)
    loss = criterion(y, t)
    losses.append(loss.cpu().data)
    pred = y.argmax(1)
    preds.extend(pred.tolist())
print('Loss: {:.3f}, Accuracy: {:.3f}'.format(
    np.mean(losses),
    metrics.accuracy_score(labels, preds, normalize=True)
))