python之相当于在 tensorflow 中枚举以在 tf.map_fn 中使用索引

我有一个张量，我使用 tf.map_fn 逐行处理。现在我想将索引作为参数包含在传递给 tf.map_fn 的函数中。在 numpy 中，我可以使用 enumerate 获取该信息并将其传递到我的 lambda 函数中。这是 numpy 中的一个示例，其中我将 0 添加到第一行，将 1 添加到第二行，依此类推:

a = np.array([[2, 1], [4, 2], [-1, 2]]) 
 
def some_function(x, i): 
    return x + i 
 
res = map(lambda (i, row): some_function(row, i), enumerate(a)) 
print(res) 
 
> [array([2, 1]), array([5, 3]), array([1, 4])] 

虽然我无法在 tensorflow 中找到与 enumerate 等效的方法，但我不知道如何在 tensorflow 中获得相同的结果。有人知道用什么让它在 tensorflow 中工作吗？这是一个示例代码，其中我将 1 添加到 a 的每一行:

import tensorflow as tf 
 
a = tf.constant([[2, 1], [4, 2], [-1, 2]]) 
 
with tf.Session() as sess: 
    res = tf.map_fn(lambda row: some_function(row, 1), a) 
    print(res.eval()) 
 
> [[3 2] 
   [5 3] 
   [0 3]] 

感谢任何能帮我解决这个问题的人。

请您参考如下方法:

tf.map_fn()可以有很多输入/输出。因此，您可以使用 tf.range()构建行索引的张量并一起使用它:

import tensorflow as tf 
 
def some_function(x, i): 
    return x + i 
 
a = tf.constant([[2, 1], [4, 2], [-1, 2]]) 
a_rows = tf.expand_dims(tf.range(tf.shape(a)[0], dtype=tf.int32), 1) 
 
res, _ = tf.map_fn(lambda x: (some_function(x[0], x[1]), x[1]),  
                   (a, a_rows), dtype=(tf.int32, tf.int32)) 
 
with tf.Session() as sess: 
    print(res.eval()) 
    # [[2 1] 
    #  [5 3] 
    #  [1 4]] 

注意:在许多情况下，“逐行处理矩阵”可以一次完成，例如通过广播，而不是使用循环:

import tensorflow as tf 
 
a = tf.constant([[2, 1], [4, 2], [-1, 2]]) 
a_rows = tf.expand_dims(tf.range(tf.shape(a)[0], dtype=tf.int32), 1) 
res = a + a_rows 
 
with tf.Session() as sess: 
    print(res.eval()) 
    # [[2 1] 
    #  [5 3] 
    #  [1 4]]