Pytorch 多 GPU 并行处理机制

Pytorch 的多 GPU 处理接口是 torch.nn.DataParallel(module, device_ids),其中 module 参数是所要执行的模型,而 device_ids 则是指定并行的 GPU id 列表。

而其并行处理机制是,首先将模型加载到主 GPU 上,然后再将模型复制到各个指定的从 GPU 中,然后将输入数据按 batch 维度进行划分,具体来说就是每个 GPU 分配到的数据 batch 数量是总输入数据的 batch 除以指定 GPU 个数。每个 GPU 将针对各自的输入数据独立进行 forward 计算,最后将各个 GPU 的 loss 进行求和,再用反向传播更新单个 GPU 上的模型参数,再将更新后的模型参数复制到剩余指定的 GPU 中,这样就完成了一次迭代计算。所以该接口还要求输入数据的 batch 数量要不小于所指定的 GPU 数量。

这里有两点需要注意:

  1. 主 GPU 默认情况下是 0 号 GPU,也可以通过 torch.cuda.set_device(id) 来手动更改默认 GPU。
  2. 提供的多 GPU 并行列表中需要包含有主 GPU。
©著作权归作者所有,转载或内容合作请联系作者
平台声明:文章内容(如有图片或视频亦包括在内)由作者上传并发布,文章内容仅代表作者本人观点,简书系信息发布平台,仅提供信息存储服务。

推荐阅读更多精彩内容

  • 我觉得我写东西到了一个瓶颈期。我开始很在意写作的工具,写作的环境,甚至是写作的格式。以至于忽略了写作的本质。我开始...
    壹贰叁肆魏碧智阅读 458评论 0 1
  • 悲伤是高贵的感情
    来吧1234阅读 284评论 0 0
  • 昨夜做了一个梦↡↡↡ 爸爸妈妈离婚后,爸爸娶了一个叫玲兰的女人,其实我很喜欢她,从第一次见到她,就很想亲近她,那种...
    弋爱的名阅读 350评论 0 0