在具有 tf.while_loop (TF 2.0) 的函数上使用 tf.vectorized 映射的问题答案

【问题标题】：Issue with using tf.vectorized_map on a function with a tf.while_loop (TF 2.0)在具有 tf.while_loop (TF 2.0) 的函数上使用 tf.vectorized 映射的问题
【发布时间】：2021-01-21 00:56:54
【问题描述】：

包含的最小工作示例尝试向量化一个函数，该函数将 10.0 添加到输入，并通过每次添加 1.0（10 次）的 while 循环来实现。该函数在使用 tf.map_fn 时运行完美，在使用 tf.vectorized_map 时运行失败。

使用矢量化地图时该函数不会运行，并且错误指向Either add a converter or set --op_conversion_fallback_to_while_loop=True，可能运行速度较慢。我究竟做错了什么？注意：这是我问过https://github.com/tensorflow/tensorflow/issues/46559 的同一问题的重复。

if __name__ == "__main__":
    import tensorflow as tf
    @tf.function()
    def add(a):       
        i = tf.constant(0, dtype = tf.int32)
        c = tf.constant(1., dtype = tf.float32)
        loop_index = lambda  i, c, a: i < 10
        def body(i, c, a):
            a = c + a
            i = i + 1
            return i,c, a
        i,c, a = tf.while_loop(loop_index, body, [i,c, a],\
             shape_invariants=[tf.TensorShape(()), tf.TensorShape(()),tf.TensorShape([1])], back_prop= False, parallel_iterations=1)

        return a

    counter =  tf.reshape(tf.range(0, 40, delta = 1, dtype = tf.float32), shape = [40,1])

    all_ = tf.vectorized_map(add, counter) # does not work
    # all_ = tf.map_fn(add, counter) # works as expected
    print(all_, '<-- should be [40,1] float32 tensor with elements [10., 11., ...49.]')

【问题讨论】：

标签： python vectorization tensorflow2.0

【解决方案1】：

不幸的是，这个问题只能在 conda 环境中使用 TF2.2.0 + CUDA 10.1 重现，并且代码没有任何问题。解决方法是使用不同版本的 tensorflow 运行代码。这是一个协作笔记本，在问题受让人 ravikyram Link 的 github 问题中展示了相同的内容@

【讨论】：