pytorch 多GPU训练 loss不能正常下降

最新推荐文章于 2024-08-14 09:34:07 发布

zzz1515151

最新推荐文章于 2024-08-14 09:34:07 发布

阅读量6.2k

点赞数 7

CC 4.0 BY-SA版权

文章标签： pytorch 多GPU loss下降异常

本文链接：https://blog.youkuaiyun.com/zzz1515151/article/details/90904941

在使用PyTorch的DataParallel进行多GPU训练时，参照某些博客可能导致loss无法正常下降，因为错误地将优化器用DataParallel包裹。问题在于未正确处理优化器清空梯度的操作。解决方案是避免对优化器使用DataParallel，并按照单卡训练的方式处理权重加载和平衡。加载权重时需注意多卡训练权重的关键字中包含'module'。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

在使用pytorch的DataParallel进行多GPU并行训练过程中，参考了很多博客，其中两个博客的内容有误，会导致多GPU训练时loss不能正常下降，且运行过程中不会有报错，下面做简单的说明。

有问题的博客链接：

https://blog.youkuaiyun.com/daniaokuye/article/details/79110365

https://blog.youkuaiyun.com/qq_19598705/article/details/80396325

以上两个博客的主要问题为将优化器也使用DataParallel包装：

optim_opt = config["optim"]
optimizer = optim.SGD(net.parameters(), lr = optim_opt["lr"], \
                momentum = optim_opt["momentum"], \
                nesterov = optim_opt["nesterov"], \
                weight_decay=optim_opt['weight_decay'])
optimizer = nn.DataParallel(optimizer, device_ids = [0, 1, 2, 3])

使用DataParallel包装后，optimizer的所有操作都需要获取module对象。以上两个博客都忽略了优化器清空梯度时的操作：