梯度剪裁源码地址以及函数的说明
我的理解：
对于存在梯度爆炸的情况， 在优化器函数之前执行这个函数，可以重新整合一遍梯度梯度缩小到指定范围。
函数需要的参数：
- parameters：计算了梯度之后的权重参数
- max_norm：认为设定的阈值
- 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
