目录
- 前言
- 一、optimizer.zero_grad()
- 二、 loss.backward()
- 三、optimizer.step()
前言
在用pytorch训练模型时,通常会在遍历epochs的过程中依次用到 optimizer.zero_grad(), loss.backward() 和 optimizer.step() 三个函数,如下所示:
1 2 3 4 5 6 7 8 9 10 11 12 13 | 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())
接下来将通过源码分别理解这三个函数的具体实现过程。在此之前,先简要说明一下函数中常用到的参数变量:
param_groups: Optimizer类在实例化时会在构造函数中创建一个param_groups列表,列表中有num_groups个长度为6的param_group字典,包含了 [‘params’, ‘lr’, ‘momentum’, ‘dampening’, ‘weight_decay’, ‘nesterov’] 这6组键值对。
param_group[‘params’]: 由模型参数组成的列表,模型参数即为实例化Optimizer类时传入的model.parameters(),每个参数是一个torch.nn.parameter.Parameter对象。
一、optimizer.zero_grad()
代码如下(示例):
1 2 3 4 5 6 7 | 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()函数会遍历模型的所有参数,通过p.grad.detach_()方法截断反向传播的梯度流,再通过p.grad.zero_()函数将每个参数的梯度值设为0,即上一次的梯度记录被清空。
因为训练的过程通常使用mini-batch方法,调用backward()函数之前都要将梯度清零,因为如果梯度不清零,pytorch中会将上次计算的梯度和本次计算的梯度累加。
这样逻辑的好处是:当我们的硬件限制不能使用更大的bachsize时,使用多次计算较小的bachsize的梯度平均值来代替,更方便。
坏处当然是:每次都要清零梯度总结就是进来一个batch的数据,计算一次梯度,更新一次网络。
总结:
- 常规情况下,每个batch需要调用一次optimizer.zero_grad()函数,把参数的梯度清零;
- 也可以多个batch只调用一次optimizer.zero_grad()函数,这样相当于增大了batch_size。
二、 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()之前。
三、optimizer.step()
以SGD为例,torch.optim.SGD().step()源码如下::
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 | 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()函数来计算梯度。