详解torch.nn.utils.clip_grad_norm_ 的使用与原理

clip_grad_norm_的原理

本文是对梯度剪裁: torch.nn.utils.clip_grad_norm_()文章的补充。所以可以先参考这篇文章

从上面文章可以看到，clip_grad_norm最后就是对所有的梯度乘以一个clip_coef，而且乘的前提是clip_coef一定是小于1的，所以，按照这个情况：clip_grad_norm只解决梯度爆炸问题，不解决梯度消失问题

clip_grad_norm_参数的选择（调参）

从上面文章可以看到，clip_coef的公式为：

$$clip\_coef=\frac{max\_norm}{total\_norm}$$

max_norm的取值：

假定忽略clip_coef > 1的情况，则可以根据公式推断出：

1. clip_coef越小，则对梯度的裁剪越厉害，即，使梯度的值缩小的越多
2. max_norm越小，clip_coef越小，所以，max_norm越大，对于梯度爆炸的解决越柔和，max_norm越小，对梯度爆炸的解决越狠

max_norm可以取小数

接下来看下total_norm和norm_type的取值：

从上面文章可以看到，total_norm受梯度大小和norm_type的影响，通过公式很难直观的感受到，这里我通过实验得出了以下结论（不严谨，欢迎勘误）：

1. 梯度越大，total_norm值越大，进而导致clip_coef的值越小，最终也会导致对梯度的裁剪越厉害，很合理
2. norm_type不管取多少，对于total_norm的影响不是太大（1和2的差距稍微大一点），所以可以直接取默认值2
3. norm_type越大，total_norm越小（实验观察到的结论，数学不好，不会证明，所以本条不一定对）

---

实验过程如下：

首先我对源码进行了一些修改，将.grad去掉，增加了一些输出，方便进行实验：

import numpy as np
import torch
from torch import nn

def clip_grad_norm_(parameters, max_norm, norm_type=2):
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]
    parameters = list(filter(lambda p: p is not None, parameters))
    max_norm = float(max_norm)
    norm_type = float(norm_type)
    if norm_type == np.inf:
        total_norm = max(p.data.abs().max() for p in parameters)
    else:
        total_norm = 0
        for p in parameters:
            param_norm = p.data.norm(norm_type)
            total_norm += param_norm.item() ** norm_type
        total_norm = total_norm ** (1. / norm_type)
    clip_coef = max_norm / (total_norm + 1e-6)
    if clip_coef < 1:
        for p in parameters:
            p.data.mul_(clip_coef)
            
    print("max_norm=%s, norm_type=%s, total_norm=%s, clip_coef=%s" % (max_norm, norm_type, total_norm, clip_coef))

只改变norm_type的情况下，各变量值的变化：

for i in range(1, 5):
    clip_grad_norm_(torch.Tensor([125,75,45,15,5]), max_norm=1, norm_type=i)
    
clip_grad_norm_(torch.Tensor([125,75,45,15,5]), max_norm=1, norm_type=np.inf)

max_norm=1.0, norm_type=1.0, total_norm=265.0, clip_coef=0.0037735848914204344
max_norm=1.0, norm_type=2.0, total_norm=153.3786163330078, clip_coef=0.006519813631054263
max_norm=1.0, norm_type=3.0, total_norm=135.16899108886716, clip_coef=0.007398146457602848
max_norm=1.0, norm_type=4.0, total_norm=129.34915161132812, clip_coef=0.007731013151704421
max_norm=1.0, norm_type=inf, total_norm=tensor(125.), clip_coef=tensor(0.0080)

只改变梯度，各变量值的变化：

for i in range(1, 5):
    clip_grad_norm_(torch.Tensor([125*i,75,45,15,5]), max_norm=1, norm_type=2)

max_norm=1.0, norm_type=2.0, total_norm=153.3786163330078, clip_coef=0.006519813631054263
max_norm=1.0, norm_type=2.0, total_norm=265.3299865722656, clip_coef=0.003768891745519864
max_norm=1.0, norm_type=2.0, total_norm=385.389404296875, clip_coef=0.0025947781289671814
max_norm=1.0, norm_type=2.0, total_norm=507.83856201171875, clip_coef=0.001969129705451148

