当前位置:   article > 正文

Pytorch CPU版本转换为GPU版本_pytorch怎么把cpu改成gpu在服务器上跑

pytorch怎么把cpu改成gpu在服务器上跑

本文主要叙述如何将一个在CPU的设备上运行的Pytorch代码移植到一台服务器上,使它能够在GPU上运行。

本文大部分内容参考莫凡Pytorch教程: GPU加速运算,并有删改和增补。

 

转换过程可以大概分如下几步吧:

1(可选). 验证GPU可用性及其块数:

  1. torch.cuda.is_available() # 是否有已经配置好可以使用的GPU (若True则有)
  2. torch.cuda.device_count() # 可用GPU块数

Output: 

  1. True
  2. 8

2. 更改网络,可以理解为将网络放入GPU:

  1. class CNN(nn.Module):
  2. def __init__(self):
  3. super(CNN,self).__init__()
  4. self.conv1 = nn.Sequential(
  5. ......
  6. )
  7. ......
  8. self.out = nn.Linear(Liner_input,2)
  9. ......
  10. def forward(self,x):
  11. x = self.conv1(x)
  12. ......
  13. output = self.out(x)
  14. return output,x
  15. cnn = CNN()
  16. # 更改,.cuda()表示将本存储到CPU的网络及其参数存储到GPU!
  17. cnn.cuda()

 

3. 更改输出数据(如向量/矩阵/张量):

  1. for epoch in range(EPOCH):
  2. epoch_loss = 0.
  3. for i, data in enumerate(train_loader2):
  4. image = data['image'] # data是字典,我们需要改的是其中的image
  5. #############更改!!!##################
  6. image = Variable(image).float().cuda()
  7. ############################################
  8. label = inputs['label']
  9. #############更改!!!##################
  10. label = Variable(label).type(torch.LongTensor).cuda()
  11. ############################################
  12. label = label.resize(BATCH_SIZE)
  13. output = cnn(image)[0]
  14. loss = loss_func(output, label) # cross entropy loss
  15. optimizer.zero_grad() # clear gradients for this training step
  16. loss.backward() # backpropagation, compute gradients
  17. optimizer.step()
  18. ... ...

上述代码就将网络输出的Tensor存储到了GPU中!

 

4. 更改其他CPU与GPU冲突的地方

有些函数必要在GPU上完成,例如将Tensor转换为Numpy,就要使用data.cpu().numpy(),其中data是GPU上的Tensor。

若直接使用data.numpy()则会报错。除此之外,plot等也需要在CPU中完成。如果不是很清楚哪里要改的话可以先不改,等到程序报错了,再哪里错了改哪里,效率会更高。例如:

  1. ... ...
  2. #################################################
  3. pred_y = torch.max(test_train_output, 1)[1].data.cpu().numpy()
  4. accuracy = float((pred_y == label.cpu().numpy()).astype(int).sum()) / float(len(label.cpu().numpy()))

假如不加.cpu()便会报错,此时再改即可。

5. (可选)切换GPU_id

若要切换GPU_id,只需要在开始处加入下述代码即可。

  1. id = 2
  2. torch.cuda.set_device(id) # id=0, 1, 2 等

 

通过以上5步就可以把CPU上的代码转换到GPU上运行啦!

 

声明:本文内容由网友自发贡献,不代表【wpsshop博客】立场,版权归原作者所有,本站不承担相应法律责任。如您发现有侵权的内容,请联系我们。转载请注明出处:https://www.wpsshop.cn/w/码创造者/article/detail/851491
推荐阅读
相关标签
  

闽ICP备14008679号