参考:
https://pytorch.org/docs/stable/nn.html
https://github.com/apachecn/pytorch-doc-zh/blob/master/docs/1.0/blitz_data_parallel_tutorial.md
一、 torch.nn.DataParallel
torch.nn.
DataParallel
(module, device_ids=None, output_device=None, dim=0)
在正向传递中,模块在每个设备上复制,每个副本处理一部分输入。在向后传递期间,来自每个副本的渐变被加到原始模块中。
- module:需要并行处理的模型
- device_ids:并行处理的设备,默认使用所有的cuda
- output_device:输出的位置,默认输出到cuda:0
例子:
>>> net = torch.nn.DataParallel(model, device_ids=[0, 1, 2])
>>> output = net(input_var) # input_var can be on any device, including