MXNet快速开始之自动微分

Basic usage

from mxnet import nd
from mxnet import autograd

示例f(x) = 2 x^2

x = nd.array([[1, 2], [3, 4]])
x

告知NDArray我们计划记录计算过程以计算梯度

x.attach_grad()

定义y=f(x)，用autograd.record()便于之后计算梯度

with autograd.record():
    y = 2 * x * x

x.grad
输出：
[[ 4.  8.]
 [12. 16.]]
<NDArray 2x2 @cpu(0)>

Using Python control flows

def f(a):
    b = a * 2
    while b.norm().asscalar() < 1000:
        b = b * 2
    if b.sum().asscalar() >= 0:
        c = b[0]
    else:
        c = b[1]
    return c

a = nd.random.uniform(shape=2)
a.attach_grad()
with autograd.record():
    c = f(a)
c.backward()

We know that b is a linear function of a, and c is chosen from b. Then the gradient with respect to a be will be either [c/a[0], 0] or [0, c/a[1]], depending on which element from b we picked. Let’s find the results:

[a.grad, c/a]
输出：
[
 [2048.    0.]
 <NDArray 2 @cpu(0)>, 
 [2048.      1008.60596]
 <NDArray 2 @cpu(0)>]