MirroredStrategy：到急切执行函数的输入不能作为符号张量

Question

我使用tensorflow 2.0来训练一个模型，我的数据管道和模型在单个gpu上用tf.keras.Model.fit()训练得很好，但是当我尝试在多个gpu上用MirroredStrategy运行它时

我的数据集是通过以下方式获得的：

dataset = tf.data.TFRecordDataset(records)
dataset.map(parse_and_decode)
map_augmentation = lambda *args: tf.py_function(
    augmentation,
    [*args],
    4*[tf.float32]
)
dataset.map(map_augmentation)

def map_ground_truth(arg1, arg2, arg3, arg4):
    ret_values = tf.py_function(
        get_ground_truth,
        [arg1, arg2, arg3, arg4],
        [tf.float32, tf.float32, tf.float32, tf.float32, tf.float32, tf.float32
    )
    # set_shape based on https://github.com/tensorflow/tensorflow/issues/24520
    for ret_val in ret_values:
        ret_val.set_shape([None for _ in range(3)])
    return *ret_values

dataset.map(map_ground_truth)

model.fit(dataset)

这在不使用MirroredStrategy的情况下可以正常工作，但当我引入它时，会出现以下错误：

_SymbolicException: Inputs to eager execution function cannot be Keras symbolic tensors, but found [<tf.Tensor 'crowd_mask:0' shape=(1, 401, 401, 1) dtype=float32>, <tf.Tensor 'unannotated_mask:0' shape=(1, 401, 401, 1) dtype=float32>, <tf.Tensor 'kp_maps_true:0' shape=(1, 401, 401, 17) dtype=float32>, <tf.Tensor 'mid_offsets/Identity:0' shape=(1, 13, 13, 64) dtype=float32>]

（在strategy.scope()中是模型和损失创建、模型compile和fit。数据集上没有对strategy.experimental_distribute_dataset进行手动调用，因为这是由fit函数单独处理的）