| 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 |
