如何在CUDA上使用PyTorch進行CIFAR-10？

Question

我正在這個pytorch頁面上關注CIFAR-10 PyTorch教程，並且無法在GPU上運行PyTorch。 代碼與教程中的完全相同。

我得到的錯誤是

Traceback (most recent call last):
  File "(file path)/CIFAR10_tutorial.py", line 116, in <module>
   outputs = net(images)
  File "/usr/local/lib/python3.5/dist-packages/torch/nn/modules/module.py", line 491, in __call__
result = self.forward(*input, **kwargs)
  File "(file path)/CIFAR10_tutorial.py", line 65, in forward
x = self.pool(F.relu(self.conv1(x)).cuda())
  File "/usr/local/lib/python3.5/dist-packages/torch/nn/modules/module.py", line 491, in __call__
result = self.forward(*input, **kwargs)
  File "/usr/local/lib/python3.5/dist-packages/torch/nn/modules/conv.py", line 301, in forward
self.padding, self.dilation, self.groups)

我的CUDA版本是9.0，Pytorch 0.4.0。 我在機器上使用過tensorflow-gpu，所以我知道CUDA設置正確。 按照教程中的建議，我到底應該在哪里使用.cuda（）和.to（device）？

Answer 1

我要留下答案，以防其他人陷入困境。

首先，將Pytorch配置為使用GPU（如果有）

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

然后，在init函數中，通過對NN的每個元素調用.cuda（）來強制轉換為gpu，例如

self.conv1 = nn.Conv2d(3, 24, 5).cuda()
self.pool = nn.MaxPool2d(2, 2).cuda()

如果不確定GPU，請在每個元素上調用.to（device） 。

在步驟的前向我（x）函數中，我做了

x = x.to(device)

創建完網絡對象后，立即將其投射到設備上

net.to(device)

在對所有輸入和標簽進行任何操作之前，應將其全部投射到設備上。

inputs, labels = inputs.to(device), labels.to(device)

我正在跳過編寫整個代碼的過程，因為問題中已經提到了鏈接。 如果似乎有一些對gpu的多余轉換，則它們沒有破壞任何內容。 我可能還會將ipynb與這些更改放在一起。