马春杰杰 Exit Reader Mode

[mcj]Pytorch单GPU改为多GPU训练

原来是:

use_cuda = (args.gpu >= 0) and torch.cuda.is_available()
device = torch.device("cuda" if use_cuda else "cpu")
model = model.to(device)

这时用的是单GPU,我们需要改为多GPU,需要在model.to(device)之前加一句

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

即可。

本文最后更新于2019年12月25日,已超过 1 年没有更新,如果文章内容或图片资源失效,请留言反馈,我们会及时处理,谢谢!