Autodiff 高阶使用

v1.5 开始，MegEngine 提供了实验性的高阶微分支持，允许多个GradManager进行协作。使用场景介绍如下。 你可以利用这个特性， 实现高阶微分、前向微分和多个一阶微分

高阶微分

对于 GradManager 来说，它只关心并如实的记录在启用范围内的全部 Tensor 计算操作，并不会区分这个 Tensor 是通过 forward 计算得到还是另一个 GradManager 求导得到（即求导器可以对 x.grad 再次进行求导）。 通过嵌套两层 GradManager ，我们可以实现求二阶导的效果。

# 这里想求x.grad对 x 的梯度：grad(x, grad(x, y))
def test_2nd_grad_with_manager():
    x_np = np.random.rand(10).astype("float32")
    x = mge.tensor(x_np)

    gm = GradManager().attach([x])       # 外层求导器，负责求二阶导
    gm2 = GradManager().attach([x])      # 内层求导器，负责求一阶导

    with gm:                             # gm 在 gm2外部，因此gm可以观测到gm2的backward行为
        with gm2:
            y = F.cos(x)                 # forward
            gm2.backward(y)              # backward，此时得到一阶导
        x_grad = x.grad                  # 保存一阶导
        x.grad = None                    # 将导数清空（否则后续会累计上来）
        np.testing.assert_almost_equal(  # 确认一阶导的正确性
            x_grad.numpy(), -np.sin(x_np), decimal=5
        )

        gm.backward(x_grad)  # 求一阶导 x_grad 对 x 的导数，即二阶导，结果会存储在 x.grad 上

    np.testing.assert_almost_equal(       # 确认二阶导的正确性
        x.grad.numpy(), -np.cos(x_np), decimal=5
    )

从外部gm的视角来看，它并不关心求导对象是在 backward 阶段产生还是从 forward 阶段产生的，只要是在 gm scope 内的计算都会被记录下来，如下图所示。

求二阶导时，需要注意嵌套顺序

我们调换 gm 与 gm2 的顺序，可以看到，在 gm.backward() 之后， x.grad 为空，也就是说 gm 没有求出 x 对 x.grad 的梯度。这是因为此时 gm 被放在内层，无法观测到 gm2.backward 这个行为。因此在 gm 眼里 x.grad 就好似凭空产生的（和 x.grad = Tensor(3) 一样），当然无法对其求导。

# expect grad(x, grad(x, y)), got None
def test_2nd_grad_fail_with_manager():
    x_np = np.random.rand(10).astype("float32")
    x = mge.tensor(x_np)

    gm = GradManager().attach([x])
    gm2 = GradManager().attach([x])

    with gm2:      # 注意此处
        with gm:   # 此时 gm 与 gm2 位置互换，放在内层
            y = F.cos(x)
            x_grad = gm2.backward(y)
            x.grad = None
            np.testing.assert_almost_equal(x_grad.numpy(), -np.sin(x_np), decimal=5)
            gm.backward(x.grad)

    assert x.grad is None

多次反向微分来模拟前向微分

# expect grad(oup_grad, grad(inp, oup, oup_grad))
def test_emulate_forward_mode_with_reverse_mode(target):
    expr = F.cos

    def jvp(inp, expr):
        with GradManager() as gm:                         # 这里不知道输出的形状，只能推迟attach
            with GradManager().attach([inp]) as gm2:
                oup = expr(inp)
                oup_grad = F.ones_like(oup)               # 随便给一些值都行
                gm.attach(oup_grad)                       # 等下要拿oup_grad.grad
                gm2.backward(oup, oup_grad)               # 算出对应的inp_grad，这里oup_grad参与了计算，等下可以求它的梯度
            gm.backward(inp.grad)                         # 算出oup_grad.grad
        return oup, oup_grad.grad

    def numerical_jvp(inp, expr):
        delta = 0.001
        return expr(inp), (expr(inp + delta) - expr(inp - delta)) / (2 * delta) # 求出近似值

    x_np = np.random.rand(10).astype("float32")
    x = mge.tensor(x_np)
    y, dy = jvp(x, target)
    y1, dy1 = numerical_jvp(x, target)

    np.testing.assert_almost_equal(y.numpy(), y1.numpy(), decimal=5)
    np.testing.assert_almost_equal(dy.numpy(), dy1.numpy(), decimal=3)

求两个一阶导

如果你希望在一次计算中，分别对不同的一批参数求导，那么你可以使用 gm | gm2 的写法，可分别求两个一阶导。 和 with gm: with gm2 写法相比，可以在不求高阶导的形式下节省一点代码，此外这两种写法其实语义上有微妙的差别： gm | gm2 中两个求导器是平等的（没有内外层之分），自然也是相互不可见的，有时可以避免一些意料之外的行为。

# expect grad(x, y) + grad(x, y)
def test_2nd_grad_with_manager_union():
    x_np = np.random.rand(10).astype("float32")
    x = mge.tensor(x_np)

    gm = GradManager().attach([x])
    gm2 = GradManager().attach([x])

    with gm | gm2:      # 这里代表将gm与gm2取并，即同时挂上这两个求导器
        y = F.cos(x)
        gm.backward(y)
        gm2.backward(y)

    np.testing.assert_almost_equal(x.grad.numpy(), -2 * np.sin(x_np), decimal=5) # 得到双份导数（因为会累加），每个gm贡献一半

此时两个 GradManager 观测到的内容如下所示：

gm | gm2 中两个求导器是平等的（没有内外层之分），相互是不可见的。因此，在下面的代码中由于 gm2 对 gm 不可见，导致无法求出 x 对于 x_grad 的梯度。

# expect grad(x, grad(x, y)), got None
def test_2nd_grad_fail_with_manager_union():
    x_np = np.random.rand(10).astype("float32")
    x = mge.tensor(x_np)

    gm = GradManager().attach([x])
    gm2 = GradManager().attach([x])

    with gm | gm2: #gm与gm2是同级的，并且互不可见
        y = F.cos(x)
        gm2.backward(y)
        x_grad = gm2.backward(y)
        x.grad = None
        np.testing.assert_almost_equal(x_grad.numpy(), -np.sin(x_np), decimal=5)
        gm.backward(x_grad)
        assert x.grad is None