梯度裁剪：torch.nn.utils.clip_grad_norm_详解

梯度裁剪是为了防止梯度爆炸。在训练FCOS算法时，因为训练过程出现了损失为NaN的情况，在github issue有很多都是这种训练过程出现loss为NaN，使用torch.nn.utils.clip_grad_norm_梯度裁剪函数，可以有效预防梯度爆炸的情况发生。

1 clip_grad_norm_介绍

1.1 函数原型

def clip_grad_norm_(
        parameters: _tensor_or_tensors, max_norm: float, norm_type: float = 2.0,
        error_if_nonfinite: bool = False, foreach: Optional[bool] = None) -> torch.Tensor:

• parameters：需要进行梯度裁剪的参数列表。通常是模型的参数列表，即model.parameters()
• max_norm：可以理解为梯度（默认是L2 范数）范数的最大阈值
• norm_type：可以理解为指定范数的类型，比如norm_type=1 表示使用L1 范数，norm_type=2 表示使用L2 范数。

这个梯度裁剪函数一般来说只需要调整max_norm 和norm_type这两个参数。clip_grad_norm_最后就是对所有的梯度乘以一个clip_coef，而且乘的前提是clip_coef一定是小于1的，所以，按照这个情况：clip_grad_norm只解决梯度爆炸问题，不解决梯度消失问题

torch.nn.utils.clip_grad_norm_和torch.nn.utils.clip_grad_norm(已弃用)的区别就是前者是直接修改原Tensor，后者不会。在Pytorch中有很多这样的函数对均是如此，在函数最后多了下划线一般都是表示直接在原Tensor上进行操作。

1.2 参数的选择

clip_coef的公式为：

max_norm的取值：

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

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

    max_norm可以取小数

total_norm受梯度大小和norm_type的影响：

• 梯度越大，total_norm值越大，进而导致clip_coef的值越小，最终也会导致对梯度的裁剪越厉害，很合理
• norm_type不管取多少，对于total_norm的影响不是太大，所以可以直接取默认值2
• norm_type越大，total_norm越小

1.3 函数执行的操作

• 对所有需要进行梯度计算的参数，收集所有参数的梯度的指定范数（通过参数norm_type进行设置，1表示绝对值，2表示二阶范数也就是平方和开根号）
• 计算所有参数的梯度范数总和（一个标量）和设定的max_norm的比值。如果max_norm/total_norm>1, 所有参数的梯度不变，可以直接反向传播。如果比值小于1，说明参数梯度需要被缩减，缩减比率为rate= max_norm/total_norm,所有反向传播的梯度变为原本的rate倍。

这样就避免权重梯度爆炸导致模型训练困难，对于大梯度的缩小，小梯度的不变。

1.4 梯度裁剪存在的问题

• 参数原本的分布很不均匀，有的梯度大有的梯度小；
• 而梯度的总体范数值大于阈值，那么所有的梯度都会被同比例缩小。

2 clip_grad_norm_使用

import torch
 
# 构造两个Tensor
x = torch.tensor([102.0, 155.0], requires_grad=True)
y = torch.tensor([201.0, 221.0], requires_grad=True)
 
# 模拟网络计算过程
z = x ** 3 + y ** 4
z = z.sum()
 
# 反向传播
z.backward()
 
# 得到梯度
print(f"gradient of x is:{x.grad}")
print(f"gradient of y is:{y.grad}")
 
# 梯度裁剪
torch.nn.utils.clip_grad_norm_([x, y], max_norm=200, norm_type=2)
 
 
# 再次打印裁剪后的梯度
# 直接修改了原x.grad的值
print("---clip_grad---")
print(f"clip_grad of x is:{x.grad}")
print(f"clip_grad of y is:{y.grad}")

运行结果显示如下：

gradient of x is:tensor([31212., 72075.])
gradient of y is:tensor([32482404., 43175444.])
---clip_grad---
clip_grad of x is:tensor([0.1155, 0.2668])
clip_grad of y is:tensor([120.2386, 159.8205])

上例中可以看出，裁剪后的梯度远小于原来的梯度。一开始变量x的梯度是tensor([31212., 72075.])，就是求z对x的偏导，变量y同理。裁剪后的梯度远小于原来的梯度，所以可以缓解梯度爆炸的问题。