Auto-Vectorizing TensorFlow Graphs: Jacobians, Auto-Batching And Beyond

Abstract

We propose a static loop vectorization optimization on top of high level dataflow IR used by frameworks like TensorFlow. A new statically vectorized parallel-for abstraction is provided on top of TensorFlow, and used for applications ranging from auto-batching and per-example gradients, to jacobian computation, optimized map functions and input pipeline optimization. We report huge speedups compared to both loop based implementations, as well as run-time batching adopted by the DyNet framework.

References in corpus (4)

• TensorFlow: Large-Scale Machine Learning on Heterogeneous Distributed Systems
• Adaptive Computation Time for Recurrent Neural Networks
• Deep Learning with Dynamic Computation Graphs
• DLVM: A modern compiler infrastructure for deep learning systems

Cited by in corpus (5)

• Not All Unlabeled Data are Equal: Learning to Weight Data in Semi-supervised Learning
• Enabling Fast Differentially Private SGD via Just-in-Time Compilation and Vectorization
• Deep Learning with Label Differential Privacy
• Automatically Batching Control-Intensive Programs for Modern Accelerators
• Self-Tuning Stochastic Optimization with Curvature-Aware Gradient Filtering