每天讲解一点PyTorch 1８多卡训练torch.nn.DataParallel
Posted knowform
tags:
篇首语：本文由小常识网(cha138.com)小编为大家整理，主要介绍了每天讲解一点PyTorch 1８多卡训练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)
)
以上是关于每天讲解一点PyTorch 1８多卡训练torch.nn.DataParallel的主要内容，如果未能解决你的问题，请参考以下文章
每天讲解一点PyTorch 18多卡训练torch.nn.DataParallel