MirroredStrategy: входные данные для активной функции исполнения не могут быть символами Keras * c тензоров

Question

Я использую tenorflow 2.0 для обучения модели, мой конвейер данных и модель тренируются нормально с tf.keras.Model.fit() на одном графическом процессоре, но когда я пытаюсь запустить его на мультипроцессоре с 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).