赞
踩
本文主要叙述如何将一个在CPU的设备上运行的Pytorch代码移植到一台服务器上,使它能够在GPU上运行。
本文大部分内容参考莫凡Pytorch教程: GPU加速运算,并有删改和增补。
转换过程可以大概分如下几步吧:
1(可选). 验证GPU可用性及其块数:
- torch.cuda.is_available() # 是否有已经配置好可以使用的GPU (若True则有)
-
- torch.cuda.device_count() # 可用GPU块数
Output:
- True
- 8
2. 更改网络,可以理解为将网络放入GPU:
- class CNN(nn.Module):
- def __init__(self):
- super(CNN,self).__init__()
- self.conv1 = nn.Sequential(
- ......
- )
-
- ......
-
- self.out = nn.Linear(Liner_input,2)
-
- ......
-
- def forward(self,x):
- x = self.conv1(x)
- ......
- output = self.out(x)
- return output,x
-
- cnn = CNN()
-
- # 更改,.cuda()表示将本存储到CPU的网络及其参数存储到GPU!
- cnn.cuda()
3. 更改输出数据(如向量/矩阵/张量):
- for epoch in range(EPOCH):
- epoch_loss = 0.
- for i, data in enumerate(train_loader2):
- image = data['image'] # data是字典,我们需要改的是其中的image
-
- #############更改!!!##################
- image = Variable(image).float().cuda()
- ############################################
-
- label = inputs['label']
- #############更改!!!##################
- label = Variable(label).type(torch.LongTensor).cuda()
- ############################################
- label = label.resize(BATCH_SIZE)
- output = cnn(image)[0]
- loss = loss_func(output, label) # cross entropy loss
- optimizer.zero_grad() # clear gradients for this training step
- loss.backward() # backpropagation, compute gradients
- optimizer.step()
- ... ...
上述代码就将网络输出的Tensor存储到了GPU中!
4. 更改其他CPU与GPU冲突的地方
有些函数必要在GPU上完成,例如将Tensor转换为Numpy,就要使用data.cpu().numpy(),其中data是GPU上的Tensor。
若直接使用data.numpy()则会报错。除此之外,plot等也需要在CPU中完成。如果不是很清楚哪里要改的话可以先不改,等到程序报错了,再哪里错了改哪里,效率会更高。例如:
- ... ...
- #################################################
- pred_y = torch.max(test_train_output, 1)[1].data.cpu().numpy()
-
- accuracy = float((pred_y == label.cpu().numpy()).astype(int).sum()) / float(len(label.cpu().numpy()))
假如不加.cpu()便会报错,此时再改即可。
5. (可选)切换GPU_id
若要切换GPU_id,只需要在开始处加入下述代码即可。
- id = 2
- torch.cuda.set_device(id) # id=0, 1, 2 等
通过以上5步就可以把CPU上的代码转换到GPU上运行啦!
Copyright © 2003-2013 www.wpsshop.cn 版权所有,并保留所有权利。