每天讲解一点PyTorch 【18】多卡训练torch.nn.DataParallel

使用多GPU训练模型：

net = ......
device_ids = [0, 1]
net = torch.nn.DataParallel(net, device_ids=device_ids)

优化器：

optimizer = torch.optim.SGD(net.parameters(), lr=lr)
optimizer = nn.DataParallel(optimizer, device_ids=device_ids)

GPU setting

parser.add_argument(‘–cpu’, action=‘store_true’, help=‘use cpu only’)
parser.add_argument(‘–n_GPUs’, type=int, default=1, help=‘number of GPUs’)
parser.add_argument(‘–gpu_id’, type=int, default=0, help=‘GPU id’)

os.environ[‘CUDA_VISIBLE_DEVICES’]=str(args.gpu_id)
device = torch.device(‘cpu’ if args.cpu else ‘cuda’)

self.feature_loss_module.to(device)

if not args.cpu and args.n_GPUs > 1:
self.feature_loss_module = nn.DataParallel(
self.feature_loss_module, range(args.n_GPUs)
)

![请添加图片描述](https://img-blog.csdnimg.cn/514c2e7481a34b2ca8834a890de65846.png)