optimizer.zero_grad()，loss.backward()，optimizer.step()的作用原理

目錄

• 前言
• 一、optimizer.zero_grad()
• 二、 loss.backward()
• 三、optimizer.step()

前言

在用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()）

接下來将通過源碼分别了解這三個函數的具體實作過程。在此之前，先簡要說明一下函數中常用到的參數變量：

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()

代碼如下（示例）：

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的資料，計算一次梯度，更新一次網絡。

總結：

1. 正常情況下，每個batch需要調用一次optimizer.zero_grad()函數，把參數的梯度清零；
2. 也可以多個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()源碼如下：：

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()函數來計算梯度。