Pytorch多GPU并行处理

原理见https://zhuanlan.zhihu.com/p/72939003

单机多卡:

DataParallel Model 

device_ids = [0,1,2,3,4,5,6,7]

#前向推理时使用并行
model = model.cuda(device_ids[0])    #设定主卡
model = nn.DataParallel(model, device_ids=device_ids)

#反向传播时使用并行
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=0.9, weight_decay=0.001)
optimizer = nn.DataParallel(optimizer, device_ids=device_ids)

多机多卡:

DataParallelCriterion

 

参考

https://blog.csdn.net/weixin_40087578/article/details/87186613

https://www.cnblogs.com/ranjiewen/p/10113532.html(还讲了梯度积累)

你可能感兴趣的:(Pytorch多GPU并行处理)