Class GradManager

imperative/python/megengine/autodiff/grad_manager.py:30–360

r"""GradManager computes gradients or more generally, vector-Jacobian product, by reverse mode automatic differentiation (a.k.a. back propagation). Reverse mode autodiff normally reuses many intermediate tensors for best computation efficiency. In a read-eval-print-loop (REPL) environme

class GradManager:
    r"""GradManager computes gradients or more generally, vector-Jacobian product, by reverse mode
    automatic differentiation (a.k.a. back propagation).

    Reverse mode autodiff normally reuses many intermediate tensors for best computation efficiency.
    In a read-eval-print-loop (REPL) environment however, it is impossible to known how the user
    would take gradients later thus which tensors to keep. To solve this problem, the user must
    somehow declare beforehand which gradient could possibly be taken. With GradManager, users are
    required to call the :meth:`attach` method on a tensor if they want to take gradients with
    respect to it later. Furthermore, any computation on a tensor before it is attached is
    completely ignored from the autodiff perspective, so :meth:`attach` must be called before any
    computation that needs differentiation.

    For example, the following symbolic differentiation code

    .. code-block::

        x = get_x()
        y = f(x)
        dy = ones_like(y)
        dx = vjp(y, x, dy) # vector-Jacobian product

    can be rewriten using GradManager for REPL environment as

    .. code-block::

        with GradManager() as gm:
            x = get_x()
            gm.attach(x) # must be placed before any computation on x that needs differentiation
            y = f(x)
            dy = ones_like(y)
            gm.backward(y, dy) # doesn't need x, already known via attach()
            dx = x.grad # backward() saves result to .grad attribute

    A more realistic example of training a neural network would be like

    .. code-block::

        gm = GradManager()
        gm.attach(model.parameters())

        for data in dataset:
            with gm:
                loss = model(data)
                gm.backward(loss)
            # gradients w.r.t. parameters is accumulated into their .grad attributes

    You can also use ``record()`` and ``release()`` method instead of ``with`` context:

    .. code-block::

        gm = GradManager()
        gm.attach(model.parameters())

        for data in dataset:
            gm.record()
            loss = model(data)
            gm.backward(loss)