GradManager

class 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 会回到非活跃的状态。

参数
record()[源代码]

开始记录操作

在此调用后,您将可以调用 backward

release()[源代码]

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

在此调用后,您将不能调用 backward