GradManager¶

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。

注解

可变张量在执行符号微分时引入歧义：我们指的是 Tensor 的哪个版本？对于 attched 的张量，GradManager 通过在第一次遇见张量时 ‘snapshoting’ 它们来解决这种歧义，要么通过 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 也接受一个回调函数的列表，在 backward 的过程中，这些回调函数会被以 Tensor 和其梯度作为参数调用。回调函数的签名应该如下：

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 调用的 Tensor 列表如果有重叠，那么这些 Tensor 对应的回调函数列表会被拼接，此操作对每个 Tensor 独立作用。例如，

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 的同时，您可能会希望对一些临时的 Tensor 求导，例如对神经网络的输入求导。考虑到这种用法，GradManager 会对 attached Tensor 持有弱引用。在大多数时候，这已经可以避免资源泄漏。但是，仍然有少数情况需要您注意：

回调函数不应该持有被 attached Tensor 的强引用，无论是直接地还是间接地。任何强引用，包括来自回调函数的强引用，都会导致 attached Tensor 无法被垃圾回收（即使运行可回收引用循环的完整垃圾回收！），直到 GradManager 对象本身被垃圾回收为止。

还需注意的一点是 GradManager 如果正在进行求导，可能会持有对被 attached Tensor 的强引用。本注解仅针对将一个 GradManager 用于多次求导可能引发的资源泄漏，并不涉及在进行一次求导时资源是否被第一时间释放的问题。

参数

tensors (Iterable[Tensor]) – 需要跟踪的 Tensor 或者 Tensor 列表
callbacks – 回调函数或回调函数的列表

attached_tensors()[源代码]¶: 返回通过 attach 方法得到的已绑定张量列表。

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

对所有 attached Tensor 计算梯度 (或向量-雅可比积)，并将其积累到对应 Tensor 的 .grad 属性中，在过程中释放资源。

backward 计算向量-雅可比积 \(dx_j = \sum_{i} dy_i J_{ij}\) ，其中 \(J_{ij} = ∂y_i/∂x_j\) 是向量变量 \(y\) 和 \(x\) 的雅可比矩阵，涉及的所有向量都在向量空间直和 (或展平并拼接) 的意义下表示为 Tensor 的列表。\(y\) and \(dy\) 分别是第一和第二个参数，\(x\) 则由所有 attached Tensor 的列表形成。计算的结果 \(dx\) 不作为返回值，而是直接积累到对应 Tensor (即 \(x\)) 的 .grad 属性上。这个对应可以无歧义地完成，因为 \(dx\) 作为 Tensor 的列表和 \(x\) 有着相同的结构。

如果 \(y\) 是一个标量而 \(dy\) 被取为 1，那么向量-雅可比积的结果恰好就是 \(y\) 对 \(x\) 的梯度。在这种情况下，您可以省略 \(dy\) 参数，backward 会自动使用 1 作为它的值并计算梯度。

backward 会随计算的进行释放 GradManager 持有的所有资源。在函数调用成功返回后，GradManager 会回到非活跃的状态。

参数

y (Union[Tensor, List[Tensor], None]) – Tensor 或 Tensor 的列表
dy (Union[Tensor, List[Tensor], None]) – Tensor 或 Tensor 的列表。如果 y 是标量，有默认值 1

record()[源代码]¶

开始记录操作

在此调用后，您将可以调用 backward。

release()[源代码]¶

停止记录操作并释放为计算梯度而保留的资源

在此调用后，您将不能调用 backward。

Function

megengine.optimizer