PyTorch 的 nn.DataParallel 是否在每个 GPU 中加载相同的模型?

Does PyTorch's nn.DataParallel load the same model in each GPU?

它似乎(逻辑上)起作用的唯一方法是将模型加载到每个 GPU 中。这意味着当更新权重时,每个 GPU 也需要更新权重,与单个 GPU 相比增加了工作量。这个推理正确吗?

首先,it is advised to use torch.nn.parallel.DistributedDataParallel代替。

您可以检查 torch.nn.DataParallel documentation where the process is described (you can also check source code and dig a little deeper on github, here 是如何执行模块复制的。

大致是这样的:

初始化

所有(或选择的)设备 id 都保存在构造函数和维度中,数据将沿着这些维度分散(几乎总是 0 意味着它将按批次拆分到设备)

前进

forward 运行:

  1. 输入是分散的(张量沿维度,tuplelistdict 浅层复制,其他数据在线程之间共享)。
  2. 如果只有一台设备就returnmodule(*args, **kwargs)
  3. 如果有多个设备,将网络从源机器复制到其他设备(每次都这样做!)
  4. 运行 在每个设备上使用各自的输入转发
  5. 将设备的输出收集到源计算机上的单个源设备(输出串联)。
  6. 运行 其余代码、反向传播、更新源机器上的权重等

源机器默认为cuda:0,但可以选择。权重也在单个 device 上更新,只有批次被分散并收集输出。