前言
链接至YOLO-V3-SPP详细解析
在显存不够用的时候,YOLO采用了梯度累积的策略
源码实现
accumulate定义
train.py文件中定义了该accumulate参数:
accumulate = max(round(24 / batch_size), 1)
- round() 方法返回浮点数 24 b a t c h s i z e \frac{24}{batchsize} batchsize24 的四舍五入值,准确的说保留值将保留到离上一位更近的一端(四舍六入)
- max取 24 b a t c h s i z e \frac{24}{batchsize} batchsize24和1之间的最大值
我将 b a t c h _ s i z e batch\_size batch_size设为12,因此accumulate为2,意味着梯度累积次数为2
如何梯度累积?
from torch.cuda import amp
这里提出了一个amp混合精度训练的知识,参考:
PyTorch 源码解读之 torch.cuda.amp: 自动混合精度详解
对混合精度需要了解FP16和FP32的解释:
losses *= 1. / accumulate # scale loss
# backward
scaler.scale(losses).backward()
# optimize
# 每训练64张图片更新一次权重
if ni % accumulate == 0:
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()
losses为一次batch求出的loss值
乘以
1
a
c
c
u
m
u
l
a
t
e
\frac{1}{accumulate}
accumulate1得本次累积loss的占比
scaler.scale(losses).backward()放大loss梯度,并且累积这次的梯度信息
scaler.step(optimizer)将本次放大的梯度,放缩回来
scaler.update()检查是否需要梯度放大
optimizer.zero_grad()正常更新权重