torch.nn.utils.clip_grad_norm

梯度剪裁源码地址以及函数的说明

函数源码
函数官方说明

我的理解：

对于存在梯度爆炸的情况， 在优化器函数之前执行这个函数，可以重新整合一遍梯度梯度缩小到指定范围。

函数需要的参数：

1. parameters：计算了梯度之后的权重参数
2. max_norm：认为设定的阈值
3. norm_type：指定的范数

函数执行的操作
1. 对所有需要进行梯度计算的参数，收集所有参数的梯度的指定范数（通过参数norm_type进行设置，1表示绝对值，2表示二阶范数也就是平方和开根号）

2. 计算所有参数的梯度范数总和（一个标量）和设定的max_norm的比值。如果max_norm/total_norm>1, 所有参数的梯度不变，可以直接反向传播。如果比值小于1，说明参数梯度需要被缩减，缩减比率为rate= max_norm/total_norm,所有反向传播的梯度变为原本的rate倍。

这样的意义就是避免权重梯度爆炸导致模型训练困难,对于大梯度的缩小，小梯度的不变。
但是存在的问题是，参数原本的分布很不均匀，有的梯度大有的梯度小；而梯度的总体范数值对于阈值，那么所有的梯度都会被同比例缩小。

import warnings
import torch
from torch._six import inf
from typing import Union, Iterable

_tensor_or_tensors = Union[torch.Tensor, Iterable[torch.Tensor]]

def clip_grad_norm_(
        parameters: _tensor_or_tensors, max_norm: float, norm_type: float = 2.0,
        error_if_nonfinite: bool = False) -> torch.Tensor:
    r"""Clips gradient norm of an iterable of parameters.

 The norm is computed over all gradients together, as if they were
 concatenated into a single vector. Gradients are modified in-place.

 Args:
 parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a
 single Tensor that will have gradients normalized
 max_norm (float or int): max norm of the gradients
 norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
 infinity norm.
 error_if_nonfinite (bool): if True, an error is thrown if the total
 norm of the gradients from :attr:``parameters`` is ``nan``,
 ``inf``, or ``-inf``. Default: False (will switch to True in the future)

 Returns:
 Total norm of the parameters (viewed as a single vector).
 """
    if isinstance(parameters, torch.Tensor):
        parameters = [parameters]
    parameters = [p for p in parameters if p.grad is not None]
    max_norm = float(max_norm)
    norm_type = float(norm_type)
    if len(parameters) == 0:
        return torch.tensor(0.)
    device = parameters[0].grad.device
    if norm_type == inf:
        norms = [p.grad.detach().abs().max().to(device) for p in parameters]
        total_norm = norms[0] if len(norms) == 1 else torch.max(torch.stack(norms))
    else:
        total_norm = torch.norm(torch.stack([torch.norm(p.grad.detach(), norm_type).to(device) for p in parameters]), norm_type)
    if error_if_nonfinite and torch.logical_or(total_norm.isnan(), total_norm.isinf()):
        raise RuntimeError(
            f'The total norm of order {norm_type} for gradients from '
            '`parameters` is non-finite, so it cannot be clipped. To disable '
            'this error and scale the gradients by the non-finite norm anyway, '
            'set `error_if_nonfinite=False`')
    clip_coef = max_norm / (total_norm + 1e-6)
    # Note: multiplying by the clamped coef is redundant when the coef is clamped to 1, but doing so
    # avoids a `if clip_coef < 1:` conditional which can require a CPU <=> device synchronization
    # when the gradients do not reside in CPU memory.
    clip_coef_clamped = torch.clamp(clip_coef, max=1.0)
    for p in parameters:
        p.grad.detach().mul_(clip_coef_clamped.to(p.grad.device))
    return total_norm