pytorch在GPU上运行模型实现并行计算_python gpu并行计算

pytorch在GPU上运行模型十分简单，只需要以下两部：

• model = model.cuda()：将模型的所有参数都转存到GPU上
• input.cuda()：将输入数据放置到GPU上

至于如何在多块GPU上进行并行计算，PyTorch也提供了两个函数，可以实现简单、高效的GPU并行计算。

• nn.parallel.data_parallel(module, inputs, device_ids=None, output_device=None, dim=0, module_kwargs=None)
• class torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)

这两个函数的参数十分相似，其中device_ids参数可以指定在哪些GPU上进行优化，output_device参数可以指定输出到哪块GPU上。

二者唯一的不同在于：前者直接利用多块GPU进行并行计算得到结果，后者返回一个新的module，能够自动在多块GPU上进行并行加速。

方法一

new_model = nn.DataParallel(model, device_ids=[0, 1])
output = new_model(input)
1
2

方法二

output = nn.parallel.data_parallel(new_model, input, device_ids=[0, 1])
1