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

Posted cv.exp

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了每天讲解一点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)





以上是关于每天讲解一点PyTorch 18多卡训练torch.nn.DataParallel的主要内容,如果未能解决你的问题,请参考以下文章

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

PyTorch多卡训练

每天讲解一点PyTorch torch.matmul

每天讲解一点PyTorch torch.matmul

每天讲解一点PyTorch 15model.load_state_dict torch.load torch.save

每天讲解一点PyTorch 15model.load_state_dict torch.load torch.save