pytorch多gpu的问题（RuntimeError: Expected tensor for argument #1 ‘input‘ to have the same device as ...）

1. pytorch正常使用多GPU的方式

# 导入需要的库
import torch
import torch.nn as nn

# 设置torch.device
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

# 假定这里直接导入了model
model = torch.load('xxx.pth')

# 利用nn.DataParallel（device_ids设置多gpu的id）
model = nn.DataParallel(model, device_ids = [0, 1])
model.to(device)
model.eval()  # 假设这里要作推理用

# 推理, 假设输入是img
result = model(img.to(device))

注：这里设置device从0号gpu开始，后面的device_ids里也必须要有0，其他id自己设置。

2. 报错：RuntimeError: Expected tensor for argument #1 ‘input’ to have the same device as …

2.1 问题描述

input和weight不再同一个device上，就是输入跟模型的参数不再同一个gpu上，所以报错。

2.2 问题分析

跟模型的定义形式有关。
（1）上图中子模块的定义只是单纯地copy了self.g的操作（不是tensor，即没有继承nn.Module的方法），所以调用多gpu时，模型的参数不会分配到其他的gpu，因此报错。

（2）贴一张“参考资料1”中比较准确的说法，是对上一张图报错的说明：

2.3 实验验证

定义一个简单的mnist模型（可以正常调用多gpu），如下图所示：

再模拟一个调用多gpu会报错的模型，把fc1变成fc3，如下图所示：
果然跟预想的一样报错了，不过报错信息是（-_-|||）：
RuntimeError: arguments are located on different GPUs at /pytorch/aten/src/THC/generic/THCTensorMathBlas.cu:255
（好吧，意思差不多，参数在不同的gpu上。。。）

把fully_connect1的定义改成如下形式，就可以正常调用多GPU:

就先写到这里把。。

3. 总结

1. pytorch想正常地调用多GPU，需要定义好模型的形式，可以都定义成如下形式：
   
   子模块也定义成这样，然后import进去，给Net调用。

2. pytorch你咋这么坑呢！！！

参考资料：

1. https://github.com/pytorch/pytorch/issues/8637
2. https://blog.csdn.net/senius/article/details/96599955

结束。