Pytorch中的梯度累加【在实验时,由于GPU显存限制,遇到batch_size不能再增大的情况。为解决该问题,使用梯度累加方法】

当GPU显存不足,无法增大batch_size时,可以采用梯度累加策略。通过设置accumulation_steps=4,将原始batch_size=32降至8,依然能达到之前的训练效果。在每个步骤中,损失先除以积累步数,然后进行反向传播,仅在积累步数的整数倍时更新网络参数并清零梯度。

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

在实验时,由于GPU显存限制,遇到batch_size不能再增大的情况。为解决该问题,使用梯度累加方法。

不进行梯度累加的方法如下:

for i,(images,target) in enumerate(train_loader):
    # 1. input output
    images = images.cuda(non_blocking=True)
    target = torch.from_numpy(np.array(target)).float().cuda(non_blocking=True)
    outputs = model(images)
    loss = criterion(outputs,target)

    # 2. backward
    optimizer.zero_grad()   # reset gradient
    loss.backward()
    optimizer.step()

使用梯度累加:

for i,(images,target) in enumerate(train_loader):
    # 1. input output
    images = images.cuda(non_blocking=True)
    target = torch.from_numpy(np.array(target)).float().cuda(non_blocking=True)
    outputs = model(images)
    loss = criterion(outputs,target)

    # 2.1 loss regularization
    loss = loss/accumulation_steps
    # 2.2 back propagation
    loss.backward()

    # 3. update parameters of net
    if((i+1)%accumulation_steps)==0:
        # optimizer the net
        optimizer.step()        # update parameters of net
        optimizer.zero_grad()   # reset gradient

原本batch size为32,使用梯度累加,设置accumulation_steps=4,此时只需将batch_size设置为8,就能达到之前的效果。




参考资料:
Pytorch中的梯度累加

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值