Method backward

imperative/python/megengine/amp/grad_scaler.py:255–285

r"""Unscale all ``grad_tensors``'s grad and check if the gradient is finite. Args: grad_tensors: Tensors needed to unscale grads. Should be all tensors that are affected by ``target`` tensor in GradManager's backward.

(
        self,
        gm: GradManager,
        scale_factor: Tensor,
        y: Union[Tensor, List[Tensor]] = None,
        dy: Union[Tensor, List[Tensor]] = None,
    )

            self.origin_params.append(Tensor(param, no_cache=True))

    def backward(
        self,
        gm: GradManager,
        scale_factor: Tensor,
        y: Union[Tensor, List[Tensor]] = None,
        dy: Union[Tensor, List[Tensor]] = None,
    ):
        r"""Unscale all ``grad_tensors``'s grad and check if the gradient is finite.

        Args:
            grad_tensors: Tensors needed to unscale grads. Should be all tensors
                that are affected by ``target`` tensor in GradManager's backward.
        """
        if y is None:
            ys = []
        elif isinstance(y, (tuple, list)):
            ys = y
        else:
            ys = [y]
        if dy is None:
            dys = [full_like(y, scale_factor) for y in ys]
        elif isinstance(dy, (tuple, list)):
            dys = [dy_ * scale_factor for dy_ in dy]
        else:
            dys = [dy * scale_factor]

        gm.backward(y=ys, dy=dys)
        self.ifinfinite = self._check_gradients(
            [x.grad for x in gm.attached_tensors() if x.grad is not None],
            1.0 / float(scale_factor),
        )

    def update(self, model, new_scale: float = None):
        r"""Update the scale factor and model paramters according to whether encountered overflow grad.