模型训练中梯度累积步数(gradient_accumulation_steps)的作用

模型训练中梯度累积步数(gradient_accumulation_steps)的作用

flyfish

在使用训练大模型时,TrainingArguments有一个参数梯度累积步数(gradient_accumulation_steps)

py 复制代码
from transformers import TrainingArguments

梯度累积是一种在训练深度学习模型时用于处理内存限制问题的技术。在每次迭代中,模型的梯度是通过反向传播计算得到的,而梯度累积步数(gradient_accumulation_steps)指定了在执行实际的参数更新之前,要累积多少个小批次(mini - batch)的梯度。

以代码来说gradient_accumulation_steps的作用

py 复制代码
import torch
from torch import nn, optim

# 生成更合理的数据集,假设目标关系是y = 3 * x + 2 加上一些噪声
def generate_dataset(num_samples):
    inputs = torch.randn(num_samples, 10)
    # 根据线性关系生成标签,添加一些随机噪声模拟真实情况
    labels = 3 * inputs.sum(dim=1, keepdim=True) + 2 + torch.randn(num_samples, 1) * 0.5
    return list(zip(inputs, labels))

# 生成数据集,这里生成2000个样本(可根据实际情况调整数据量)
your_dataset = generate_dataset(2000)

# 模型、损失和优化器
model = nn.Linear(10, 1)
# 使用Xavier初始化方法来初始化模型参数,有助于缓解梯度消失和爆炸问题,提升训练效果
nn.init.xavier_uniform_(model.weight)
nn.init.zeros_(model.bias)
criterion = nn.MSELoss()
# 适当调整学习率,这里改为0.1,可根据实际情况进一步微调
optimizer = optim.Adam(model.parameters(), lr=0.1)

# 配置梯度累积步数
gradient_accumulation_steps = 4
global_step = 0

# 模拟训练循环
for epoch in range(20):  # 训练20个周期
    for step, (inputs, labels) in enumerate(torch.utils.data.DataLoader(your_dataset, batch_size=8)):
        
        # 前向传播
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        
        # 反向传播(累积梯度)
        loss.backward()
        
        # 执行梯度更新
        if (step + 1) % gradient_accumulation_steps == 0:
            optimizer.step()
            optimizer.zero_grad()
            global_step += 1
            print(f"更新了模型参数,当前全局步数: {global_step}, 当前损失: {loss.item()}")

解释:

  • batch_size=8:每个梯度计算时,模型会处理 8 张图像。
  • gradient_accumulation_steps=4:表示每次参数更新前累积 4 次梯度。

因此:

  • 每个 step: 处理 8 张图像。
  • 累积 4 个 step: 共处理 8 × 4 = 32 8 \times 4 = 32 8×4=32 张图像。

关键点:

  • 一个 step: 是指一次前向和后向传播(不包含参数更新)。
  • 一次参数更新: 在累积 4 个 step 后,进行一次模型参数更新。

等效有效批次:

有效批次大小 = batch_size × gradient_accumulation_steps

即: 8 × 4 = 32 8 \times 4 = 32 8×4=32。

这意味着,即使显存有限,模型仍然能以有效批次大小 32 的方式进行训练

相关推荐
LLM精进之路9 小时前
上海AI实验室突破扩散模型!GetMesh融合点云与三平面,重塑3D内容创作
人工智能·深度学习·机器学习·语言模型·transformer
点云SLAM14 小时前
PyTorch中 nn.Linear详解和实战示例
人工智能·pytorch·python·深度学习·cnn·transformer·mlp
苏苏susuus2 天前
NLP:Transformer各子模块作用(特别分享1)
人工智能·自然语言处理·transformer
什么都想学的阿超2 天前
【大语言模型 17】高效Transformer架构革命:Reformer、Linformer、Performer性能突破解析
语言模型·架构·transformer
兔子的倔强4 天前
Transformer在文本、图像和点云数据中的应用——经典工作梳理
人工智能·深度学习·transformer
sinat_286945194 天前
AI服务器介绍
服务器·人工智能·算法·chatgpt·transformer
Virgil1395 天前
【TrOCR】模型预训练权重各个文件解读
人工智能·pytorch·计算机视觉·自然语言处理·ocr·transformer
聚客AI7 天前
深度拆解AI大模型从训练框架、推理优化到市场趋势与基础设施挑战
图像处理·人工智能·pytorch·深度学习·机器学习·自然语言处理·transformer
Coovally AI模型快速验证7 天前
农田扫描提速37%!基于检测置信度的无人机“智能抽查”路径规划,Coovally一键加速模型落地
深度学习·算法·yolo·计算机视觉·transformer·无人机
AIGC安琪7 天前
Transformer中的编码器和解码器是什么?
人工智能·深度学习·ai·语言模型·大模型·transformer·ai大模型