TensorFlow 2.0：带有@tf.function 装饰器的函数不采用 numpy 函数

Question

我正在编写一个函数来在 TensorFlow 2.0 中实现一个模型。它需要 image_batch（一批 numpy RGB 格式的图像数据）并执行一些我需要的特定数据增强任务。导致我出现问题的行是：

@tf.function
def augment_data(image_batch, labels):
    import numpy as np
    from tensorflow.image import flip_left_right

    image_batch = np.append(image_batch, flip_left_right(image_batch), axis=0)

    [ ... ]

当我将 @tf.function 装饰器放在它上面时，

numpy 的 .append() 函数不再起作用。它returns:

ValueError: zero-dimensional arrays cannot be concatenated

当我在函数外部使用 np.append() 命令时，或者顶部没有 @tf.function 时，代码运行没有问题。

这正常吗？我是否被迫移除装饰器以使其工作？或者这是一个错误，因为 TensorFlow 2.0 仍然是测试版？那样的话，我该如何解决呢？

Answer 1

只需将 numpy ops 包装到 tf.py_function

def append(image_batch, tf_func):
    return np.append(image_batch, tf_func, axis=0)

@tf.function
def augment_data(image_batch):
    image = tf.py_function(append, inp=[image_batch, tf.image.flip_left_right(image_batch)], Tout=[tf.float32])
    return image

TensorFlow 2.0：带有@tf.function 装饰器的函数不采用 numpy 函数

TensorFlow 2.0: Function with @tf.function decorator doesn't take numpy functions

python

numpy

tensorflow

tensorflow2.0