Multi-GPU layers

class torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)[source]

在模块级别上实现数据并行。

此容器通过将mini-batch划分到不同的设备上来实现给定module的并行。在forward过程中，module会在每个设备上都复制一遍，每个副本都会处理部分输入。在backward过程中，副本上的梯度会累加到原始module上。

batch的大小应该大于所使用的GPU的数量。还应当是GPU个数的整数倍，这样划分出来的每一块都会有相同的样本数量。

请看: Use nn.DataParallel instead of multiprocessing

除了Tensor，任何位置参数和关键字参数都可以传到DataParallel中。所有的变量会通过指定的dim来划分（默认值为0）。原始类型将会被广播，但是所有的其它类型都会被浅复制。所以如果在模型的forward过程中写入的话，将会被损坏。

参数说明：

• module – 要被并行的module
• device_ids – CUDA设备，默认为所有设备。
• output_device – 输出设备（默认为device_ids[0]）

例子：

net = torch.nn.DataParallel(model, device_ids=[0, 1, 2]) output = net(input_var)