Pytorch——多卡GUP训练原理(torch.nn.DataParallel)
本文将记录如何使用单机多卡GPU的方式进行训练,主要是采用DP模式(DDP模式一般用于多机多卡训练)。
1、DP模式基本原理
DP模型进行单机多卡训练基本步骤如下:
(1)将模型复制到各个GPU中,并将一个batch的数据划分成mini_batch并分发给给个GPU
(2)各个GPU独自完成mini_batch的前向传播,并把获得的output传递给GPU_0(主GPU)
(3)GPU_0整合各个GPU传递过来的output,并计算loss。此时GPU_0可以对这些个loss进行一些聚合操作。
(4)GPU_0归并loss之后,并进行后向传播以及梯度下降从而完成模型参数的更新(此时只有GPU_0上的模型参数得到了更新),GPU_0将更新好的模型参数又传递给其余GPU
以上就是DP模式下多卡GPU进行训练的方式。其实可以看到GPU_0不仅承担了前向传播的任务,还承担了收集loss,并进行梯度下降。因此在使用DP模式进行单机多卡GPU训练的时候会有一张卡的显存利用会比其他卡更多,那就是你设置的GPU_0。
2、Pytorch如何进行单机多卡训练
Pytorch中可以直接使用torch.nn.DataParallel(model)就可以完成了。具体的代码模块见下图。
1 2 3 4 5 | import torch.nn as nn import os os.environ[ "CUDA_VISIBLE_DEVICES" ] = '2,3,1' # 指定该程序可以识别的物理GPU编号 model = nn.DataParallel(model) # 在执行该语句之前最好加上model.cuda(),保证你的模型存在GPU上即可 |
第三行代码建议一定要加上,这一行代码是为了指定有哪些物理GPU是对该程序可见的,并且此处的顺序将会是该程序在执行的时候识别GPU的顺序(例如此处该程序的GPU_0=物理GPU2号,GPU_1=物理GPU3号,GPU_2=物理GPU1号),所以物理GPU2号将会作为该程序的主要GPU(用于计算loss,梯度下降,参数更新)。
对于数据,我们只需要按照平常的方式使用.cuda()放置在GPU上即可,内部batch的拆分已经被封装在了DataParallel模块中。要注意的是,由于我们的model被nn.DataParallel()包裹住了,所以如果想要储存模型的参数,需要使用model.module.state_dict()的方式才能取出(不能直接是model.state_dict())。
参考网站:
pytorch多gpu并行训练 - 知乎 (zhihu.com)
解决了PyTorch 使用torch.nn.DataParallel 进行多GPU训练的一个BUG:模型(参数)和数据不在相同设备上_senius的博客-CSDN博客
pytorch遇到的bug记录—2_羊藤枝的博客-CSDN博客
【分布式训练】单机多卡的正确打开方式(三):PyTorch - 知乎 (zhihu.com)
分布式训练 - 单机多卡(DP和DDP)_love1005lin的博客-CSDN博客
【推荐】编程新体验,更懂你的AI,立即体验豆包MarsCode编程助手
【推荐】凌霞软件回馈社区,博客园 & 1Panel & Halo 联合会员上线
【推荐】抖音旗下AI助手豆包,你的智能百科全书,全免费不限次数
【推荐】博客园社区专享云产品让利特惠,阿里云新客6.5折上折
【推荐】轻量又高性能的 SSH 工具 IShell:AI 加持,快人一步
· 浏览器原生「磁吸」效果!Anchor Positioning 锚点定位神器解析
· 没有源码,如何修改代码逻辑?
· 一个奇形怪状的面试题:Bean中的CHM要不要加volatile?
· [.NET]调用本地 Deepseek 模型
· 一个费力不讨好的项目,让我损失了近一半的绩效!
· PowerShell开发游戏 · 打蜜蜂
· 在鹅厂做java开发是什么体验
· 百万级群聊的设计实践
· WPF到Web的无缝过渡:英雄联盟客户端的OpenSilver迁移实战
· 永远不要相信用户的输入:从 SQL 注入攻防看输入验证的重要性