TensorFlow 2.0: функция с декоратором @ tf.function не принимает функции numpy

Question

Я пишу функцию для реализации модели в TensorFlow 2.0.Он берет image_batch (пакет данных изображения в формате 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)

    [ ... ]

numpy Функция .append() больше не работает, когда я помещаю декоратор @tf.function поверх него.Он возвращает:

ValueError: невозможно объединить нульмерные массивы

Когда я использую команду np.append() вне функции или без @tf.function onсверху, код работает без проблем.

Это нормально?Я вынужден удалить декоратор, чтобы он работал?Или это ошибка из-за того, что 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