TensorFlow 1.8.0 Launches with Powerful Multi-GPU Training via MirroredStrategy

Introduction

TensorFlow has become one of the most popular and widely used open-source libraries for machine learning. Developed by Google Brain and released in 2015, TensorFlow provides a flexible ecosystem of tools, libraries, and community resources to push the boundaries of ML research and bring cutting-edge models into production.

The TensorFlow team‘s commitment to innovation was on full display at the 2018 TensorFlow Developer Summit. Over the course of the two-day event, Googlers and community contributors showcased exciting new features for eager execution, a simplified Keras API, and accelerated training pipelines.

Just one week later, the team delivered on those promises with the release of TensorFlow 1.8.0. While technically a minor point release, the 1.8 update delivers major enhancements that address pain points for researchers and practitioners alike. Let‘s dive into the features that make TensorFlow 1.8 a must-have upgrade.

Distributed Training Made Easy with MirroredStrategy

The headline feature of TensorFlow 1.8 is the new tf.contrib.distribute.MirroredStrategy API for training models on multiple GPUs. Utilizing the power of distributed training has traditionally been one of the biggest challenges for TensorFlow users. Configuring multi-GPU environments required significant code changes and expertise to manage the intricacies of device assignment, gradient averaging, and batch splitting.

MirroredStrategy drastically simplifies the process by providing an abstraction layer on top of TensorFlow‘s existing distributed runtime. With just a few lines of code, you can now train your model on all available GPUs on a single machine:

mirrored_strategy = tf.contrib.distribute.MirroredStrategy()

with mirrored_strategy.scope():
  model = tf.keras.Sequential([
    tf.keras.layers.Dense(16, activation=‘relu‘, input_shape=(10,)),
    tf.keras.layers.Dense(1, activation=‘sigmoid‘)
  ])

  model.compile(loss=‘binary_crossentropy‘, optimizer=‘adam‘)

When you call fit() on a model defined within a MirroredStrategy scope, the training is automatically distributed across the specified devices. MirroredStrategy takes care of splitting each batch of input data, replicating the model on each device, and combining the gradients using an all-reduce algorithm.

The result is effortless speedup from parallelism without sacrificing model performance. The TensorFlow team published benchmarks showing near-linear scaling for a variety of models on up to 8 Tesla V100 GPUs:

Source: TensorFlow blog

MirroredStrategy currently supports synchronous distributed training on a single machine. The TensorFlow team is actively working on a more general DistributionStrategy API that will enable a variety of synchronous and asynchronous distributed training configurations for both single-machine and multi-machine clusters.

Eager Execution Enhancements

TensorFlow‘s eager execution mode continues to mature, gaining features that make it more suitable for production workloads in the 1.8 release. Eager execution provides an intuitive imperative programming model, ideal for rapid development and debugging.

One key addition is support for automatic device placement. When running eagerly, TensorFlow will now automatically assign ops to the available CPU or GPU devices. For example:

import tensorflow as tf
import tensorflow.contrib.eager as tfe

tf.enable_eager_execution()

with tf.device(‘/gpu:0‘):
  x = tf.random_normal([1000, 1000])

y = tf.random_normal([1000, 1000])

z = tf.matmul(x, y)  # Runs on GPU

Here the x and y tensors are created on the CPU, but TensorFlow is smart enough to place the matmul op on the GPU to optimize performance. Eager execution also now supports manual device placement using a tf.device scope as shown above.

Another notable eager enhancement is improved integration with tf.data and tf.keras. The tf.data.Dataset API now fully supports eager execution, enabling features like iterating over datasets and creating multiple iterators. There are also a number of new tf.keras layers and utilities with eager support.

The TensorFlow team‘s efforts to optimize eager performance have paid off as well. A new set of microbenchmarks published alongside the 1.8 release demonstrate significant speedups compared to graph execution mode for many common building blocks:

Source: TensorFlow blog

For models that fit in memory on the accelerator, eager execution with XLA JIT compilation now matches or exceeds the performance of graph execution. When combined with the developer productivity benefits, the 1.8 release makes eager mode compelling for an even wider variety of real-world applications.

tf.data and tf.keras Improvements

The tf.data and tf.keras modules continue to evolve in TensorFlow 1.8, gaining features that make them more flexible and performant. Several new tf.data.Dataset methods enable optimized GPU input pipelines and simplified ETL from various data formats:

• tf.contrib.data.prefetch_to_device allows you to prefetch dataset elements directly to GPU memory, improving training throughput
• tf.contrib.data.make_csv_dataset provides an easy way to create datasets from CSV files
• tf.contrib.data.AUTOTUNE argument can be passed to map, prefetch, batch, and other transformation methods to automatically tune parallelism parameters

tf.keras, TensorFlow‘s implementation of the popular Keras API, also gets upgrades in this release. New application-specific preprocessing layers make it simple to build tf.keras input pipelines for image, time series, and text data.

There are also major improvements to model exporting in tf.keras with full support for the SavedModel format. This enables tf.keras models to be deployed with TensorFlow Serving and other production environments.

Rapid Release Cadence Draws Mixed Reactions

With four releases in the last three months, TensorFlow has dramatically accelerated its development cycle. Some in the community have raised concerns about the cognitive overhead of keeping up with such a fast-moving target. A few frustrated users have even threatened to switch to competing frameworks like PyTorch.

In my view, the TensorFlow team has done an admirable job responding to these concerns by designating some releases as "long-term support" versions. These LTS releases, starting with TensorFlow 1.7, will receive extended support and compatibility guarantees for at least one year.

For users that want maximum stability, the LTS releases provide a safe choice. Those on the cutting edge can update with each new release to stay on top of the latest features. This hybrid model balances the needs of both research and production use cases.

It‘s also worth noting that TensorFlow is far from the only open-source library moving at such a rapid clip. The PyTorch team has had four releases in a similar timeframe, and MXNet published eleven releases in 2017. As deep learning matures, we should expect all of the major frameworks to iterate and evolve quickly.

Conclusion

TensorFlow 1.8 is a feature-packed release that pushes the framework forward on multiple fronts. The MirroredStrategy API democratizes multi-GPU training, eager execution gains production-grade functionality, and Keras and tf.data continue to evolve.

While some have criticized the release velocity, there are valid arguments on both sides. For many organizations and users, access to cutting-edge features is well worth the upgrade overhead. The LTS releases provide stability for those that need it.

From my perspective, the 1.8 release is a milestone that showcases the strength of TensorFlow‘s ecosystem. The features that landed in this version are the result of close collaboration between Google and the broader open-source community.

I‘m personally excited to put the new MirroredStrategy feature to work in my own projects. The complexity of multi-GPU training has been a stumbling block for myself and many others for far too long. With a simple API that works with both eager and graph execution, we can all reap the benefits of distributed training without the usual headaches.

Looking ahead, I expect the TensorFlow team to continue rapidly iterating and delivering innovative features. The DistributionStrategy API previewed in this release is a taste of things to come as TensorFlow evolves to support even larger-scale distributed training and inference.

At the same time, I hope the project can find a sustainable balance between the competing demands of researchers and practitioners. Continued investment in eager execution, the Keras API, and other high-level interfaces is critical to growing the framework‘s mindshare.

I believe this 1.8 release strikes the right balance, delivering powerful new features for experts while continuing to lower the barrier to entry for beginners. As TensorFlow marches towards the 2.0 milestone, this latest update shows the framework is well-positioned to remain the platform of choice for deep learning.