megengine.autodiff 模块

megengine.autodiff.grad_manager

class megengine.autodiff.grad_manager.AttachSpec[源代码]

基类:object

__slots__ = ('tensor', 'callbacks')
callbacks
tensor
class megengine.autodiff.grad_manager.GradManager[源代码]

基类:object

GradManager 负责计算梯度,或更一般地,矢量-雅可比积,通过反向模式做自动微分(又称为反向传播)。

反向模式自动微分为了最佳计算效率通常会重用许多中间张量。然而在交互式(REPL)环境中,很难知道用户之后将如何使用梯度从而预先保存某些张量。要解决此问题,用户必须以某种方式事先声明什么梯度需要被保留。在 GradManager 中,用户如果希望后续计算张量的梯度,则需要在该张量上调用 attach() 方法。此外,在张量被 attach 之前对张量的任何计算都会从 autodiff 的角度完全被忽略,所以 attach() 必须在需要微分的任何计算之前被调用。

例如,下面的自动微分代码

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

可用 GradManager 在交互式环境下重写微

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

训练神经网络的一个更现实的例子是

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

你也可以使用 record()release() 而不使用 with 语义:

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()

为方便起见,可以(并非必须)重用 GradManager 。如在这些示例中,您只需要 attach 一个张量一次,而 GradManager 将永远记住它。但是,一个 GradManager 只能记录一次计算历史。要同时进行多种微分或进行高阶微分,可根据需要创建尽可能多的 GradManager

注解

可变张量在执行符号微分时引入歧义:我们指的是张量的哪个版本?对于 attched 的张量,GradManager 通过在第一次遇见张量时 ‘snapshoting’ 它们来解决这种歧义,要么通过 record()(或者输入 ‘statement’),若张量是在 record() 之前就被 attach(),要么通过 attach() 如果 GradManager 已经在记录。attched 的张量将会被解释为它们的快照版本以进行区分。对 backward() 的第一个参数的相同歧义是只需使用最新版本即可解决。

通常,在数据并行的时候,我们需要跨进程计算梯度的均值。使用者最终可获得平均的梯度若一个 “AllReduce” 回调函数像下面一样被注册的话:

import megengine.distributed as dist

gm = GradManager()
gm.attach(model.parameters(), callback=dist.make_allreduce_cb("MEAN"))
attach(tensors, callbacks=None)[源代码]

指示 GradManager 跟踪张量上的操作,以便那些张量上的梯度,可以在之后进行计算。

attach() also accepts a list of callbacks, which will be called with the tensor and its gradient during backward(). The signature of callbacks should look like:

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

attach() 对于重叠张量的调用将导致其回调函数针对每个张量独立地级联被调。例如,

gm.attach([x, y], callbacks=[f])
gm.attach([y], callbacks=[g])

等价于

gm.attach([x], callbacks=[f])
gm.attach([y], callbacks=[f, g])

attach() 的影响将在多次使用 GradManager 时持续存在。当重用一个 GradManager,在同一张量和回调上重复调用 attach() 可能是错误的,这可能会无限期地增加回调列表。

注解

重用一个 GradManager 时,有时希望每次都 attach 临时张量,例如,用于计算神经网络输入的梯度。GradManager 尝试通过保持对 attached 张量的弱引用来适应此类用法。大多数时候,这足以防止资源泄漏。 不幸的是,仍然存在一些陷阱:

  • 回调不应直接或间接对 attached 张量持有强引用,任何强引用,包括来自于回调函数的,都会阻止``attached`` 张量的垃圾回收(甚至是通过循环收集器),直到 GradManager 对象被垃圾收集为止。

还请注意,GradManager 使用过程中可能包含附加到 attached 张量的多余强引用。本说明仅覆盖多次使用同一个 GradManager 时潜在的资源泄漏,这与是否在一次使用中及时释放资源无关。

参数
  • tensors (list) – 将追踪的张量或者其列表

  • callbacks – 回调函数或其列表

backward(y=None, dy=None)[源代码]

为所有连接的张量计算梯度(或矢量-雅可比积),累加到相应的 .grad 属性中,并沿途释放相应资源。

backward() 计算向量-雅各比积: \(dx_j = \sum_{i} dy_i J_{ij}\),其中 \(J_{ij} = ∂y_i/∂x_j\) 是张量 \(y\)\(x\) 之间的雅各比矩阵,所涉及的所有张量都表示为张量列表, 通过直接或者拼接的方式。\(y\)\(dy\) 分别作为第一和第二个参数被传递,而 \(x\) 直接从所有被 attached 的张量列表中获取。 结果 \(dx\) 不会被返回.。相反,它直接被加到对应张量(也称为 \(x\))的 .grad 属性中。这个过程没有歧义因为 \(dx\) 作为一个张量列表有着与 \(x\) 一样的结构。

如果 \(y\) 是一个标量,然后 \(dy\) 被选为1,则作为特例,向量-雅可比积推出 \(y\) 相对于 \(x\) 的梯度。在这种情况下,您可以省略 \(dy\) 参数并且 backward() 会自动使用1,然后计算梯度。

backward() 使用所有当前 GradManager 的资源并在当前调用进程中释放它们。调用完成后,GradManager 将回到非活动状态。

参数
  • y – 张量或者张量列表

  • dy – 张量或者张量列表。默认为1如果 y 是标量的话。

record()[源代码]

开始记录前向运算。

调用此函数后,可继续调用 backward()

release()[源代码]

停止记录并释放用于梯度计算的资源。

调用此函数后,将无法调用 backward()

megengine.autodiff.grad_manager.get_backwarding_grad_manager()[源代码]