Passing a tf.data.Dataset that returns dicts to Keras

Viewed 1038

Background

As far as I can tell, Keras will become the preferred way of writing eager/lazy agnostic models (source). There also seems to be support for tf.data with Keras (source) which is great because that's the ETL I'm using for my current tf.estimator project.

Problem

I'd like to keep my tf.data pipeline but replace my tf.estimator with a tf.keras.Model, however after porting my graph I get tripped up by a

ValueError: Please do not pass a dictionary as model inputs.

upon calling model.fit(dataset, steps_per_epoch=n) which seems weird to me because tf.data/tf.estimator was built to work with (features: dict, labels: dict) where values in the dicts are tensors.

Summary

How can I use tf.keras.Model with dicts?

1 Answers

A temporary workaround is to work with tuples of tensors instead of dicts, for example by calling

def flatten(dataset):
    return dataset.map(lambda features, labels: (tuple(features.values()), tuple(labels.values())))

before model.fit, while making sure Model.call goes by indices instead of keys (e.g. return tuple(outputs.values()) etc.).

Unfortunately: The above hack only works with training data. For model.fit(validation_data=flatten(validation_dataset), ...) it won't work due to some internal bug: https://github.com/tensorflow/tensorflow/issues/21729

Related