Multi-GPU training with Estimators, tf.keras and tf.data

At Zalando Research, as in most AI research departments, we realize the importance of experimenting and quickly prototyping ideas. With datasets getting bigger it thus becomes useful to know how to train deep learning models quickly and efficiently on the shared resources we have. TensorFlow's Estimators API is useful for training models in a distributed environment with multiple GPUs. Here, we'll present this workflow by training a custom estimator written with tf.keras for the tiny Fashion-MNIST dataset, and then show a more practical use case at the end. Note: there's also a cool new feature the TensorFlow team has been working on, (which at the time of writing is still in master), that lets you train a tf.keras model without first needing to convert it to an Estimator, with just a couple lines of additional code!