深度学习笔记（3）-pytorch模型训练流程&实现小GPU显存跑大Batchsize

近期在进行pytorch模型的训练，对pytorch的流程进行一次简单梳理作为笔记。此外，由于GPU显存有限，数据的Batchsize一般只能到2，而相关资料显示较大的Batchsize有利于提高模型训练效果，经查阅资料，找到通过梯度累加的方式来等效增大Batchsize。

一、pytorch模型训练流程

在用pytorch训练模型时，通常会在遍历epochs的过程中依次用到optimizer.zero_grad(), loss.backward()和optimizer.step()三个函数，如下所示：

model = MyModel()
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.001, momentum=0.9, weight_decay=1e-4)
 
for epoch in range(1, epochs):
    for i, (inputs, labels) in enumerate(train_loader):
        output= model(inputs)
        loss = criterion(output, labels)
        
        # compute gradient and do SGD step
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

总得来说，这三个函数的作用是先将梯度归零（optimizer.zero_grad()），然后反向传播计算得到每个参数的梯度值（loss.backward()），最后通过梯度下降执行一步网络参数更新（optimizer.step()）。

1. optimizer.zero_grad

   训练的过程通常使用mini-batch方法，如果不将梯度清零的话，梯度会与上一个batch的数据相关，因此该函数要写在反向传播和梯度下降之前。

2. loss.backward()

   PyTorch的反向传播(即tensor.backward())是通过autograd包来实现的，autograd包会根据tensor进行过的数学运算来自动计算其对应的梯度。

   梯度可以理解为：函数各个自变量求偏导，这些偏导数写成向量的形式就是梯度。

   梯度方向各点处函数值减小最多的方向，梯度的方向不一定指向函数的最小值，但一定指向函数值减小最多的方向。

   具体来说，torch.tensor是autograd包的基础类，如果你设置tensor的requires_grads为True，就会开始跟踪这个tensor上面的所有运算，如果你做完运算后使用tensor.backward()，所有的梯度就会自动运算，tensor的梯度将会累加到它的.grad属性里面去。

   更具体地说，损失函数loss是由模型的所有权重w经过一系列运算得到的，若某个w的requires_grads为True，则w的所有上层参数（后面层的权重w）的.grad_fn属性中就保存了对应的运算，然后在使用loss.backward()后，会一层层的反向传播计算每个w的梯度值，并保存到该w的.grad属性中。

   如果没有进行tensor.backward()的话，梯度值将会是None，因此loss.backward()要写在optimizer.step()之前。

3. optimizer.step()
   step()函数的作用是执行一次优化步骤，通过梯度下降法来更新参数的值。因为梯度下降是基于梯度的，所以在执行optimizer.step()函数前应先执行loss.backward()函数来计算梯度。

   注意：optimizer只负责通过梯度下降进行优化，而不负责产生梯度，梯度是tensor.backward()方法产生的。

二、小GPU显存跑大Batchsize

对于模型训练来说，batch_size越大，模型效果会越好。但是某些环境下，没有足够的GPU来支撑起大的batch_size，因此这时可以考虑使用accumulate_steps来达到类似的效果。

具体地，原来训练过程中每个batch_size都会进行梯度更新，这时我们可以采取每训练（叠加）accumulate_steps个batch_size再更新梯度（这个操作就相当于将batch_size扩大accumulate_steps倍）

一般训练函数，一个batch训练过程如下：

for i, (image, label) in enumerate(train_loader):
    # 1. input output
    pred = model(image)
    loss = criterion(pred, label)

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

步骤如下：
1.正向传播，将数据传入网络，得到预测结果；
2.根据预测结果与label，计算损失；
3.将前一个batch计算之后的网络梯度清零；
4.利用损失进行反向传播，计算参数梯度；
5.利用计算的参数梯度更新网络参数。

当进行手动梯度累加的时候，代码如下：

for i,(image, label) in enumerate(train_loader):
    # 1. input output
    pred = model(image)
    loss = criterion(pred, label)

    # 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

步骤如下：

1.正向传播，将数据传入网络，得到预测结果；
2.根据预测结果与label，计算损失；
3. 根据累加值，计算平均loss。平均的目的在于：loss平均得到平均梯度，loss.backward计算时会释放计算图，如果不平均，梯度过大，内存爆掉；
4. loss.backward 反向传播计算梯度；
5.当不足一个accumulation_steps值时，梯度累加，实现等效大batchsize的目的；
6. 当满足一个accumulation_steps值时，optimizer.step进行网络参数更新，optimizer.zero_grad清零上一组梯度，为下一轮梯度累加做准备。

Tip：累加的是梯度而不是loss，每次计算loss并backward()后，计算图就已经释放了，所以不会引起爆内存的问题，但是此时梯度保留了下来，一起帮助参数更新。梯度的regularization我的理解是可以帮助许多小batch的loss加在一起反向传播时更加稳定。

问题1：使用梯度累加时，对应LR要怎么变？

原因：bs增大之后，每个batch的可信度更高，同时，bs增大之后整个训练过程的迭代次数就少了，因此需要增加学习率在更少的迭代次数内完成同样的工作。

New LR = LR * accumulation_steps

(等效Batchsize一般不超过64？)

问题2：为什么要loss = loss / accumulation_steps ？
loss/steps主要影响到的梯度，如果不除以steps，也就相当于直接将8次的梯度累加到模型上，在反传的时候，loss会波动很大,但是如果除以steps，相当于平均了这steps次梯度的累积过程，很平滑。有点类似于滑动平均的样子，只不过这个滑动平均是在steps次积累上做的.

Reference：
1. https://www.cnblogs.com/sddai/p/14598018.html