81 lines
3.2 KiB
Python
Executable File
81 lines
3.2 KiB
Python
Executable File
from typing import *
|
|
import torch
|
|
import numpy as np
|
|
import torch.utils
|
|
|
|
|
|
class AdaptiveGradClipper:
|
|
"""
|
|
Adaptive gradient clipping for training.
|
|
"""
|
|
def __init__(
|
|
self,
|
|
max_norm=None,
|
|
clip_percentile=95.0,
|
|
buffer_size=1000,
|
|
):
|
|
self.max_norm = max_norm
|
|
self.clip_percentile = clip_percentile
|
|
self.buffer_size = buffer_size
|
|
|
|
self._grad_norm = np.zeros(buffer_size, dtype=np.float32)
|
|
self._max_norm = max_norm
|
|
self._buffer_ptr = 0
|
|
self._buffer_length = 0
|
|
|
|
def __repr__(self):
|
|
return f'AdaptiveGradClipper(max_norm={self.max_norm}, clip_percentile={self.clip_percentile})'
|
|
|
|
def state_dict(self):
|
|
return {
|
|
'grad_norm': self._grad_norm,
|
|
'max_norm': self._max_norm,
|
|
'buffer_ptr': self._buffer_ptr,
|
|
'buffer_length': self._buffer_length,
|
|
}
|
|
|
|
def load_state_dict(self, state_dict):
|
|
self._grad_norm = state_dict['grad_norm']
|
|
self._max_norm = state_dict['max_norm']
|
|
self._buffer_ptr = state_dict['buffer_ptr']
|
|
self._buffer_length = state_dict['buffer_length']
|
|
|
|
def log(self):
|
|
return {
|
|
'max_norm': self._max_norm,
|
|
}
|
|
|
|
def __call__(self, parameters, norm_type=2.0, error_if_nonfinite=False, foreach=None):
|
|
"""Clip the 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
|
|
norm_type (float): 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)
|
|
foreach (bool): use the faster foreach-based implementation.
|
|
If ``None``, use the foreach implementation for CUDA and CPU native tensors and silently
|
|
fall back to the slow implementation for other device types.
|
|
Default: ``None``
|
|
|
|
Returns:
|
|
Total norm of the parameter gradients (viewed as a single vector).
|
|
"""
|
|
max_norm = self._max_norm if self._max_norm is not None else float('inf')
|
|
grad_norm = torch.nn.utils.clip_grad_norm_(parameters, max_norm=max_norm, norm_type=norm_type, error_if_nonfinite=error_if_nonfinite, foreach=foreach)
|
|
|
|
if torch.isfinite(grad_norm):
|
|
self._grad_norm[self._buffer_ptr] = grad_norm
|
|
self._buffer_ptr = (self._buffer_ptr + 1) % self.buffer_size
|
|
self._buffer_length = min(self._buffer_length + 1, self.buffer_size)
|
|
if self._buffer_length == self.buffer_size:
|
|
self._max_norm = np.percentile(self._grad_norm, self.clip_percentile)
|
|
self._max_norm = min(self._max_norm, self.max_norm) if self.max_norm is not None else self._max_norm
|
|
|
|
return grad_norm |