东阳的学习记录，坚持就是胜利！

nn.DataParallel

torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)

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

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

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

参数

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

警告

建议使用DistributedDataParallel而不是此类进行多GPU训练，即使只有一个节点也是如此。 请参阅：使用nn.parallel.DistributedDataParallel而不是多处理或nn.DataParallel和分布式数据并行

使用

使用下面的代码就可以实现在多块GPU上训练：

device_ids = [0, 1, 2, 3]
if cuda:net = net.cuda(device_ids[0])net = nn.DataParallel(net, device_ids=device_ids)

nn.parallel.DistributedDataParallel

torch.nn.parallel.DistributedDataParallel(module, device_ids=None, output_device=None, dim=0, broadcast_buffers=True, process_group=None, bucket_cap_mb=25, find_unused_parameters=False, check_reduction=False)

在模块级别基于torch.distributed包实现分布式数据并行性。

该容器通过按批处理维度分块指定设备之间的输入来并行化给定模块的应用程序。 该模块在每台机器和每台设备上复制，并且每个这样的副本处理一部分输入。 在向后传递过程中，将平均每个节点的梯度。

批处理大小应大于本地使用的GPU数量。

例子

>>> torch.distributed.init_process_group(backend='nccl', world_size=4, init_method='...')
>>> model = DistributedDataParallel(model, device_ids=[i], output_device=i)

>>> import torch.distributed.autograd as dist_autograd
>>> from torch.nn.parallel import DistributedDataParallel as DDP
>>> from torch import optim
>>> from torch.distributed.optim import DistributedOptimizer
>>> from torch.distributed.rpc import RRef
>>>
>>> t1 = torch.rand((3, 3), requires_grad=True)
>>> t2 = torch.rand((3, 3), requires_grad=True)
>>> rref = rpc.remote("worker1", torch.add, args=(t1, t2))
>>> ddp_model = DDP(my_model)
>>>
>>> # Setup optimizer
>>> optimizer_params = [rref]
>>> for param in ddp_model.parameters():
>>>     optimizer_params.append(RRef(param))
>>>
>>> dist_optim = DistributedOptimizer(
>>>     optim.SGD,
>>>     optimizer_params,
>>>     lr=0.05,
>>> )
>>>
>>> with dist_autograd.context() as context_id:
>>>     pred = ddp_model(rref.to_here())
>>>     loss = loss_func(pred, loss)
>>>     dist_autograd.backward(context_id, loss)
>>>     dist_optim.step()