Как правильно структурировать генератор ввода поезда / проверки в Tensorflow 2.1

Question

Я изо всех сил пытаюсь понять, почему мой генератор не работает. Я работаю в Tensorflow 2.1 с Tensorflow Transform (tfx), созданными осколками TFRecord. Приведенный ниже пример работает для моего конвейера tf.estimator.

def tfrecords_input_fn(files_name_pattern, transformed_metadata,
                       mode=tf.estimator.ModeKeys.EVAL,  
                       num_epochs=1, 
                       batch_size=500):
    features, domains = dataset_schema.schema_utils.schema_as_feature_spec(transformed_metadata.schema)
    dataset = tf.data.experimental.make_batched_features_dataset(
        file_pattern=files_name_pattern,
        batch_size=batch_size,
        features=features,
        reader=tf.data.TFRecordDataset,
        num_epochs=num_epochs,
        shuffle=True if mode == tf.estimator.ModeKeys.TRAIN else False,
        shuffle_seed = 36,
        shuffle_buffer_size=1+(batch_size*2),
        prefetch_buffer_size=1,
        reader_num_threads = 4,
        drop_final_batch = True
    )
    iterator =  tf.compat.v1.data.make_one_shot_iterator(dataset)
    features = iterator.get_next()
    target = features.pop(TARGET_FEATURE_NAME)
    return features, target

Я передаю это в TrainSpe c следующим образом:

train_spec = tf.estimator.TrainSpec(
          input_fn = lambda: tfrecords_input_fn(TRAIN_DATA_FILES,transformed_metadata,
            mode=tf.estimator.ModeKeys.TRAIN,
            num_epochs= hparams[NUM_EPOCHS],
            batch_size = BATCH_SIZE
          ),
          max_steps=hparams[MAX_TRAIN_STEPS],
        )

И затем в tf.estimator.train_and_evaluate (). Никаких проблем.

Я пытаюсь написать код для обучения пользовательской модели Keras без использования в настоящее время неисправной функции tf.keras.model_to_estimator ( не работает со слоем tf.keras.DenseFeatures ) , Я попробовал два следующих метода:

Метод 1:

    train_dataset_gen = lambda: tfrecords_input_fn(TRAIN_DATA_FILES,transformed_metadata,
        mode=tf.estimator.ModeKeys.TRAIN,
        num_epochs= hparams[NUM_EPOCHS],
        batch_size = BATCH_SIZE
        ) 

Аналогично с генератором проверочных наборов, а затем напрямую передан keras_model.fit (train_dataset_gen, .....).

Метод 2:

def create_input_fn(files_list, train=False,
                        batch_size=500,
                       num_epochs=1,
                   ):

    def tfrecords_input_fn(files_list, train=False, batch_size=500, num_epochs=1):
        dataset = tf.data.TFRecordDataset(files_list, num_parallel_reads=4)
        dataset.prefetch(tf.data.experimental.AUTOTUNE)
        if train == True:
            dataset = dataset.shuffle(buffer_size=TRAIN_SAMPLE_SIZE)
        dataset = dataset.repeat(num_epochs)
        dataset = dataset.batch(batch_size)
        return dataset

    return tfrecords_input_fn
# and then used like so:
train_dataset_gen = create_input_fn(TRAIN_DATA_FILES_LIST, train=True, batch_size=BATCH_SIZE, num_epochs=hparams[NUM_EPOCHS])
#validation gen created the same way...

#and then passed to keras .fit model 

keras_model.fit(train_dataset_gen, ....)

Оба метода в вызове fit дают мне:

/usr/local/lib/python3.5/dist-packages/tensorflow_core/python/keras/engine/training.py in fit(self, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_freq, max_queue_size, workers, use_multiprocessing, **kwargs)
    817         max_queue_size=max_queue_size,
    818         workers=workers,
--> 819         use_multiprocessing=use_multiprocessing)
    820 
    821   def evaluate(self,

/usr/local/lib/python3.5/dist-packages/tensorflow_core/python/keras/engine/training_v2.py in fit(self, model, x, y, batch_size, epochs, verbose, callbacks, validation_split, validation_data, shuffle, class_weight, sample_weight, initial_epoch, steps_per_epoch, validation_steps, validation_freq, max_queue_size, workers, use_multiprocessing, **kwargs)
    233           max_queue_size=max_queue_size,
    234           workers=workers,
--> 235           use_multiprocessing=use_multiprocessing)
    236 
    237       total_samples = _get_total_number_of_samples(training_data_adapter)

/usr/local/lib/python3.5/dist-packages/tensorflow_core/python/keras/engine/training_v2.py in _process_training_inputs(model, x, y, batch_size, epochs, sample_weights, class_weights, steps_per_epoch, validation_split, validation_data, validation_steps, shuffle, distribution_strategy, max_queue_size, workers, use_multiprocessing)
    531                      'at same time.')
    532 
--> 533   adapter_cls = data_adapter.select_data_adapter(x, y)
    534 
    535   # Handle validation_split, we want to split the data and get the training

/usr/local/lib/python3.5/dist-packages/tensorflow_core/python/keras/engine/data_adapter.py in select_data_adapter(x, y)
    996         "Failed to find data adapter that can handle "
    997         "input: {}, {}".format(
--> 998             _type_name(x), _type_name(y)))
    999   elif len(adapter_cls) > 1:
   1000     raise RuntimeError(

ValueError: Failed to find data adapter that can handle input: <class 'function'>, <class 'NoneType'>

Что не так с этими методами генератора? Он говорит, что получает класс 'function', класс 'NoneType', но мои генераторы возвращают именно то, что ожидается, по крайней мере, как я их интерпретирую:

Набор данных tf.data. Должен возвращать кортеж из (входы, цели) source

Я делаю ложные предположения о том, как работают генераторы Keras? мое возвращаемое значение от генератора недействительно? Как исправить? Я делаю это правильно?