做网站爱,开源网站推广,企业网站推广排名,网站数据库连接失败文章目录 1. 基本原理单机多卡训练教程——DP模式 2. Pytorch进行单机多卡训练步骤1. 指定GPU2. 更改模型训练方式3. 更改权重保存方式 摘要#xff1a;多GPU同时训练#xff0c;能够解决单张GPU显存不足问题#xff0c;同时加快模型训练。
1. 基本原理
单机多卡训练教程—… 文章目录 1. 基本原理单机多卡训练教程——DP模式 2. Pytorch进行单机多卡训练步骤1. 指定GPU2. 更改模型训练方式3. 更改权重保存方式 摘要多GPU同时训练能够解决单张GPU显存不足问题同时加快模型训练。
1. 基本原理
单机多卡训练教程——DP模式
(1将模型复制到各个GPU中并将一个batch的数据划分成mini_batch(平均分配 并分发给每个GPU; 注意:这里的batch_size要大于device数。 (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进行单机多卡训练步骤
只需要在你的代码中改三个地方就可实现
1. 指定GPU 如上所示在导入各种库下面使用os.environ[CUDA_VISIBLE_DEVICES]来指定可识别的GPU该语句在程序开始前使用。 代码如下
import torch.nn as nn
import os
os.environ[CUDA_VISIBLE_DEVICES] 2,3,1#指定该程序可以识别的物理GPU编号这里的你主机上的2号GPU就是训练程序中的主GPUO这里最好—定要自己指定你自己可以用的gpu号。2. 更改模型训练方式 平常的模型训练方式只需要model.cuda()语句即可在单机多卡训练中只需要在该语句下面添加一行nn.DataParallel语句即可。 代码如下
model.cuda()
model nn.DataParallel(modeldevise [0,1,2])#在执行该语句之前最好加上model.cuda(),保证你的模型存在GPU上即可3. 更改权重保存方式
对于数据我们只需要按照平常的方式使用.cuda()放置在GPU上即可内部batch的拆分已经被封装在了DataPanallel模块中。要注意的是由于我们的model被nn.DataPanallel()包裹住了所以如果想要储存模型的参数需要使用model.module.state_dict()的方式才能取出不能直接是model.state_dict() 代码如下 使用单机多卡训练的模型权重保存方式torch.save(model.module.state_dict(),fbest.pth) 作为参考将平常的权重保存方式也写上 平常的权重保存方式torch.save(model.state_dict(),fbest.pth)
