Pytorch网络训练流程的作用原理：源码分析optimizer.zero_grad()loss.backward()optimizer.step()

Posted 2022-06-01 呆呆象呆呆

常见参数训练流程的作用原理

1 总述

在用pytorch训练模型时，通常会在循环epoch的过程中，不断循环遍历所有训练数据集。

依次用到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() 

（学习率更新的模块lr_scheduler非必要所以暂时不放在这里讲，想了解的可以看如下文章：pytorch 动态调整学习率，学习率自动下降，根据loss下降主要集中在第6节）

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

接下来将通过源码分别理解这三个函数的具体实现过程。在此之前，先简要说明一下函数中常用到的参数变量：

• param_groups： Optimizer类在实例化时会在构造函数中创建一个param_groups列表，列表中的元素为param_group字典，此时param_group字典的数量就是param_groups列表的长度，也就是变量num_groups的含义。
• param_group字典：此字典长度为6，包含了params，lr，momentum，dampening，weight_decay，nesterov这6组键值对。
• param_group['params']字典中的键值params： 由模型参数组成的迭代器，模型参数即为实例化Optimizer类时传入并注册在模型的成员属性_parameters中，每个参数是一个torch.nn.parameter.Parameter对象。

2 optimizer.zero_grad()

代码如下（示例）：

def zero_grad(self):
    r"""Clears the gradients of all optimized :class:`torch.Tensor` s."""
    for group in self.param_groups:
        for p in group['params']:
            if p.grad is not None:
                p.grad.detach_()
                p.grad.zero_()

optimizer.zero_grad()函数会遍历模型的所有参数，这里所说的参数都是之前总述中所叙述过的torch.nn.parameter.Parameter类型变量。也就是之后的p。通过p.grad.detach_()方法截断反向传播的梯度流，再通过p.grad.zero_()函数将每个参数的梯度值设为0，即上一次的梯度记录被清空。

因为训练的过程通常使用mini-batch方法，调用backward()函数之前都要将梯度清零，因为如果梯度不清零，pytorch中会将上次计算的梯度和本次计算的梯度累加。

好处

当我们的硬件限制不能使用更大的bachsize时，代码构造成多个batchsize进行一次optimizer.zero_grad()函数调用，这样就可以使用多次计算较小的bachsize的梯度平均值来代替，更方便。

坏处

每次都要清零梯度：进来一个batch的数据，计算一次梯度，更新一次网络。

总结

1. 常规情况下，每个batch需要调用一次optimizer.zero_grad()函数，把参数的梯度清零；
2. 也可以多个batch只调用一次optimizer.zero_grad()函数，这样相当于增大了batch_size。

3 loss.backward()

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

具体来说，torch.tensor是autograd包的基础类，如果你设置tensor的requires_grads为True，就会开始跟踪这个tensor上面的所有运算。

如果做完运算后使用tensor.backward()，所有的梯度就会自动运算，tensor的梯度将会累加到它的.grad属性里面去。如果没有进行tensor.backward()的话，梯度值将会是None，因此loss.backward()要写在optimizer.step()之前。

4 optimizer.step()

以SGD为例，torch.optim.SGD().step()源码如下：

def step(self, closure=None):
    """Performs a single optimization step.
    Arguments:
        closure (callable, optional): A closure that reevaluates the model
            and returns the loss.
    """
    loss = None
    if closure is not None:
        loss = closure()

    for group in self.param_groups:
        weight_decay = group['weight_decay']
        momentum = group['momentum']
        dampening = group['dampening']
        nesterov = group['nesterov']

        for p in group['params']:
            if p.grad is None:
                continue
            d_p = p.grad.data
            if weight_decay != 0:
                d_p.add_(weight_decay, p.data)
            if momentum != 0:
                param_state = self.state[p]
                if 'momentum_buffer' not in param_state:
                    buf = param_state['momentum_buffer'] = torch.clone(d_p).detach()
                else:
                    buf = param_state['momentum_buffer']
                    buf.mul_(momentum).add_(1 - dampening, d_p)
                if nesterov:
                    d_p = d_p.add(momentum, buf)
                else:
                    d_p = buf

            p.data.add_(-group['lr'], d_p)

    return loss

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

注意：

不一样的优化器的.step()函数具体的流程基本相似

主要不同的地方是在于由于优化手段的不同，会通过所计算出来的梯度，用不同的公式去更新参数，仅此而已

具体公式如何不同或者有什么区别可以参考

深度学习笔记（四）：常用优化算法分析比较及使用策略（SGD、SGDM、SGD with Nesterov Acceleration、AdaGrad、AdaDelta、Adam、Nadam）_呆呆象呆呆的博客-CSDN博客

LAST 参考文献

optimizer.zero_grad()，loss.backward()，optimizer.step()的作用原理 | 码农家园

深度学习笔记（四）：常用优化算法分析比较及使用策略（SGD、SGDM、SGD with Nesterov Acceleration、AdaGrad、AdaDelta、Adam、Nadam）_呆呆象呆呆的博客-CSDN博客