RabbitRUI/ruispace

1from typing import Dict, List2 3import torch4 5if torch.__version__ < '1.9':6    Iterable = torch._six.container_abcs.Iterable7else:8    import collections9 10    Iterable = collections.abc.Iterable11from torch.cuda.amp import GradScaler12 13 14class _MultiDeviceReplicator(object):15    """16    Lazily serves copies of a tensor to requested devices.  Copies are cached per-device.17    """18 19    def __init__(self, master_tensor: torch.Tensor) -> None:20        assert master_tensor.is_cuda21        self.master = master_tensor22        self._per_device_tensors: Dict[torch.device, torch.Tensor] = {}23 24    def get(self, device) -> torch.Tensor:25        retval = self._per_device_tensors.get(device, None)26        if retval is None:27            retval = self.master.to(device=device, non_blocking=True, copy=True)28            self._per_device_tensors[device] = retval29        return retval30 31 32class MaxClipGradScaler(GradScaler):33    def __init__(self, init_scale, max_scale: float, growth_interval=100):34        GradScaler.__init__(self, init_scale=init_scale, growth_interval=growth_interval)35        self.max_scale = max_scale36 37    def scale_clip(self):38        if self.get_scale() == self.max_scale:39            self.set_growth_factor(1)40        elif self.get_scale() < self.max_scale:41            self.set_growth_factor(2)42        elif self.get_scale() > self.max_scale:43            self._scale.fill_(self.max_scale)44            self.set_growth_factor(1)45 46    def scale(self, outputs):47        """48        Multiplies ('scales') a tensor or list of tensors by the scale factor.49 50        Returns scaled outputs.  If this instance of :class:`GradScaler` is not enabled, outputs are returned51        unmodified.52 53        Arguments:54            outputs (Tensor or iterable of Tensors):  Outputs to scale.55        """56        if not self._enabled:57            return outputs58        self.scale_clip()59        # Short-circuit for the common case.60        if isinstance(outputs, torch.Tensor):61            assert outputs.is_cuda62            if self._scale is None:63                self._lazy_init_scale_growth_tracker(outputs.device)64            assert self._scale is not None65            return outputs * self._scale.to(device=outputs.device, non_blocking=True)66 67        # Invoke the more complex machinery only if we're treating multiple outputs.68        stash: List[_MultiDeviceReplicator] = []  # holds a reference that can be overwritten by apply_scale69 70        def apply_scale(val):71            if isinstance(val, torch.Tensor):72                assert val.is_cuda73                if len(stash) == 0:74                    if self._scale is None:75                        self._lazy_init_scale_growth_tracker(val.device)76                    assert self._scale is not None77                    stash.append(_MultiDeviceReplicator(self._scale))78                return val * stash[0].get(val.device)79            elif isinstance(val, Iterable):80                iterable = map(apply_scale, val)81                if isinstance(val, list) or isinstance(val, tuple):82                    return type(val)(iterable)83                else:84                    return iterable85            else:86                raise ValueError("outputs must be a Tensor or an iterable of Tensors")87 88        return apply_scale(outputs)89