loss.backward() 和optimizer.step()的关系及灵活运用

在deep learning的模型训练中，我们经常看到如下的代码片段：

loss.backward()
optimizer.step()

那么，这两个函数到底是怎么联系在一起的呢？

loss.backward()的作用

我们都知道，loss.backward()函数的作用是根据loss来计算网络参数的梯度，其对应的输入默认为网络的叶子节点，即数据集内的数据，叶子节点如下图所示:

同样的，该梯度信息也可以用函数torch.autograd.grad()计算得到

x = torch.tensor(2., requires_grad=True)
y = torch.tensor(3., requires_grad=True)

z = x * x * y
z.backward()
print(x.grad)
>>> tensor(12.)

x = torch.tensor(2., requires_grad=True)
y = torch.tensor(3., requires_grad=True)

z = x * x * y
x_grad = torch.autograd.grad(outputs=z, inputs=x)
print(x_grad[0])
>>> tensor(12.)

以上内容引自https://zhuanlan.zhihu.com/p/279758736

optimizer.step()的作用

优化器的作用就是针对计算得到的参数梯度对网络参数进行更新，所以要想使得优化器起作用，主要需要两个东西：

• 优化器需要知道当前的网络模型的参数空间
• 优化器需要知道反向传播的梯度信息（即backward计算得到的信息）

观察一下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'] = d_p.clone()
                    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

我们可以看到里面有如下的代码

for p in group['params']:
    if p.grad is None:
        continue
        d_p = p.grad.data

说明，step()函数确实是利用了计算得到的梯度信息，且该信息是与网络的参数绑定在一起的，所以optimizer函数在读入是先导入了网络参数模型’params’，然后通过一个.grad()函数就可以轻松的获取他的梯度信息。

如何验证该关系的正确性

我们想通过改变梯度信息来验证该关系的正确性，即是否可以通过一次梯度下降后，再通过一次梯度上升来得到初始化的参数

import torch
import torch.nn as nn

#  Check if we have a CUDA-capable device; if so, use it
device = 'cuda' if torch.cuda.is_available() else 'cpu'
print('Will train on {}'.format(device))

#  为了让参数恢复成初始化状态，使用最简单的SGD优化器
optimizer = torch.optim.SGD(net.parameters(), lr=0.1)

#  定义模型
class CNN(nn.Module):
    def __init__(self):
        super(CNN, self).__init__()
        
        self.linear = nn.Linear(3,1)
        
    def forward(self, x):
        y = self.linear(x)
        return y
    
#  载入模型与输入，并打印此时的模型参数
x = (torch.rand(3)).to(device)
net = CNN().to(device)
print('the first output!')
for name, parameters in net.named_parameters():
    print(name, ':', parameters)
    
print('-------------------------------------------------------------------------------')    
#  做梯度下降
optimizer.zero_grad()
y = net(x)
loss = (1-y)**2

loss.backward()
optimizer.step()
#  打印梯度信息
for name, parameters in net.named_parameters():
    print(name, ':', parameters.grad)
#  经过第一次更新以后，打印网络参数
for name, parameters in net.named_parameters():
    print(name, ':', parameters)
    
print('-------------------------------------------------------------------------------')
#  我们直接将网络参数的梯度信息改为相反数来进行梯度上升
for name, parameters in net.named_parameters():
    parameters.grad *= -1
#  打印
for name, parameters in net.named_parameters():
    print('the second output!')
    print(name, ':', parameters.grad)

经过对比，我们发现最后的结果与我们的设想一样，网络参数恢复成初始化状态，因此可以证明optimizer.step()与loss.backward()之间的关系。