【Pytorch】梯度裁剪——torch.nn.utils.clip_grad_norm_的原理及计算过程


众所周知,梯度裁剪是为了防止梯度爆炸。在训练FCOS算法时,因为训练过程出现了损失为NaN的情况,在github issue有很多都是这种训练过程出现loss为NaN,作者也提出要调整梯度裁剪的超参数,于是理了理梯度裁剪函数torch.nn.utils.clip_grad_norm_ 的计算过程,方便调参。


一、torch.nn.utils.clip_grad_norm_

torch.nn.utils.clip_grad_norm_(parameters, max_norm, norm_type) ,这个梯度裁剪函数一般来说只需要调整max_normnorm_type这两个参数。
parameters参数是需要进行梯度裁剪的参数列表。通常是模型的参数列表,即model.parameters()
max_norm参数可以理解为梯度(默认是L2 范数)范数的最大阈值
norm_type参数可以理解为指定范数的类型,比如norm_type=1 表示使用L1 范数,norm_type=2 表示使用L2 范数。
同时,torch.nn.utils.clip_grad_norm_torch.nn.utils.clip_grad_norm(该函数已被弃用)的区别就是前者是直接修改原Tensor,后者不会(在Pytorch中有很多这样的函数对均是如此,在函数最后多了下划线一般都是表示直接在原Tensor上进行操作)。

import torch

# 构造两个Tensor
x = torch.tensor([99.0, 108.0], requires_grad=True)
y = torch.tensor([45.0, 75.0], requires_grad=True)

# 模拟网络计算过程
z = x ** 2 + y ** 3
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=100, 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([198., 216.])
gradient of y is:tensor([ 6075., 16875.])
---clip_grad---
clip_grad of x is:tensor([1.1038, 1.2042])
clip_grad of y is:tensor([33.8674, 94.0762])
"""

上例中可以看出,裁剪后的梯度远小于原来的梯度。一开始变量x的梯度是tensor([198., 216.]),这个很好计算,就是求zx的偏导,也就是2*x 。变量y同理。裁剪后的梯度远小于原来的梯度,所以可以缓解梯度爆炸的问题。

二、计算过程

梯度裁剪的计算过程参考源码是不难的,
SOURCE CODE FOR TORCH.NN.UTILS.CLIP_GRAD
结合代码转换成数学公式,计算过程如下:
第一步:依然以上面的代码为例,构造Tensor反向传播,得到参数xy 的梯度,也就是torch.nn.utils.clip_grad_norm_(parameters, max_norm, norm_type) 中的parameters参数。
在这里插入图片描述

import torch

# 构造两个Tensor
x = torch.tensor([99.0, 108.0], requires_grad=True)
y = torch.tensor([45.0, 75.0], requires_grad=True)

# 模拟网络计算过程
z = x ** 2 + y ** 3
z = z.sum()

# 反向传播
z.backward()

# 得到梯度
print(f"gradient of x is:{
      
      x.grad}") 
print(f"gradient of y is:{
      
      y.grad}") 

# 输出
"""
gradient of x is:tensor([198., 216.])
gradient of y is:tensor([ 6075., 16875.])
"""

第二步:计算每个变量梯度的L2 范数(以L2 范数为例)
在这里插入图片描述

这段代码的意思就是定义一个空列表norms ,用来存储每个参数梯度的L2 范数。随后利用torch.stacknorms中的所有Tensor(计算所得的L2 范数)合并成一个Tensor,最后又再求合并后Tensor的L2 范数得到total_norm(总范数)。

norm_type是inf即无穷范数时,total_norm会直接取参数梯度最大的那一个

# 相当于把x_L2norm、y_L2norm放入代码中的norms空列表

# x的梯度的L2 范数
x_L2norm = torch.sum(x.grad ** 2) ** 0.5
# y的梯度的L2 范数
y_L2norm = torch.sum(y.grad ** 2) ** 0.5

# 相当于遍历norms列表合并成一个Tensor
total_norm = torch.sum(torch.stack([x_L2norm, y_L2norm]) ** 2) ** 0.5

"""
等价过程
x_L2norm = sum([198 ** 2, 216 ** 2]) ** 0.5
y_L2norm = sum([6075 ** 2, 16875 ** 2]) ** 0.5
total_norm = sum([x_L2norm ** 2, y_L2norm ** 2]) ** 0.5
"""

第三步:计算梯度裁剪系数
在这里插入图片描述

# 1e-6防止分母为0
# clip_coef = max_norm / (total_norm + 1e-6)
max_norm = 100
clip_coef = max_norm / total_norm 

第四步:将原始梯度乘以梯度裁剪系数得到裁剪后的梯度,这与函数计算的结果是一致的。
在这里插入图片描述

print(f"clip_grad of x is: is {
      
      x.grad * clip_coef }")
print(f"clip_grad of x is: is {
      
      y.grad * clip_coef }")

# 输出
"""
clip_grad of x is: is tensor([1.1038, 1.2042])
clip_grad of x is: is tensor([33.8674, 94.0762])
"""

整合一下代码:

import torch

# 构造两个Tensor
x = torch.tensor([99.0, 108.0], requires_grad=True)
y = torch.tensor([45.0, 75.0], requires_grad=True)

# 模拟网络计算过程
z = x ** 2 + y ** 3
z = z.sum()

# 反向传播
z.backward()

# 得到梯度
print(f"gradient of x is:{
      
      x.grad}") 
print(f"gradient of y is:{
      
      y.grad}") 

x_L2norm = torch.sum(x.grad ** 2) ** 0.5
y_L2norm = torch.sum(y.grad ** 2) ** 0.5
total_norm = torch.sum(torch.stack([x_L2norm, y_L2norm]) ** 2) ** 0.5

max_norm = 100
clip_coef = max_norm / total_norm 

print(f"clip_grad of x is: is {
      
      x.grad * clip_coef }")
print(f"clip_grad of x is: is {
      
      y.grad * clip_coef }")

# 输出如下
"""
gradient of x is:tensor([198., 216.])
gradient of y is:tensor([ 6075., 16875.])
clip_grad of x is:tensor([1.1038, 1.2042])
clip_grad of y is:tensor([33.8674, 94.0762])
"""

三、确定max_norm

根据上述计算过程,梯度裁剪最主要的就是计算出裁剪系数得出裁剪后的梯度。clip_coef = max_norm / total_norm 公式中,clip_coef 越小,裁剪的梯度越大。
max_norm越小,裁剪的梯度越大,得到的梯度就越小,防止梯度爆炸的效果越明显。

在训练模型时,我们可以根据total_norm的值大概确定max_norm的一个取值范围。调用torch.nn.utils.clip_grad_norm_([x, y], max_norm=100, norm_type=2)函数时,该函数会返回total_norm的值。比如最开始参数x、y,这时的total_norm17937.5879,值非常大,那么为了防止梯度爆炸我们就可以把max_norm设置得稍微小一些。

# 以最开始的x、y为例
total_norm = torch.nn.utils.clip_grad_norm_([x, y], max_norm=100, norm_type=2)
print(total_norm)

#输出
# tensor(17937.5879) 

猜你喜欢

转载自blog.csdn.net/m0_46412065/article/details/131396098