Bootstrap

YOLO-V3-SPP 梯度累积的实现方式accumulate

前言

链接至YOLO-V3-SPP详细解析
在显存不够用的时候,YOLO采用了梯度累积的策略

源码实现

accumulate定义

train.py文件中定义了该accumulate参数:

accumulate = max(round(24 / batch_size), 1)
  1. round() 方法返回浮点数 24 b a t c h s i z e \frac{24}{batchsize} batchsize24 的四舍五入值,准确的说保留值将保留到离上一位更近的一端(四舍六入)
  2. 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()正常更新权重

;