ローリングコンバットピッチなう!

AIとか仮想化とかペーパークラフトとか

【PyTorch始めました】PyTorchで書いたMNIST用のスクリプトをGPU化

[technology]PyTorchでGPUを使う


以下のエントリーの続きです。Chainer使いがPyTorchに移行するための備忘録です。
rc30-popo.hatenablog.com
rc30-popo.hatenablog.com
rc30-popo.hatenablog.com

PyTorchでGPUを使う場合(NVIDIAのドライバーやCUDAはインストール済みの前提で)、トレーニング前にモデルとデータ(トレーニングデータ、教師データ)をto('cuda')でGPUに転送するのみです。
推論結果を確認する場合は、モデルからforward()された結果に対してto('cpu')もしくはcpu()でCPUにデータを戻します。

下記のコードでは一応torch.cuda.is_avaiable()でGPU使用可否をチェックし、GPU(CUDA IF)が無ければCPUを使う様にしてあります。

また前回までのサンプルでは、学習結果のテスト時にテスト用データを1件づつ推論させていて効率が悪いので、纏めて推論させて、その後に結果を照合する様に変更しました。

GPU化したコード:

# -*- Coding: utf-8 -*-

# Numpy
import numpy as np
# Chainer
# ★Chainerのdatasetsからmnistを引っ張るためchainerは残す
import chainer
# ★Pytorchをインポート
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim

print('MNIST Deep Learning sample by PyTorch')
print(' PyTorch Version: ',torch.__version__)

# GPU availability check
# GPUは使用可能ならdeviceを'cuda'に設定、GPU無し環境では'cpu'を使用
if torch.cuda.is_available():
    device = 'cuda'
else:
    device = 'cpu'
print(' Device: ',device)

train, test = chainer.datasets.get_mnist()

x_train = []
x_test  = []
t_train = []
t_test  = []

for line in train:
    x_train.append(line[0])
    t_train.append(int(line[-1]))

for line in test:
    x_test.append(line[0])
    t_test.append(int(line[-1]))

x_train = np.array(x_train,dtype=np.float32)
x_test = np.array(x_test,dtype=np.float32)
# ★PyTorchでは分類クラスはintではなくlong
t_train = np.array(t_train,dtype=np.int64)
t_test = np.array(t_test,dtype=np.int64)

x_train = x_train.reshape(len(x_train),1,28,28)
x_test = x_test.reshape(len(t_test),1,28,28)

        
# DNN class ★モデルは前回と同じ
class DNN(nn.Module):
    def __init__(self):
        super(DNN,self).__init__()
        self.cv1 = nn.Conv2d(1,5,4)
        self.cv2 = nn.Conv2d(5,10,4)
        self.l1 = nn.Linear(10 * 4 * 4,100)
        self.l2 = nn.Linear(100,10)
    def forward(self,x):
        c = F.max_pool2d(F.relu(self.cv1(x)),2)
        c = F.max_pool2d(F.relu(self.cv2(c)),2)
        # Chainerと異なり、全結合層に入力する前にshape変換が必要
        c = c.view(-1,10 * 4 * 4)
        h = F.relu(self.l1(c))
        h = self.l2(h)
        return h

# Create DNN class instance
# ★ to(device)を追加してGPUに転送
model = DNN().to(device)

# Set optimizer
# ★optimizerをセット
optimizer = optim.Adam(model.parameters())

# Number of epochs
n_epoch = 20

# batch size
batch_size = 1000

# Execute Training
for epoch in range(n_epoch):
    sum_loss = 0
    perm = np.random.permutation(60000)
    for i in range(0, 60000, batch_size):
# ★x_trainとy_trainをpytorchのTensorに変換
#   ここで変換は効率悪い気もするけど..元のコードベースを崩さないため
# ★ to(device)を追加してGPUに転送
        x = torch.from_numpy(x_train[perm[i:i+batch_size]]).to(device)
        t = torch.from_numpy(t_train[perm[i:i+batch_size]]).to(device)
        y = model.forward(x)
# ★cleargrads()はzero_grad()に
        model.zero_grad()
# ★softmax_cross_entropy()→cross_entropy()
        loss = F.cross_entropy(y,t)
        loss.backward()
#  ★optimizer.update()はoptimizer.step()に
        optimizer.step()
        sum_loss += loss.data*batch_size

    print("epoch: {}, mean loss: {}".format(epoch, sum_loss/60000))

# Execute Test
cnt = 0
# ★テストデータを纏めて推論するように改善
x = torch.from_numpy(x_test).to(device)
y = model.forward(x)
# ★照合用に推論結果yをCPUに戻し、numpyの行列に変換
y = y.cpu().detach().numpy()
# ★1点づつ答え合わせ
for i in range(10000):
    t = t_test[i]
    yi = np.argmax(y[i])
    if t == yi:
        cnt += 1

# Display Result
print("accuracy: {}".format(cnt/(10000)))

GPU版MNIST実行結果(20epoch):

$ time python mnist_cnn_test_torch_gpu.py
MNIST Deep Learning sample by PyTorch
 PyTorch Version:  1.3.1
 Device:  cuda
epoch: 0, mean loss: 1.8325331211090088
epoch: 1, mean loss: 0.4075491726398468
epoch: 2, mean loss: 0.23716536164283752
epoch: 3, mean loss: 0.18152500689029694
epoch: 4, mean loss: 0.14973215758800507
epoch: 5, mean loss: 0.12893804907798767
epoch: 6, mean loss: 0.11485324054956436
epoch: 7, mean loss: 0.10463400930166245
epoch: 8, mean loss: 0.09558816999197006
epoch: 9, mean loss: 0.08927364647388458
epoch: 10, mean loss: 0.08328855782747269
epoch: 11, mean loss: 0.07930488139390945
epoch: 12, mean loss: 0.07505807280540466
epoch: 13, mean loss: 0.07105353474617004
epoch: 14, mean loss: 0.06764773279428482
epoch: 15, mean loss: 0.06459520757198334
epoch: 16, mean loss: 0.06436379253864288
epoch: 17, mean loss: 0.059575922787189484
epoch: 18, mean loss: 0.058406319469213486
epoch: 19, mean loss: 0.05537751689553261
accuracy: 0.9814

real	0m23.489s
user	0m18.544s
sys	0m5.716s

CPU版MNIST実行結果(20epoch):

$ time python mnist_cnn_test_torch.py 
epoch: 0, mean loss: 1.6286011934280396
epoch: 1, mean loss: 0.4382973611354828
epoch: 2, mean loss: 0.2794622778892517
〜中略〜
epoch: 17, mean loss: 0.06488480418920517
epoch: 18, mean loss: 0.06439197808504105
epoch: 19, mean loss: 0.06153048202395439
accuracy: 0.9814

real	1m38.296s
user	7m17.184s
sys	3m44.128s

GPUで処理時間1/4くらいに短縮されています。