Pytorch多卡训练踩坑记录——inputs on different devices

pytorch中设置多卡训练时，操作比较简便，只要定义了模型，然后加上如下一行指令就可以。

model = nn.DataParallel(model).cuda()

但是，在pytorch中进行多卡训练还是会遇到一些其他的问题。执行的上述指令后，pytorch会自动在每块卡上都复制一份模型，同时将input batch size等分。比如设置了8张GPU训练，那么每块卡上的batch size为所设置的input batch size的1/8，每块卡上也都复制了相同的一份模型，进行相同的计算。

但是，若在模型初始化时从外部传递参数，pytorch则不会复制8份，一般会默认储存在0号GPU。如果这些变量只存储和读取，不参与运算，则不会有什么问题，但是当这些参数需要参与到网络计算时（例如点乘），这时便会出现“inputs on different devices”的问题。

例如出现如下问题：

RuntimeError: binary_op(): expected both inputs to be on same device, but input a is on cuda:1 and input b is on cuda:0

以上是进行两个tensor的点对点相乘操作，其中一个tensor是模型初始化时的外部传参，默认储存在GPU 0，另一个tensor是网络的中间变量，每张GPU上均有一份，这样，在与除GPU 0 以外的其他卡上的tensor进行相乘时，就会发生以上报错。

 

解决方案：torch.nn.Parameter

最简单易行的解决方案便是通过torch.nn.Parameter的方式将所传参数变成网络的模块参数，即成为网络的一部分，这时pytorch会自动将其在每块GPU上都复制一份，可解决上述GPU不匹配问题。

例如模型初始化时传入的参数为tensor_input，则在module中可进行如下操作：

self.tensor_input = torch.nn.parameter(tensor_input)

以上方案亲测可用，若后期发现更优方案，会再来进行补充。