只改变max_norm，各变量值的变化：

for i in range(1, 5):
    clip_grad_norm_(torch.Tensor([125,75,45,15,5]), max_norm=i, norm_type=2)

max_norm=1.0, norm_type=2.0, total_norm=153.3786163330078, clip_coef=0.006519813631054263
max_norm=2.0, norm_type=2.0, total_norm=153.3786163330078, clip_coef=0.013039627262108526
max_norm=3.0, norm_type=2.0, total_norm=153.3786163330078, clip_coef=0.01955944089316279
max_norm=4.0, norm_type=2.0, total_norm=153.3786163330078, clip_coef=0.02607925452421705

clip_grad_norm_使用演示

在上面文章还提到一个重要的事情：clip_grad_norm_要放在backward和step之间。接下来我会实际演示梯度在训练过程中的变化，并解释要这么做的原因：

首先定义一个测试模型：

class TestModel(nn.Module):
    
    def __init__(self):
        super().__init__()
        
        self.model = nn.Sequential(
            nn.Linear(1,1, bias=False),
            nn.Sigmoid(),
            nn.Linear(1,1, bias=False),
            nn.Sigmoid(),
            nn.Linear(1,1, bias=False),
            nn.Sigmoid(),
            nn.Linear(1,1, bias=False),
            nn.Sigmoid(),
        )
    
    def forward(self, x):
        return self.model(x)

model = TestModel()

定义好模型后，固定一下模型参数：

for param in model.parameters():
    param.data = torch.Tensor([[0.5]])
    print("param=%s" % (param.data.item()))

param=0.5
param=0.5
param=0.5
param=0.5

可以看目前四个线性层的权重参数都为0.5。之后对模型进行一轮训练，并进行反向传播：

criterion = nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=1)
predict_y = model(torch.Tensor([0.1]))
loss = criterion(predict_y, torch.Tensor([1]))
model.zero_grad()
loss.backward()

反向传播过后，再次打印模型参数，可以看到反向传播后计算好的各个参数的梯度：

for param in model.parameters():
    print("param=%s, grad=%s" % (param.data.item(), param.grad.item()))

param=0.5, grad=-3.959321111324243e-05
param=0.5, grad=-0.0016243279678747058
param=0.5, grad=-0.014529166743159294
param=0.5, grad=-0.11987950652837753

重点来了，各个参数的梯度如上图所示（越靠近输入的位置，梯度越小，虽然没有出现梯度爆炸，反而出现了梯度消失，但不影响本次实验），现在对其进行梯度裁剪：

nn.utils.clip_grad_norm_(model.parameters(), max_norm=0.01208, norm_type=2)

tensor(0.1208)

在上面，我传入的max_norm=0.01208，而total_norm=0.1208，所以可得clip_coef=0.1，即所有的梯度都会缩小一倍，此时我们再打印一下梯度：

for param in model.parameters():
    print("param=%s, grad=%s" % (param.data.item(), param.grad.item()))

param=0.5, grad=-3.960347839893075e-06
param=0.5, grad=-0.00016247491294052452
param=0.5, grad=-0.001453293371014297
param=0.5, grad=-0.01199105940759182

看到没，所有的梯度都减小了10倍。之后我们执行step()操作，其就会将进行param=param-lr*grad操作来进行参数更新。再次打印网络参数：

optimizer.step()

for param in model.parameters():
    print("param=%s, grad=%s" % (param.data.item(), param.grad.item()))

param=0.5000039339065552, grad=-3.960347839893075e-06
param=0.5001624822616577, grad=-0.00016247491294052452
param=0.5014532804489136, grad=-0.001453293371014297
param=0.5119910836219788, grad=-0.01199105940759182

可以看到，在执行step后，执行了param=param-grad操作（我设置的lr为1）。同时，grad并没有清0，所以这也是为什么要显式的调用zero_grad的原因。

---

参考资料

梯度剪裁: torch.nn.utils.clip_grad_norm_()：https://blog.csdn.net/Mikeyboi/article/details/119522689

什么是范数（Norm），其具有哪些性质: https://blog.csdn.net/zhaohongfei_358/article/details/122818616