目录
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. 总结
-
pytorch想正常地调用多GPU,需要定义好模型的形式,可以都定义成如下形式:
子模块也定义成这样,然后import进去,给Net调用。 -
pytorch你咋这么坑呢!!!
参考资料:
- https://github.com/pytorch/pytorch/issues/8637
- https://blog.csdn.net/senius/article/details/96599955
结束。