TensorFlow 2.1 makes Keras play nice with TPUs


The TensorFlow team has finished work on version 2.1 of the numerical computation library, with the result offering Keras enhancements and improvements for distributed training.

Amongst those are support for automatic data distribution and sharding in distributed environments for tf.data.Dataset, as well as ways to tune distribution policies in this abstraction and alterations for better performance. Apart from that, global distribution is now an option for tf.distribute.experimental_set_strategy() users, while those working with Google’s TPUs and TPU pods now have ways to create custom training loops available.

Since the team behind machine learning project Keras decided to concentrate on its TensorFlow module and doesn’t invest in the original project anymore, most enhancements can be found in tf.keras this time around. Most have to do with making it cooperate better with Cloud TPUs, which isn’t that surprising if you remember that TensorFlow was started by Google.

But there are also new functional additions, such as a TextVectorization layer which can be used for text standardisation, tokenisation, n-gram generation, and vocabulary indexing given raw string input. Reference implementations of popular models have been added to TF’s Model Garden, which should make life easier for Keras neophytes.


Changes in TensorFlow 2.1, that could lead to old TF code no longer working, include the removal of Operation.traceback_with_start_lines, and tf.config.experimental_list_devices. The latter can be replaced by tf.config.list_logical_devices. Devs using tf.config.experimentalVirtualDeviceConfiguration should alter that to tf.config.LogicalDeviceConfiguration after updating to the new version.

Windows users should be aware that the official pip packages of TF 2.1 have been built with a newer Visual Studio 2019 version in order to use a new compiler flag that reduces the build time. This means, however, that “Microsoft Visual C++ Redistributable for Visual Studio 2015, 2017 and 2019” has to be installed on a person’s system to use the packages otherwise a warning will appear.

As announced last year, TensorFlow 2.1 is the last release supporting Python 2, since official support for the language ceased on 1 January, 2020.

- Advertisement -