Pytorch CPU版本转换为GPU版本

本文主要叙述如何将一个在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上运行啦！