6.2 梯度裁剪#

在3.3.7节中，我们首次介绍了深度学习中的梯度爆炸问题，其根本原因在于反向传播算法在求解模型梯度时累乘的计算特性导致越靠近输入层的权重参数越容易出现梯度爆炸的现象。通常来说解决梯度爆炸最直接的两种方法分别是使用较小的学习率和对梯度的大小进行限制。在接下来的这节内容中，我们将介绍深度学习中两种使用最为广泛的梯度裁剪（Gradient Clip）策略以及各自对应的使用方法[1]。

6.2.1 基于阈值裁剪#

基于阈值的梯度裁剪方法是梯度裁剪中最直接也是最简单的方法，其核心思想便是根据给定的区间范围对于现有的梯度进行约束，对于超过该范围的梯度值将直接重新被赋值为区间的端点值。

例如对于梯度值$[-3.8,-1.2,1.5,2.8]$，给定梯度的最大值为2.0，那么根据该梯度值将会被限定在区间$[-2.0,2.0]$之间，则裁剪后的梯度值便为$[-2.0,-1.2,1.5,2.0]$。下面借用PyTorch框架中的clip_grad_value_函数通过一段简单的代码进行示例：

 1 def test_grad_clip(clip_value=0.8):
 2     w = torch.tensor([[1.5, 0.5, 3.0],[0.5, 1., 2.]],
 3                      dtype=torch.float32, requires_grad=True)
 4     b = torch.tensor([2., 0.5, 3.5], dtype=torch.float32, requires_grad=True)
 5     x = torch.tensor([[2, 3.]], dtype=torch.float32)
 6     y = torch.mean(torch.matmul(x, w ** 2) + b ** 2)
 7     y.backward()
 8     print("# 梯度裁剪前: ")
 9     print(f"grad_w: {w.grad}")
10     print(f"grad_b: {b.grad}")
11     torch.nn.utils.clip_grad_value_([w, b], clip_value)
12     print(f"# 梯度裁剪后: ")
13     print(f"grad_w: {w.grad}")
14     print(f"grad_b: {b.grad}")

在上述代码中，第2~6行是定义了一个简单的线性变换计算过程（平方只是为了让每个权重参数求解得到的梯度不同）。第7行是计算y关于参数w和b的梯度。第11行便是对计算完成的梯度进行裁剪。

以下便是上述代码运行结束后的结果：

1 # 梯度裁剪前: 
2 grad_w: tensor([[2.0000, 0.6667, 4.0000],
3                 [1.0000, 2.0000, 4.0000]])
4 grad_b: tensor([1.3333, 0.3333, 2.3333])
5 # 梯度裁剪后: 
6 grad_w: tensor([[0.8000, 0.6667, 0.8000],
7                 [0.8000, 0.8000, 0.8000]])
8 grad_b: tensor([0.8000, 0.3333, 0.8000])

从上述结果可以看出，在进行梯度裁剪后参数w和b的梯度均被约束在范围$[-0.8,0.8]$中。

6.2.2 基于范数裁剪#

在介绍完基于阈值的梯度裁剪策略后，我们再来看第2种基于范数的裁剪方法。基于范数的裁剪方法其核心思想是：①先计算所有参数梯度各自的$P$范数；②然后再计算第①步中得到的各个参数梯度$P$范数的$P$范数；③进一步根据给定的最大范数同第②步得到的$P$范数计算得到一个缩放系数；④最后再将该系数作用于原始各个参数的梯度得到最终裁剪后的结果。

下面以第6.2.1节示例中梯度裁剪前的结果为例，且采用2范数，最大范数$\text{max\_norm}=1.2$进行示例。

根据式(6-7)可得，首先可以计算得到两个参数梯度各自的2范数，然后再计算整体的2范数，此时有

进一步，根据式(6-8)中的范数总和同$\text{max\_norm}=1.2$计算得到缩放系数，并作用于式(6-7)中的原始值得到最后裁剪后的结果，此时有

从式(6-9)可以看出，基于范数的裁剪策略关键在于计算得到裁剪的缩放系数，且给定的最大范数越小则梯度被裁剪得越小。同时，这需要注意的是式(6-9)中的缩放系数还会再做一次裁剪，即当超过1时仍旧取1，因为裁剪的目的便是为了缩小梯度所以不能乘以一个大于1的系数。

同样，上述结果也可以用PyTorch框架中的clip_grad_norm_函数进行计算，并且只需要上面第11行改为如下代码即可

1 torch.nn.utils.clip_grad_norm_([w, b], max_norm=1.2, norm_type=2.0)

其中norm_type用来指定使用什么样的范数。

在使用基于范数的裁剪方法后，第6.2.1节中的示例将会得到如下结果

1 # 梯度裁剪前: 
2 grad_w: tensor([[2.0000, 0.6667, 4.0000],
3                 [1.0000, 2.0000, 4.0000]])
4 grad_b: tensor([1.3333, 0.3333, 2.3333])
5 # 梯度裁剪后: 
6 grad_w: tensor([[0.3436, 0.1145, 0.6873],
7                 [0.1718, 0.3436, 0.6873]])
8 grad_b: tensor([0.2291, 0.0573, 0.4009])

可以发现，上述计算结果便是式(6-9)中所示的结果。

6.2.3 使用示例#

在介绍完两种梯度裁剪策略后，我们最后再来看如何在模型训练时进行使用。以下完整示例代码可以参见Code/Chapter06/C02_GradClip/train.py文件。

对于梯度裁剪，我们只需要每次在执行梯度下降之前对计算得到的梯度进行裁剪即可，示例代码如下所示：

1 for epoch in range(epochs):
2     for i, (x, y) in enumerate(train_iter):
3         loss, logits = model(x, y)
4         optimizer.zero_grad()
5         loss.backward()
6         # torch.nn.utils.clip_grad_norm_(model.parameters(), 0.5)
7         torch.nn.utils.clip_grad_value_(model.parameters(), 0.5)
8         optimizer.step()  # 执行梯度下降

在上述代码中，第6~7行便是基于范数和基于阈值的梯度裁剪策略，其余部分的代码也不需要做任何调整。

6.2.4 小结#

在本节内容中，我们首先分别介绍了两种梯度裁剪策略的基本原理；然后介绍了两种方法在PyTorch中的使用方法；最后介绍了如何将其加入到模型的训练过程中。

引用#

[1] Paszke A, Gross S, Massa F, et al. Pytorch: An imperative style, high-performance deep learning library[J]. Advances in neural information processing systems, 2019, 32.