megengine.autodiff.grad_manager 源代码

import weakref
from typing import Callable, Iterable, List, Union

from ..core._imperative_rt.core2 import (
from ..core.autodiff.grad import Grad
from ..core.tensor.dtype import is_differentible_dtype
from ..logger import get_logger
from ..tensor import Tensor
from ..utils.future import Future

logger = get_logger(__name__)

backwarding_grad_manager = None

def get_backwarding_grad_manager():
    return backwarding_grad_manager

class AttachSpec:
    __slots__ = "tensor", "callbacks"

[文档]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) # backward() will clear recorded history and free resources # call release() if backward() is not called # gm.release() For your convenience, GradManager may (not must) be reused. As shown in the examples, you only need to attach a tensor once and GradManager will remember it afterwards. However, a single GradManager can record only one computation history at a time. To run multiple differentiations simultaneously or perform high order differentiation, create as many GradManager as you need. .. note:: Mutable tensors introduce ambiguities when doing symbolic differentiation: which version of the tensor are we referring to? For attached tensors, GradManager resolves this ambiguity by "snapshoting" them on first encounter, either on :meth:`record` (or entering with statement) if tensor is attached before :meth:`record`, or on :meth:`attach` if GradManager is already recording. Attached tensors will then be interpreted as their snapshotted version for differentiation purpose. The same ambiguity on the first parameter of :meth:`backward` is simply resolved by using the latest version. Typically, in data parallel, we would like to average the gradients across processes. Users will finally get the averaged gradients if an "AllReduce" callback is registered as follows: .. code-block:: import megengine.distributed as dist gm = GradManager() gm.attach(model.parameters(), callback=dist.make_allreduce_cb("MEAN")) """ def __init__(self): self._attach_specs = {} # id(Tensor) -> AttachSpec self._recording = False self._grad = None self._after_backward_callback = [] self._gradients = {}
[文档] def attached_tensors(self): r"""Return attached tensor list from :meth:`attach`.""" return [spec.tensor() for spec in self._attach_specs.values()]
[文档] def attach(self, tensors: Iterable[Tensor], callbacks=None): r"""Instruct GradManager to track operations on tensors, so that gradients with respect to those tensors could be evaluated later. :meth:`attach` also accepts a list of callbacks, which will be called with the tensor and its gradient during :meth:`backward`. The signature of callbacks should look like: .. code-block:: def callback(tensor: Tensor, grad: Tensor) -> Tensor: ... # returned grad is passed to subsequent callbacks # and finally accumulated to the .grad attribute of tensor return grad :meth:`attach` calls with overlapping tensors will result in their callbacks concatenated, independently for each tensor. For example, .. code-block:: gm.attach([x, y], callbacks=[f]) gm.attach([y], callbacks=[g]) is equivalent to .. code-block:: gm.attach([x], callbacks=[f]) gm.attach([y], callbacks=[f, g]) The effect of :meth:`attach` will persist across multiple uses of the GradManager. When reusing a GradManager, it is likely a mistake to call :meth:`attach` on the same set of tensors and callbacks repeatedly, which may grow the callback list indefinitely. .. note:: When reusing a GradManager, it is sometimes desirable to attach temporary tensors each time, e.g. for computing gradients of inputs of a neural network. GradManager tries to accommodate such usages by holding weak references to attached tensors. Most of the times, this should be enough to prevent resource leak. Unfortunately, there are still some pitfalls left: - Callbacks should not hold strong references, directly or indirectly, to attached tensors. Any strong reference, including those from callbacks, will prevent garbage collection (even by the cycle collector!) of a attached tensor, until the GradManager object is garbage collected. Please also note that GradManager might hold additional strong references to attached tensors when it is in use. This note only covers potential resource leaks across multiple uses of a GradManager, which is unrelated to whether resources is timely released within a single use. Args: tensors: tensor or list of tensors to track callbacks: callback or list of callbacks """ if callbacks is None: callbacks = [] if isinstance(callbacks, Callable): callbacks = [callbacks] if isinstance(tensors, Tensor): tensors = [tensors] def make_spec(tensor): selfref = weakref.ref(self) key = id(tensor) def deleter(_): self = selfref() if self is not None: del self._attach_specs[key] spec = AttachSpec() spec.tensor = weakref.ref(tensor, deleter) spec.callbacks = [] return spec for x in tensors: assert isinstance(x, Tensor), "Object to be attached should be Tensor" assert is_differentible_dtype(x.dtype), ( "Only tensors of floating point dtype can be attached to get gradients, " "get tensor dtype: {} and shape: {}".format(x.dtype, x.shape) ) spec = self._attach_specs.get(id(x)) new_attach = spec is None if spec is None: spec = make_spec(x) self._attach_specs[id(x)] = spec spec.callbacks.extend(callbacks) if new_attach and self._recording: self._do_record(spec) return self
def _register_after_backward_callback(self, callback): self._after_backward_callback.append(callback) return self
[文档] def backward( self, y: Union[Tensor, List[Tensor]] = None, dy: Union[Tensor, List[Tensor]] = None, ): r"""Compute gradients (or vector-Jacobian product) for all attached tensors, accumulate to corresponding .grad attribute, and release resources along the way. :meth:`backward` computes the vector-Jacobian product :math:`dx_j = \sum_{i} dy_i J_{ij}` where :math:`J_{ij} = ∂y_i/∂x_j` is the Jacobian matrix between vector variables :math:`y` and :math:`x`, with all vectors involved represented as a list of tensors, in the sense of direct sums (or flatten-and-concatenate). :math:`y` and :math:`dy` are passed as the first and second parameter respectively, whereas :math:`x` is directly taken from the list of all attached tensors. The result :math:`dx` is also not returned. Instead, it is directly accumulated into the .grad attribute of matching attached tensors (a.k.a. :math:`x`). This can be done unambiguously since :math:`dx` as a list of tensors has the same structure as :math:`x`. If :math:`y` is a scalar and :math:`dy` is chosen to be 1, the vector-Jacobian product yield gradient of :math:`y` with repect to :math:`x` as a special case. In that case, you will be able to omit the :math:`dy` parameter and :meth:`backward` will automatically use 1 for it and compute the gradient. :meth:`backward` consumes all resources held by this GradManager and releases them in the process of this call. When the call successfully finishes, the GradManager will be put back to an inactive state. Args: y: tensor or list of tensors dy: tensor or list of tensors. Defaults to 1 if y is scalar """ push_scope("backward") set_option("record_computing_path", 0) _origin_auto_format = get_auto_format_convert() from ..functional import ones_like global backwarding_grad_manager cache = backwarding_grad_manager backwarding_grad_manager = self if not self._recording: raise RuntimeError( "no computation history. " "did you forget record() or " "call a method that clears the history?" ) assert self._grad is not None # These checks should be consistent with GradScaler's if y is None: ys = [] elif isinstance(y, (tuple, list)): ys = y else: ys = [y] if dy is None: dys = [ones_like(y) for y in ys] elif isinstance(dy, (tuple, list)): dys = dy else: dys = [dy] try: self._grad(ys, dys) for callback in self._after_backward_callback: callback() for id_, grad in self._gradients.items(): if isinstance(grad, Future): grad = grad.get() spec = self._attach_specs.get(id_) tensor = spec and spec.tensor() if tensor is not None: if tensor.grad is None: tensor.grad = grad else: tensor.grad += grad finally: self.release() backwarding_grad_manager = cache set_option("record_computing_path", 1) pop_scope("backward")
[文档] def record(self): r"""Start recording operations After this call, you will be able to call :meth:`backward`. """ if self._recording: raise RuntimeError("already recording") grad = Grad() self._recording = True self._grad = grad grad.__enter__() for spec in self._attach_specs.values(): self._do_record(spec)
def _do_record(self, spec): tensor = spec.tensor() if tensor is None: return def callback(grad, callbacks=spec.callbacks): from ..functional import ones_like for cb in callbacks: grad = cb(tensor, grad) self._gradients[id(tensor)] = grad # NOTE: override prev callback wrt when called serval times self._grad.wrt(tensor, callback=callback)
[文档] def release(self): r"""Stop recording operations and release resources kept for gradient computation After this call, you will not be able to call :meth:`backward`. """ if self._grad is not None: self._grad.__exit__(None, None, None) self._grad = None self._recording = False self._gradients = dict()
def __enter__(self): self.record() return self def __exit__(self, exc_type, exc_val, exc_tb): self.release() def __or__(self, other): if isinstance(other, GradManager): return GradManagerGroup([self, other]) return NotImplemented __ror__ = __or__
class GradManagerGroup: def __init__(self, gms) -> None: self._gms = list(gms) def merge_with(self, other): if isinstance(other, GradManager): other = GradManagerGroup([other]) elif not isinstance(other, GradManagerGroup): return NotImplemented return GradManagerGroup([*self._gms, *other._gms]) __or__ = merge_with __ror__ = merge_with def __enter__(self): Grad.stack.append([]) Grad.begin_group() for gm in self._gms: gm.record() assert gm._grad is not None Grad.end_group() def __exit__(self, exc_type, exc_val, exc_tb): for gm in reversed(self._gms): gm.release() assert gm._grad is None