在 Tensorflow 中使用列表推导式

Question

我们可以在 Tensorflow 图（非 eager）执行中使用 Python 理解吗？

IOU: tf.Tensor = tf.concat(            
    values=[
        intersection_over_union(
            box_pred[..., j, 1:5],     
            box_true[..., 1:5]         
        )
        for j in range(self.B)         # <--- predicted boxes
    ],
    axis=-1,
    name="IOU"
)

我认为 Graph 默认情况下不能有循环，需要使用

@tf.function

或

tf.while_loop

，但有些帖子使用了理解，因此想确认一下。

张量流中的等效列表理解

vals = [dict[tensor1[k]] for k in range(tensor1.get_shape().as_list()[0])]
tensor2 = tf.stack(vals, axis=0)

Answer 1

是的，确实如此。我提供了下面的代码来执行此操作：

# tested on TF 2.13.1
import tensorflow as tf
        
@tf.function
def list_comprehension(l):
    return tf.concat([tf.reshape(ll, (1,-1)) for ll in l], axis = 0)
l = [tf.random.uniform(shape=[5,2]) for _ in range(6)]

print(list_comprehension(l).shape)
# output (6, 10)

我定义了一个具有形状 (5,2) 张量的列表。在函数中，我更改了它们中每一个的形状，并将它们沿轴 0 连接起来。

如果列表理解很难追踪，是否有不使用列表理解的解决方案？更接近 TF ops 并且更优化的东西？

我在这篇post中问这个问题。