pytorch 使用单个GPU与多个GPU进行训练与测试

pytorch 使用单个GPU与多个GPU进行训练与测试

使用多个GPU的原理就是通过上面这句代码将model在每个GPU上分别保存一份，然后对model的输入tensor进行自动的分割，每个GPU计算tensor的一部分，这样就能实现计算量的平均分配。在每个model计算完成之后，DataParallel将这些结果进行收集和融合，之后再将结果返回。

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") # 判断是否有GPU，没有的话使用CPU
model = Model()  #使用模型
if torch.cuda.device_count() > 1:
    model = nn.DataParallel(model)
model.to(device)  #模型送入GPU种
mytensor = my_tensor.to(device)  # 数据送入GPU