pytorch使用GPU上的方法与心得

方法和心得都写在代码注释里了！

以下的代码中有两个，CPU版本是只能在CPU上运行，没有问题，但在GPU上会出现数据位置不一致的问题；GPU版本在CPU上和GPU上运行都没有问题。

import torch
import torch.nn as nn

'''
把模型放到指定机器上的方法非常简单，以下创建的模型在CPU和GPU上都可以运行，有GPU的时候就会用GPU：
假定我们创建一个模型my_model，然后
1. 第一步创建device
    Device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

2. 模型并行化，这样如果有多个GPU的情况可以实现多GPU并行
    parallel_model = nn.DataParallel(my_model)  # 将模型并行化，以便于使用所有的GPU

3. 将并行化的模型用to方法发送到第一步创建的设备上，会返回一个新的模型对象
    my_model = parallel_model.to(Device)

4. 输入模型的数据也要发送到相同的设备上
    data = data.to(Device)
    label = label.to(Device)
    
5. 此时模型和数据都是在Device指定的设备上
   my_model(data,label)就可以正常运行了
'''
#CPU版本，在模型的forward中创建了新的模块，导致将模型发送到GPU时，forward中的模块是不会被发送到GPU上，从而数据与模块不在同一个地方而无法运行
class CPUModel(nn.Module):
    def __init__(self):
        super(CPUModel, self).__init__()
        self.lieaner1 = nn.Linear(8,4)

    def forward(self,data):
        l1 = self.lieaner1(data)
        print('++ l1:',l1.device)  #在GPU上时输出cuda:0,在cpu上运行时输出cpu,因为self.linear1是模型创建之后用to发送到指定设备了
        self.lieaner2 = nn.Linear(8,4)
        l2 = self.lieaner2(data)  #在CPU上运行没问题，但在GPU上运行时用出错，因为data是在GPU上，但self.lieaner2是在模型已经发送到设备之后创建的，因此数据不在同一个地方而无法运行
        print('++ l2:',l2.device)

#GPU版本，在模型构造时创建好所需要所有模块，在forward中只作运算
class GPUModel(nn.Module):
    def __init__(self):
        super(GPUModel, self).__init__()
        self.lieaner1 = nn.Linear(8,4)
        self.lieaner2 = nn.Linear(8, 4)

    def forward(self,data):
        l1 = self.lieaner1(data)
        print('++ l1:',l1.device)  #在GPU上时输出cuda:0,在cpu上运行时输出cpu
        l2 = self.lieaner2(data)   #在GPU上时输出cuda:0,在cpu上运行时输出cpu
        print('++ l2:',l2.device)

if __name__=='__main__':
    Device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    data = torch.randn(16, 5, 8)
    data = data.to(Device)

    my_cpumodel = CPUModel()
    parallel_model = nn.DataParallel(my_cpumodel)  # 将模型并行化，以便于使用所有的GPU
    my_cpumodel = parallel_model.to(Device)
    try:
        my_cpumodel(data)
    except Exception as e:
        print(e)

    print('+'*50)
    my_gpumodel = GPUModel()
    parallel_model = nn.DataParallel(my_gpumodel)  # 将模型并行化，以便于使用所有的GPU
    my_gpumodel = parallel_model.to(Device)
    my_gpumodel(data)

CPU上运行结果截图：

GPU上运行结果截图（本例的服务器上3个GPU）：可以看到CPU版本的代码不能正常输出"l2"