TensorFlow py_function 设置返回值形状答案

【问题标题】：Tensorflow py_function setting return value shapeTensorFlow py_function 设置返回值形状
【发布时间】：2021-03-10 13:59:08
【问题描述】：

我正在尝试设置从 tf.py_function 调用返回的 numpy 数组的形状。函数返回一个数组元组及其形状，所以我可以在tf.function 中设置它的形状。但是我不能在 set_shape 调用中使用返回的形状数组。我尝试将其转换为 int32 或通过索引等访问它，但所有结果都会导致 TypeError: Dimension value must be integer or None or have an __index__ method, got value '<tf.Tensor 'strided_slice:0' shape=<unknown> dtype=int32>' with type '<class 'tensorflow.python.framework.ops.Tensor'>' 错误。

@tf.function
def samples(img, mask) -> tuple:
    xs, xs_shape = tf.py_function(process,[128, 128], [tf.float32, tf.int32])
    xs.set_shape(tf.TensorShape(xs_shape))
    return xs

ds = ds.map(samples)

【问题讨论】：

标签： python tensorflow keras tensorflow2.0

【解决方案1】：

这是一个人为的例子。我认为这可能是您的要求。

def f(x):
    print("print ", x)
    # Include Logic using 'x' which is a tensor.
    # Sample  return values
    return [tf.convert_to_tensor([1,2]), tf.convert_to_tensor([1,2]).shape]


@tf.function
def samples(x) -> tuple:
    xs, xs_shape = f(x)
    xs.set_shape(xs_shape)
    print( xs_shape)
    return xs

dataset = tf.data.Dataset.range(1, 6)
dataset = dataset.map(  samples )

【讨论】：

这仍然会导致相同的错误。问题是使用 tf.py_function 调用 f 时会导致问题。我需要 py_function 因为我需要访问 numpy
简化代码在问题中。在您的示例中，如果您只是将 f(x) 调用更改为 tf.py_function 调用，您将得到相同的错误。 tf.py_function(f,[x], [tf.float32, tf.int32])
我不确定那是不是possible。 numpy 可以在我展示的函数中使用，但我不知道你的逻辑。