TensorFlow 2.2 kicks Python 2 to the curb, adds new profiler and tf.distribute improvements


Machine learning platform TensorFlow has added a new profiler to its project and looked into ways of improving distributed tasks and model training in Keras, making version 2.2 ready for downloading. 

The standout feature of the new release seems to be a new TensorFlow 2 profiler. It is meant to help with analysing performance on CPUs, GPUs, TPUs and combinations thereof, so that users have an easier time identifying performance bottlenecks. Optimisation guidance “whenever possible” is also part of the new tool.

One way to speed up the training process would be to distribute computations across devices, which can be done by using the tf.distribute module that also saw some progress in this release. Working with GPUs, specifically, is said to have been improved through updating the version of NCCL used, and adding support for gradient allreduce in float16 as well as CompositeTensor in DistributedIterators. 

The module also comes with an implementation of all reduce gradient packing, which is experimental for now, while the run method experimental_run_v2 has graduated preview stage and has been renamed to run.


Although the TensorFlow team has been quite busy, most enhancements can be found in the Keras module, which has been the official place for the advancement of the deep learning library since September 2019. Back then, the Keras team released its last iteration of multi-backend Keras, encouraging users to switch to the better cultivated tf.keras. Maintenance for the independent project ended just last month. 

With more focus on just this module, the tf.keras devs were able to improve the model training process a fair bit, allowing users for example to use custom logic with Model.fit. They also don’t have to manually call Model._set_inputs anymore when using custom training loops, since SavedModel now uses its own attribute instead of relying on Model.inputs and Model.input_names.

Starting with this release, metrics, preprocessing layers, and other built-in Keras layers are supported in the SavedModel format. Users have to be aware, though, that the top layer in tf.keras.applications is now called predictions by default, which could break code relying on the name being something else. Also, the Huber loss function has been reworked for consistency with other functions, so results may not compare well with earlier computations.

As most Python related projects, TensorFlow took the opportunity of the new release to drop Python 2 support, which hit its end of life on 1 January, 2020. Since all TensorFlow Docker images now use Python 3, those containing p3 in their tag will no longer be updated or canned altogether.

TensorFlow 2.2 requires gast 0.3.3 for the platform to run and a minimum bazel version 2.0. XLA_CPU and XLA_GPU have been deprecated with this release. More details on breaking changes and bug fixes can be found in the project repository.

- Advertisement -