TF-Replicator: Distributed Machine Learning for Researchers

Abstract

We describe TF-Replicator, a framework for distributed machine learning designed for DeepMind researchers and implemented as an abstraction over TensorFlow. TF-Replicator simplifies writing data-parallel and model-parallel research code. The same models can be effortlessly deployed to different cluster architectures (i.e. one or many machines containing CPUs, GPUs or TPU accelerators) using synchronous or asynchronous training regimes. To demonstrate the generality and scalability of TF-Replicator, we implement and benchmark three very different models: (1) A ResNet-50 for ImageNet classification, (2) a SN-GAN for class-conditional ImageNet image generation, and (3) a D4PG reinforcement learning agent for continuous control. Our results show strong scalability performance without demanding any distributed systems expertise of the user. The TF-Replicator programming model will be open-sourced as part of TensorFlow 2.0 (see https://github.com/tensorflow/community/pull/25).

References in corpus (22)

• Very Deep Convolutional Networks for Large-Scale Image Recognition
• Continuous control with deep reinforcement learning
• Spectral Normalization for Generative Adversarial Networks
• Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour
• MXNet: A Flexible and Efficient Machine Learning Library for Heterogeneous Distributed Systems
• Improved Techniques for Training GANs
• Mixed Precision Training
• Density estimation using Real NVP
• Don't Decay the Learning Rate, Increase the Batch Size
• IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures
• DeepMind Control Suite
• Horovod: fast and easy distributed deep learning in TensorFlow
• Distributed Prioritized Experience Replay
• cGANs with Projection Discriminator
• Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes
• Extremely Large Minibatch SGD: Training ResNet-50 on ImageNet in 15 Minutes
• Learning in Implicit Generative Models
• Connecting Generative Adversarial Networks and Actor-Critic Methods
• SparkNet: Training Deep Networks in Spark
• Decoupled Neural Interfaces using Synthetic Gradients
• High-Accuracy Low-Precision Training
• Scale out for large minibatch SGD: Residual network training on ImageNet-1K with improved accuracy and reduced time to train

Cited by in corpus (13)

• Large Scale Adversarial Representation Learning
• Adversarial Video Generation on Complex Datasets
• Stabilizing Transformers for Reinforcement Learning
• Deep Learning without Weight Transport
• Object Discovery with a Copy-Pasting GAN
• V-MPO: On-Policy Maximum a Posteriori Policy Optimization for Discrete and Continuous Control
• End-to-End Adversarial Text-to-Speech
• Two Routes to Scalable Credit Assignment without Weight Symmetry
• Data Movement Is All You Need: A Case Study on Optimizing Transformers
• Efficient Replication for Straggler Mitigation in Distributed Computing
• The Case for Strong Scaling in Deep Learning: Training Large 3D CNNs with Hybrid Parallelism
• Identifying Learning Rules From Neural Network Observables
• Non-Differentiable Supervised Learning with Evolution Strategies and Hybrid Methods