这里记录用pytorch 多GPU训练 踩过的许多坑 仅针对单服务器多gpu 数据并行 而不是 多机器分布式训练
一、官方思路包装模型
这是pytorch 官方的原理图 按照这个官方的原理图 修改应该参照
https://blog.csdn.net/qq_19598705/article/details/80396325
上文也用dataParallel 包装了optimizer, 对照官方原理图中第二行第二个,将梯度分发出去,将每个模型上的梯度更新(第二行第三个),然后再将更新完梯度的模型参数合并到主gpu(第二行最后一个步骤)
其实完全没必要,因为每次前向传播的时候都会分发模型,用不着反向传播时将梯度loss分发到各个GPU,单独计算梯度,再合并模型。可以就在主GPU 上根据总loss 更新模型的梯度,不用再同步其他GPU上的模型,因为前向传播的时候会分发模型。
所以 上述链接里 不用 dataParallel 包装 optimizer。
DataParallel并行计算只存在在前向传播
总结步骤:
import os
import torch
args.gpu_id=&#