souging/TRELLIS_TextTo3D

1from typing import *2import torch3import numpy as np4import torch.utils5 6 7class AdaptiveGradClipper:8    """9    Adaptive gradient clipping for training.10    """11    def __init__(12        self,13        max_norm=None,14        clip_percentile=95.0,15        buffer_size=1000,16    ):17        self.max_norm = max_norm18        self.clip_percentile = clip_percentile19        self.buffer_size = buffer_size20        21        self._grad_norm = np.zeros(buffer_size, dtype=np.float32)22        self._max_norm = max_norm23        self._buffer_ptr = 024        self._buffer_length = 025 26    def __repr__(self):27        return f'AdaptiveGradClipper(max_norm={self.max_norm}, clip_percentile={self.clip_percentile})'28        29    def state_dict(self):30        return {31            'grad_norm': self._grad_norm,32            'max_norm': self._max_norm,33            'buffer_ptr': self._buffer_ptr,34            'buffer_length': self._buffer_length,35        }36 37    def load_state_dict(self, state_dict):38        self._grad_norm = state_dict['grad_norm']39        self._max_norm = state_dict['max_norm']40        self._buffer_ptr = state_dict['buffer_ptr']41        self._buffer_length = state_dict['buffer_length']42 43    def log(self):44        return {45            'max_norm': self._max_norm,46        }47 48    def __call__(self, parameters, norm_type=2.0, error_if_nonfinite=False, foreach=None):49        """Clip the gradient norm of an iterable of parameters.50 51        The norm is computed over all gradients together, as if they were52        concatenated into a single vector. Gradients are modified in-place.53 54        Args:55            parameters (Iterable[Tensor] or Tensor): an iterable of Tensors or a56                single Tensor that will have gradients normalized57            norm_type (float): type of the used p-norm. Can be ``'inf'`` for58                infinity norm.59            error_if_nonfinite (bool): if True, an error is thrown if the total60                norm of the gradients from :attr:`parameters` is ``nan``,61                ``inf``, or ``-inf``. Default: False (will switch to True in the future)62            foreach (bool): use the faster foreach-based implementation.63                If ``None``, use the foreach implementation for CUDA and CPU native tensors and silently64                fall back to the slow implementation for other device types.65                Default: ``None``66 67        Returns:68            Total norm of the parameter gradients (viewed as a single vector).69        """70        max_norm = self._max_norm if self._max_norm is not None else float('inf')71        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)72        73        if torch.isfinite(grad_norm):74            self._grad_norm[self._buffer_ptr] = grad_norm75            self._buffer_ptr = (self._buffer_ptr + 1) % self.buffer_size76            self._buffer_length = min(self._buffer_length + 1, self.buffer_size)77            if self._buffer_length == self.buffer_size:78                self._max_norm = np.percentile(self._grad_norm, self.clip_percentile)79                self._max_norm = min(self._max_norm, self.max_norm) if self.max_norm is not None else self._max_norm80        81        return grad_norm