Pytorch的nn.DataParallel 多 GPU 训练 只用到单个GPU/只用到一个GPU/不起作用

最近尝试使用 pytorch 的 nn.DataParallel 进行多 GPU 训练。按照官方文档修改代码后发现只使用了一块 GPU，最后经过查阅论坛，找到了原因。这里总结一下，希望能帮到大家。

顺便一提 pytorch 官网推荐使用 DistributedDataParallel 进行多 GPU 训练，而这篇博客只是分享一些问题的原因。

nn.DataParallel 原理简述：在前向传播过程中，将输入的 batch 平均分配到用户指定的 GPU 中，每个 GPU 中有一份 module 的副本，每个 GPU 分别处理分配到的 batch，最后将计算后的梯度汇总。

只用到单个GPU/只用到一个GPU/不起作用可能的原因：

1. batch_size 过小
如果你的 batch_size 小于你的 GPU 数量，那么 DataParallel 不会起作用。
2.未使用 module 的 forward 方法进行前向传播
当你像下面的方式使用 DataParallel 时:

model = nn.DataParallel(model)

那么，只有在你调用 module 时才会将 batch 分配到其它的GPU，即：

model(input)# 这里是调用了 model.forward() 函数

如果你使用自定义方法计算的话，则不会将 batch 分配到其它的GPU，如：

model.my_forward_method(input)

3.forward()参数不正确
此原因不确定，仅供参考
如果forward函数的输入是 tensor ，那么该 tensor 会按照你指定的维度（默认是0）分配到其它 GPU 中，而如果输入是其它类型，则会直接复制到其它的GPU中。
例如 tensor size : (batch_size,3,128,128)