如何在tensorflow tf.fun()中返回多个值

Question

我使用tfrecord作为数据集。我想对图像和标签做一些增强。但是由于图形模式的关系，我想我不能同时返回它们。

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)

以为例，我使用encode_out对标签进行一些转换，效果很好。但是，当在一个函数中同时处理图像和标签时，它会报告一个错误。

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)

operatornotallowedingraphherror:在图形执行中不允许迭代tf.Tensor。使用即时执行或使用@tf.function修饰此函数。

但是我可以在data_augument function中打印image和label的真实值。有没有办法通过tf.py_function

返回两个张量这里有一个简单的代码，我认为与这个问题有关。在for循环中，我希望它打印两次。

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))

谢谢

Answer 1

我看到您没有在data_augment_out函数中指定标签的输出类型，它也可以帮助设置形状。

我想试试这个:

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

这里有一个类似的例子:https://www.tensorflow.org/guide/data#applying_arbitrary_python_logic