梯度剪裁: torch.nn.utils.clip_grad_norm_()

一、原理

pytorch中梯度剪裁方法为 torch.nn.utils.clip_grad_norm_(parameters, max_norm, norm_type=2)1。三个参数：

parameters：希望实施梯度裁剪的可迭代网络参数
max_norm：该组网络参数梯度的范数上限
norm_type：范数类型

我们来逐段分析其实现代码：

def clip_grad_norm_(parameters, max_norm, norm_type=2):if isinstance(parameters, torch.Tensor):parameters = [parameters]parameters = list(filter(lambda p: p.grad is not None, parameters))max_norm = float(max_norm)norm_type = float(norm_type)

该部分处理了传入的三个参数。首先将parameters中的非空网络参数存入一个列表，然后将max_norm和norm_type类型强制为浮点数。

    if norm_type == inf:total_norm = max(p.grad.data.abs().max() for p in parameters)

该句对无穷范数进行了单独计算，即取所有网络参数梯度范数中的最大值，定义为total_norm：

    else:total_norm = 0for p in parameters:param_norm = p.grad.data.norm(norm_type)total_norm += param_norm.item() ** norm_typetotal_norm = total_norm ** (1. / norm_type)

对于其他范数，我们计算所有网络参数梯度范数之和，再归一化，即等价于把所有网络参数放入一个向量，再对向量计算范数。将结果定义为total_norm：

    clip_coef = max_norm / (total_norm + 1e-6)if clip_coef < 1:for p in parameters:p.grad.data.mul_(clip_coef)return total_norm

最后定义了一个“裁剪系数”变量clip_coef，为传入参数max_norm和total_norm的比值（+1e-6防止分母为0的情况）。如果max_norm > total_norm，即没有溢出预设上限，则不对梯度进行修改。反之则以clip_coef为系数对全部梯度进行惩罚，使最后的全部梯度范数归一化至max_norm的值。注意该方法返回了一个 total_norm，实际应用时可以通过该方法得到网络参数梯度的范数，以便确定合理的max_norm值。

二、使用方法

每一次迭代中，梯度处理的过程应该是：

因此 torch.nn.utils.clip_grad_norm_() 的使用应该在loss.backward() 之后，**optimizer.step()**之前：

...
loss = crit(...)
optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(parameters=model.parameters(), max_norm=10, norm_type=2)
optimizer.step()
...