欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

PyTorch——PyTorch也支持通过累加操作实现大的BatchSize的训练

程序员文章站 2022-03-21 19:49:06
...

1 前言

这几天在学习的时候,一直在想一个问题,为什么PyTorch不支持大的BatchSize的计算呢?

按照道理来说,这并不是一个问题,因为梯度值是可以累加的,

如果设置的BatchSize较大,而GPU的显存一次性无法装入这么多的数据进行训练的话,

只要分批次计算loss就可以了;

后来发现原来PyTorch也是支持通过累加梯度实现大BatchSize训练的操作的。

2 PyTorch通过累加梯度实现大BatchSize的训练

这是我在微信群聊中其他同学共享的代码,

loss.backward()  # 计算梯度
if ((i + 1) % accumulation_steps) == 0:
    optimizer.step()  # 反向传播,更新网络参数
    optimizer.zero_grad()  # 清空梯度