【深度学习训练小技巧】1080ti与2080ti区别、apex与梯度累加

当我们没有足够的显卡训练模型时，apex和梯度累加是有效的减少显卡资源使用的手段。

1080ti与2080ti区别

1080ti与2080ti均为11G显存，在训练同样的模型时batch_size可以调到一样大。它们的主要区别在于2080ti具有很强的半精度浮点数计算能力。1080ti和2080ti的计算能力如下：

gtx 1080ti:
半精度浮点数：0.17TFLOPS
单精度浮点数：11.34TFLOPS
双精度浮点数：0.33TFLOPS
rtx 2080ti:
半精度浮点数：20.14TFLOPS
单精度浮点数：10.07TFLOPS
双精度浮点数：0.31TFLOPS

半精度浮点数即FP16，单精度浮点数即FP32，双精度浮点数即FP64。
在不使用apex的pytorch训练过程中，一般默认均为单精度浮点数，从上面的数据可以看到1080ti和2080ti的单精度浮点数运算能力差不多，因此不使用apex时用1080ti和2080ti训练模型时间上差别很小。
在文章：【庖丁解牛】从零实现RetinaNet（六）：RetinaNet的训练与测试中，不使用apex时用2个2080ti训练时一个epoch要2h38min，而使用apex时用1个2080ti训练时一个epoch是2h31min，两者时间几乎一样，但是却少用了一张2080ti。这是因为在pytorch训练中使用apex时，此时大多数运算均为半精度浮点数运算，而2080ti的半精度浮点数运算能力是其单精度浮点数运算能力的两倍（1080ti几乎没有半精度浮点数运算能力）。这就为什么少用了一张2080ti但是一个epoch的训练时间差不多的原因。对于1080ti，由于其几乎没有半精度浮点数运算能力，即使使用apex，最终实际也是按单精度浮点数来计算的，因此起不到训练提速效果。

在目标检测和分割任务中使用apex

使用apex能够显著的降低显存占用25-30%，通常情况下同样显卡可以把batch_size扩大接近一倍，非常实用。
类似于base model第七弹中所述，我们先安装apex工具库。如果你是Python3.7环境，使用下列命令安装：

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./

如果你是Python3.6环境，则使用下列命令安装：

git clone https://github.com/NVIDIA/apex
cd apex
pip install -v --no-cache-dir ./

安装完成后，在train.py中进行下列修改：

from apex import amp
......
......
# 定义好Model和optimizer后，增加下面这行代码：
amp.register_float_function(torch, 'sigmoid')
amp.register_float_function(torch, 'softmax')
model, optimizer = amp.initialize(model, optimizer, opt_level='O1')
# 反传梯度原本的代码是loss.backward()，改为下面两行代码：
with amp.scale_loss(loss, optimizer) as scaled_loss:
    scaled_loss.backward()

这里和base model第七弹中的区别是增加了两行amp.register_float_function函数。这个函数的作用是将模型所有sigmoid和softmax层强制使用FP32精度计算。因为FP16的位数较少，能保存数据的上限和下限的绝对值也变小了，在检测和分割问题中sigmoid和softmax层很容易在求和时导致数据溢出。在我的实验中，不使用amp.register_float_function时几乎每次都会在epoch5到8时loss变为nan，因此这个操作是十分必要的。对于分类任务，一般不加amp.register_float_function也不会溢出。

梯度累加（一般不在目标检测中使用）

当我们的显卡资源不够时，还可以考虑使用梯度累加。所谓梯度累加就是把一个大的batchsize分成若干次小batchsize计算梯度，每次梯度除以次数，然后进行梯度累加，这样若干次后累加的梯度就近似为大batchsize计算出来的梯度。需要说明的是，这种模拟是近似的，因为BN层仍然按照实际的小batchsize更新。
我们一般可以在分类任务中使用梯度累加，因为分类任务中即使划分为小batchsize时小batchsize也比较大（一般大于32），对BN层的影响有限。对于目标检测任务，由于总的batchsize本身就比较小，使用梯度累加会导致BN层的误差较大，所以我们一般不在目标检测中使用梯度累加。
在train.py中，进行下列修改，就可以实现梯度累加：

inputs, labels = inputs.cuda(), labels.cuda()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss = loss / args.accumulation_steps
loss.backward()

if iter_index % args.accumulation_steps == 0:
	optimizer.step()
	optimizer.zero_grad()

其中，loss = loss / args.accumulation_steps根据累加次数把loss做除法，使得在不修改学习率的情况下最终累加的反传梯度数量级与原来一致（由于我们在累加的过程中考虑了多次小batchsize产生的loss，但又除以了累加次数，所以在loss.backward()多次累加后就近似相当于大batchsize计算出的loss产生的梯度）。最后，我们经过累加的梯度次数后才使用optimizer.step()更新一次网络参数，然后使用optimizer.zero_grad()将梯度清零。

torch.no_grad()

在测试时，我们不需要对所有变量计算梯度，只需要计算前向传播即可。我们可以用with torch.no_grad():包裹不需要计算梯度的那部分代码。虽然测试时因为没有loss.backward()，即使产生了梯度也不会更新变量，但使用后可以减少这部分计算梯度的计算量和梯度占用的显存。
要注意model.eval()和with torch.no_grad():是不同的，model.eval()只是把模型中所有BN层设置为评估模式，这样有前向计算时模型的BN层均值和标准差也不会更新。