A Data-Centric Optimization Framework for Machine Learning
  Rapid progress in deep learning is leading to a diverse set of quickly
changing models, with a dramatically growing demand for compute. However, as
frameworks specialize performance optimization to patterns in popular networks,
they implicitly constrain novel and diverse models that drive progress in
research. We empower deep learning researchers by defining a flexible and
user-customizable pipeline for optimizing training of arbitrary deep neural
networks, based on data movement minimization. The pipeline begins with
standard networks in PyTorch or ONNX and transforms computation through
progressive lowering. We define four levels of general-purpose transformations,
from local intra-operator optimizations to global data movement reduction.
These operate on a data-centric graph intermediate representation that
expresses computation and data movement at all levels of abstraction, including
expanding basic operators such as convolutions to their underlying
computations. Central to the design is the interactive and introspectable
nature of the pipeline. Every part is extensible through a Python API, and can
be tuned interactively using a GUI. We demonstrate competitive performance or
speedups on ten different networks, with interactive optimizations discovering
new opportunities in EfficientNet.
