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