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)>]