参考来源：
CSDN：理解optimizer.zero_grad(), loss.backward(), optimizer.step()的作用及原理
博客园：Autograd: 自动求导

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

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

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

**

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


**param_groups**：Optimizer 类在实例化时会在构造函数中创建一个 param_groups 列表，列表中有 num_groups 个长度为 6 的 param_group 字典（num_groups 取决于你定义 optimizer 时传入了几组参数），每个 param_group 包含了 ['params', 'lr', 'momentum', 'dampening', 'weight_decay', 'nesterov'] 这 6 组键值对。
**param_group['params']**：由传入的模型参数组成的列表，即实例化 Optimizer 类时传入该 group 的参数，如果参数没有分组，则为整个模型的参数 model.parameters()，每个参数是一个 torch.nn.parameter.Parameter 对象。

1. optimizer.zero_grad()：

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

optimizer.zero_grad() 函数会遍历模型的所有参数，通过 p.grad.detach_() 方法截断反向传播的梯度流，再通过 p.grad.zero_() 函数将每个参数的梯度值设为 0，即上一次的梯度记录被清空。
因为训练的过程通常使用 **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()：

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

  1.     def step(self, closure=None):
  2.         """Performs a single optimization step.
  3.         Arguments:
  4.             closure (callable, optional): A closure that reevaluates the model
  5.                 and returns the loss.
  6.         """
  7.         loss = None
  8.         if closure is not None:
  9.             loss = closure()
  11.         for group in self.param_groups:
  12.             weight_decay = group['weight_decay']
  13.             momentum = group['momentum']
  14.             dampening = group['dampening']
  15.             nesterov = group['nesterov']
  17.             for p in group['params']:
  18.                 if p.grad is None:
  19.                     continue
  20.                 d_p = p.grad.data
  21.                 if weight_decay != 0:
  22.                     d_p.add_(weight_decay, p.data)
  23.                 if momentum != 0:
  24.                     param_state = self.state[p]
  25.                     if 'momentum_buffer' not in param_state:
  26.                         buf = param_state['momentum_buffer'] = torch.clone(d_p).detach()
  27.                     else:
  28.                         buf = param_state['momentum_buffer']
  29.                         buf.mul_(momentum).add_(1 - dampening, d_p)
  30.                     if nesterov:
  31.                         d_p = d_p.add(momentum, buf)
  32.                     else:
  33.                         d_p = buf
  35.                 p.data.add_(-group['lr'], d_p)
  37.         return loss

step() 函数的作用是执行一次优化步骤，通过梯度下降法来更新参数的值。因为梯度下降是基于梯度的，所以在执行 optimizer.step() 函数前应先执行 loss.backward() 函数来计算梯度。
注意：optimizer 只负责通过梯度下降进行优化，而不负责产生梯度，梯度是 tensor.backward() 方法产生的。