MXNet学习之mx.sym.BlockGrad理解

MXNet学习之mx.sym.BlockGrad理解

1、解释

我们知道现在深度学习的框架是计算图，由节点和路径组成。在前向和反向的时候都是通过图路径传递的，那么这个函数是用在反向传播的时候，字面意思就是阻塞梯度传播，对某个节点不进行梯度计算。

2、示例

输入两个点，输出一个点，如下图所示：

2.1、正常反向传播

上图表示的公式是：y=3a+4b，在正常的反向传播时，a的偏导就是3，b的偏导为4（假设都是标量的情况下）。
代码测试：

v1 = np.array([[1, 2]])
v2 = np.array([[0, 1]])
a = mx.sym.Variable('a')
b = mx.sym.Variable('b')
b_grad = 4 * b
a_grad = 3 * a
loss = mx.sym.MakeLoss(b_grad + a_grad)
executor = loss.simple_bind(ctx=mx.cpu(), a=(1, 2), b=(1, 2))
executor.forward(is_train=True, a=v1, b=v2)
executor.backward()
print(executor.grad_arrays)
# [
# [[4. 4.]]
# <NDArray 1x2 @cpu(0)>,
# [[3. 3.]]
# <NDArray 1x2 @cpu(0)>]

2.2、阻塞传播

我们有时候需要阻塞一个op（操作）的反向，那么就需要利用BlockGrad函数。例如，我们阻塞上图中的v4节点，简单来说就是反向传播的时候，y->v5->v4->input此路不通了，也就没有相应的梯度值了。那么此时的结果应该是：a的偏导为3，b的偏导为0。

v1 = np.array([[1, 2]])
v2 = np.array([[0, 1]])
a = mx.sym.Variable('a')
b = mx.sym.Variable('b')
b_grad = 4 * b
b_stop_grad = mx.sym.BlockGrad(b_grad)
a_grad = 3 * a
loss = mx.sym.MakeLoss(b_stop_grad + a_grad)
executor = loss.simple_bind(ctx=mx.cpu(), a=(1, 2), b=(1, 2))
executor.forward(is_train=True, a=v1, b=v2)
executor.backward()

print(executor.grad_arrays[0])
# [[0. 0.]]
# <NDArray 1x2 @cpu(0)>

print(executor.grad_arrays[1])
# [[3. 3.]]
# <NDArray 1x2 @cpu(0)>

print(executor.grad_arrays)
# [
# [[0. 0.]]
# <NDArray 1x2 @cpu(0)>,
# [[3. 3.]]
# <NDArray 1x2 @cpu(0)>]]