0
点赞
收藏
分享

微信扫一扫

pytorch在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)

方法二

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


举报

相关推荐

0 条评论