Tensorflow map_fn Error PartialTensorShape：合并期间的秩不兼容

Question

以下代码给了我一个我找不到答案的错误。我正在尝试将 python 函数应用于张量的每个元素，它将元素转换为形状为 3 的向量，因此我可以计算自定义评估指标。它需要是一个 Python 函数，因为它也在其他地方使用。

错误（下面的日志）是Invalid argument: PartialTensorShape: Incompatible ranks during merge: 1 vs. 0，我认为它与 map_fn 的结果及其形状有关。但是，它仅在运行时发生，就好像我有任何其他形状一样，然后在我执行 model.compile() 时会引发形状不兼容的错误。我是否误解了如何使用 map_fn？有什么建议吗？

提前致谢！

2021-04-09 12:19:31.357542: W tensorflow/core/framework/op_kernel.cc:1767] OP_REQUIRES failed at list_kernels.h:101 : Invalid argument: PartialTensorShape: Incompatible ranks during merge: 1 vs. 0
Traceback (most recent call last):
  File "test.py", line 93, in <module>
    validation_data=(val_input, val_output))
  File "/home/user/anaconda3/envs/tf_models/lib/python3.6/site-packages/tensorflow/python/keras/engine/training.py", line 108, in _method_wrapper
    return method(self, *args, **kwargs)
  File "/home/user/anaconda3/envs/tf_models/lib/python3.6/site-packages/tensorflow/python/keras/engine/training.py", line 1098, in fit
    tmp_logs = train_function(iterator)
  File "/home/user/anaconda3/envs/tf_models/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py", line 780, in __call__
    result = self._call(*args, **kwds)
  File "/home/user/anaconda3/envs/tf_models/lib/python3.6/site-packages/tensorflow/python/eager/def_function.py", line 840, in _call
    return self._stateless_fn(*args, **kwds)
  File "/home/user/anaconda3/envs/tf_models/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 2829, in __call__
    return graph_function._filtered_call(args, kwargs)  # pylint: disable=protected-access
  File "/home/user/anaconda3/envs/tf_models/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 1848, in _filtered_call
    cancellation_manager=cancellation_manager)
  File "/home/user/anaconda3/envs/tf_models/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 1924, in _call_flat
    ctx, args, cancellation_manager=cancellation_manager))
  File "/home/user/anaconda3/envs/tf_models/lib/python3.6/site-packages/tensorflow/python/eager/function.py", line 550, in call
    ctx=ctx)
  File "/home/user/anaconda3/envs/tf_models/lib/python3.6/site-packages/tensorflow/python/eager/execute.py", line 60, in quick_execute
    inputs, attrs, num_outputs)
tensorflow.python.framework.errors_impl.InvalidArgumentError: 2 root error(s) found.
  (0) Invalid argument:  PartialTensorShape: Incompatible ranks during merge: 1 vs. 0
         [[node map/TensorArrayV2Stack/TensorListStack (defined at test.py:27) ]]
         [[map_1/while/LoopCond/_50/_64]]
  (1) Invalid argument:  PartialTensorShape: Incompatible ranks during merge: 1 vs. 0
         [[node map/TensorArrayV2Stack/TensorListStack (defined at test.py:27) ]]
0 successful operations.
0 derived errors ignored. [Op:__inference_train_function_823]

Function call stack:
train_function -> train_function

这是使用 Tensorflow 2.3.1 和 Python 3.6 重现问题的代码。

from typing import List

import numpy as np

import tensorflow as tf
from tensorflow.keras.layers import Dense, Input, Flatten

INPUT_SHAPE = (2, 10, 10)


class CustomMetric(tf.keras.metrics.Metric):

    def __init__(self, name='custom_metric', **kwargs):
        super().__init__(name=name, **kwargs)
        self.mean_custom_metric = self.add_weight(name='mean_custom_metric', initializer='zeros', dtype=float)

    def update_state(self, y_true, y_pred, sample_weight=None):
        # y_true is a probability distribution (batch, 2*10*10), so find index of most likely position
        y_pred = tf.argmax(y_pred, axis=1)
        # y_pred and y_true are both tensors with shape (batch, 1)
        print(f"y_pred: {y_pred}")

        # apply python func to convert each value to a 3D value (single scalar to vector with 3 scalars)
        # according to docs: map_fn(fn, elems).shape = [elems.shape[0]] + fn(elems[0]).shape.
        # So: elems.shape[0] == batch | fn(elems[0]).shape == 3,
        # error happens when trying to do anything with the result of map_fn below
        y_true_positions = tf.map_fn(self.wrapper, y_true, fn_output_signature=tf.float32)
        y_pred_positions = tf.map_fn(self.wrapper, y_pred, fn_output_signature=tf.float32)
        # y_true_positions, y_pred_positions: tensors with shape (batch, 3)
        print(f"y_true_positions: {y_true_positions}")

        # do something with y_true_positions and y_pred_positions
        y_final = y_true_positions
        mean = tf.reduce_sum(y_final)

        print('---')
        self.mean_custom_metric.assign(mean)

    def result(self):
        return self.mean_custom_metric

    def reset_states(self):
        self.mean_custom_metric.assign(0.0)

    def wrapper(self, x):
        # x: tensor with shape (1,)
        print(f"x: {x}")

        result = tf.py_function(python_function, [int(x)], tf.float32)
        # result is a tensor of shape unknown
        print(f"result: {result}")
        result.set_shape(tf.TensorShape(3))
        # result: tensor with shape (3,)
        print(f"result: {result}")

        return result


def python_function(index: int) -> List[float]:
    # dummy function
    return [0, 0, 0]


# dummy model
block_positions = Input(shape=(*INPUT_SHAPE, 1), dtype=tf.float32)

block_positions_layer = Flatten()(block_positions)
target_output_layer = Dense(128, activation='relu')(block_positions_layer)

target_output = Dense(np.prod(INPUT_SHAPE), activation='softmax', name='regions')(target_output_layer)

model = tf.keras.models.Model(
    inputs=[block_positions],
    outputs=(target_output))

custom_metric = CustomMetric()
model.compile(
    loss=tf.keras.losses.SparseCategoricalCrossentropy(from_logits=False),
    optimizer=tf.optimizers.Adam(learning_rate=0.001),
    metrics=['accuracy', custom_metric])

print(model.summary())

# placeholder data
train_input = np.zeros(shape=(100, *INPUT_SHAPE), dtype=np.float32)
train_output = np.zeros(shape=(100, 1), dtype=np.int32)

val_input = np.zeros(shape=(100, *INPUT_SHAPE), dtype=np.float32)
val_output = np.zeros(shape=(100, 1), dtype=np.int32)

history = model.fit(
    train_input, train_output, epochs=10, verbose=1,
    validation_data=(val_input, val_output))

Answer 1

一段时间后我找到了解决方案。 wrapper 函数返回形状为 (3,) 的张量，而 map_fn 应用于形状为 (batch, 1) 的张量。我不完全明白为什么，但似乎 map_fn 需要一个形状为 (batch, 1,) 的返回张量，而不是文档建议的 fn(elems[0]).shape。

换行： result.set_shape(tf.TensorShape(3)) 为了 result = tf.reshape(tf.concat(result, 1), (1, 3)) 在wrapper 所以返回值是 (1, 3) 而不是 (3) 解决了这个问题。在map_fn 之后，你最终会得到一个形状为 (batch, 1, 3) 的张量，我将其重新整形为 (batch, 3)。