[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くらいに短縮されています。