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)