PyTorch - 使用 GPU 加速複雜的 model 訓練

在往後的練習，NN 越來越 “Deep” ，且使用的 model 越來越複雜(CNN, RNN, ResNet,…)，總是使用CPU做訓練會相當耗時，所以要先來練習使用GPU，後面再練習其他 model 訓練時會比較輕鬆。

使用 GPU 前，在安裝 PyTorch 時，是否有選擇正確? 以下是我安裝的規格。

接著我們利用上一篇 CNN MNIST 練習來使用 GPU 加速。

在使用GPU加速幾個要注意的地方，先列出來，方便以後使用時可以快速查找：
1.model  model.cuda() or model.to(device) 兩種方法(後續都使用.to(device))
2. Variable 容器後面 ex. train = Variable(images.view(input_shape))  train = Variable(images.view(input_shape)).to(device)
3. predicted 計算後面 ex. predicted = torch.max(outputs.data, 1)[1]  predicted = torch.max(outputs.data, 1)[1].to(device)

將上一篇 CNN MNIST 代碼修改一下:

首先在1. Import Libraries 後面加入

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)

print 出 cuda:0 即為 GPU device 被偵測可使用

cuda:0

接著，在3. 建立模型 的地方修改如下:

model = CNN_Model()
print(model)
model.to(device) #修改model，使用GPU
optimizer = torch.optim.Adam(model.parameters(), lr=LR) # optimize all cnn parameters
loss_func = nn.CrossEntropyLoss() # the target label is not one-hotted
input_shape = (-1,1,28,28)

最後，4. 訓練模型 中 def fit_model():函式內有 Variable 跟 predicted 計算後面都加上 .to(device)

挑出以下這幾行修改即可:

train = Variable(images.view(input_shape)).to(device) #使用GPU
labels = Variable(labels).to(device) #使用GPU
predicted = torch.max(outputs.data, 1)[1].to(device) #使用GPU
test = Variable(images.view(input_shape)).to(device) #使用GPU
labels = Variable(labels).to(device) #使用GPU
predicted = torch.max(outputs.data, 1)[1].to(device) #使用GPU

完整代碼放在 my Github，可以直接下載來跑看看。

使用GPU 跑了約 70 秒 ，只使用CPU的話大概跑了 542 秒，差了約7倍的時間，之後 model 更複雜所差距的時間會更大。