mxnet 符号 API 是否支持条件控制流？

Question

我想在我的符号中添加一些条件控制，似乎 if-else 是在符号构建时评估的。但我希望它在符号运行时间内进行评估。

a = mx.symbol.Variable(name='a')
b = mx.symbol.Variable(name='b')

if a>b:
    c = a-b
else:
    c = a+b

TensorFlow提供了tf.cond()算子来处理，mxnet有对应的吗？

Answer 1

您可以使用 mx.symbol.where。

您可以计算 a_minus_b 和 a_plus_b 和 return 一个数组，其中每个元素来自 a_minus_b 或 a_plus_b，具体取决于中的相应值另一个 condition 数组。这是一个例子：

a = mx.symbol.Variable(name='a')
b = mx.symbol.Variable(name='b')

a_minus_b = a - b
a_plus_b  = a + b

# gt = a > b
gt = a.__gt__(b) 

result = mx.sym.where(condition=gt, x=a_minus_b, y=a_plus_b)

ex = result.bind(ctx=mx.cpu(), args={'a':mx.nd.array([1,2,3]), 'b':mx.nd.array([3,2,1])})
r = ex.forward()

print(r[0].asnumpy()) #result should be [1+3, 2+2, 3-1]

mxnet 符号 API 是否支持条件控制流？

Does the mxnet symbol API support conditional control flow?

mxnet