pytorch 多GPU处理过程

多GPU的处理机制:

使用多GPU时,pytorch的处理逻辑是:

1.在各个GPU上初始化模型。

2.前向传播时,把batch分配到各个GPU上进行计算。

3.得到的输出在主GPU上进行汇总,计算loss并反向传播,更新主GPU上的权值。

4.把主GPU上的模型复制到其它GPU上。