Tensorflow 将 op 应用于 2d 张量的每个元素

Question

我所追求的是能够将 tensorflow op 应用于 2d 张量的每个元素，例如

input = tf.Variable([[1.0, 2.0], [3.0, 4.0])
myCustomOp = ... # some kind of custom op that operates on 1D tensors
finalResult = tf.[thing I'm after](input, myCustomOp)
# when run final result should look like: [myCustomOp([1.0, 2.0]), myCustomOp([3.0, 4.0)]

有任何想法吗？

Answer 1

TensorFlow 的下一个版本(0.8，如果您从源代码构建或下载每晚构建，则当前可用)包括 higher-order operators包括 tf.map_fn() 和 tf.scan() 允许您将一个由 TensorFlow ops 组成的函数应用于更大张量的子张量。
tf.map_fn(fn, elems, ...)函数解包 N维输入elems沿第一个维度分成多个 N-1维子张量和应用 fn到每个子张量。这似乎非常适合您的用例:

input = tf.Variable([[1.0, 2.0], [3.0, 4.0]])
function_to_map = lambda x: f(x)  # Where `f` instantiates myCustomOp.
final_result = tf.map_fn(function_to_map, input)