tensorflow2.0 - how to return multiple values in tensorflow tf.fun()

Question

I am using tfrecord for dataset. I want to do some augmentation for the image and label. But I cannot return both of them due to the graph mode, i think.

def encode_out(label):
    res = tf.py_function(encode_label,[label], tf.float32)
    return res

def encode_label(label):
    label = label.numpy()
    encoder = SSDInputEncoder(img_height=IMG_SIZE[0],
                            img_width=IMG_SIZE[1],
                            n_classes=N_CLASSES,
                            predictor_sizes=PREDICTOR_SIZES,
                            scales=SCALES,
                            aspect_ratios_per_layer=ASPECT_RATIOS_PER_LAYER,
                            steps=STEPS,
                            offsets=OFFSETS)
    return encoder(label)

label = encode_out(label)

for example I use encode_out to do some transformation for label, and it works well. However when dealing both image and label in one function, it reports an error.

def data_augument_out(image, label):
    image, label = tf.py_function(data_augument,
                                    [image, label], tf.float32)
    return image, label

def data_augument(image, label):
    image = image.numpy()
    label = label.numpy()

    augumentation = SSDDataAugmentation(img_height=IMG_SIZE[0], 
                                    img_width=IMG_SIZE[1])
    image, label = augumentation(image,label)
    return image, label

image, label = data_augument_out(image, label)

OperatorNotAllowedInGraphError: iterating over tf.Tensor is not allowed in Graph execution. Use Eager execution or decorate this function with @tf.function.

howver I can print the true value of image and label in data_augument function. so is there a way to return two tensor via tf.py_function

here is a simple code i think related to the question. in the for loop, i want it print two times.

import tensorflow as tf

@tf.function
def square_if_positive(x):
    for i in x:
        print('*********************',i)

a=tf.constant([2,3,4])
b=tf.constant([[2,3], [4,5], [5,6]])
square_if_positive(tf.tuple(a,b))

thank you

question from:https://stackoverflow.com/questions/66049750/how-to-return-multiple-values-in-tensorflow-tf-fun

Answer 1

From what I see you did not specify the output type for the label in your data_augment_out function and it can also help to set the shape.

I would try this:

def data_augment_out(image, label):
    im_shape = image.shape
    label_shape = label.shape
    [image, label,] = tf.py_function(data_augument, [image, label], [tf.float32, tf.float32])
    image.set_shape(im_shape)
    label.set_shape(label_shape)
    return image, label

There is a similar example here: https://www.tensorflow.org/guide/data#applying_arbitrary_python_logic