Performance Analysis of Just-in-Time Compilation for Training TensorFlow Multi-Layer Perceptrons

The TensorFlow system [1] has been developed to provide a general, efficient and scalable framework for writing Machine Learning (ML) applications. With the rapid advancement and popularity of ML, and deep artificial neural networks in particular, TensorFlow has found success in both research and industry, particularly as the main ML engine for Google’s Cloud Platform [2]. TensorFlow targets a large, diverse range of application areas, from ML research on a single machine, to large, distributed, and heterogeneous production environments. It is therefore crucial that TensorFlow is able to achieve high performance in production executions, while also retaining its capability for fast and efficient prototyping. Furthermore, TensorFlow’s userbase encompasses a wide variety of expertise in application performance optimization, thus the TensorFlow run-time must achieve these goals without overly burdening the user with complex performance concerns. However, TensorFlow presents a very large parameter optimization space to the user, where programming and configuration decisions can have a significant impact on performance. Similarly to other popular ML run-times with high-level language interfaces [3], [4], [5], TensorFlow places a high degree of emphasis on programmability and portability, which then forms a trade-off with performance. To achieve high performance, programmers need to make decisions ranging from the choice of a parallelization strategy for ML models, to allocating computation to hardware resources, to the format, representation, and communication patterns of input and output data, as well as a host of options around operation compilation and fusion. These decisions are non-trivial, as they are highly dependent on the nature of the interaction between the ML model and the hardware and software configuration of the underlying system. To avoid this complexity, programmers may leave some decisions to the TensorFlow run-time, e.g., by not explicitly specifying the mapping of computations to the hardware. However, there is no automatic way for the run-time to reliably and consistently apply optimal mappings. This often results in the user engaging in an inefficient parameter search when attempting to optimize a particular TensorFlow application, further complicated when considering the potential to trade ML model accuracy for increased runtime performance. It is our position that performance analysis tooling for TensorFlow could greatly aid this process, and may potentially lead to automated program optimization techniques. As a first step, in this extended abstract, we investigate the performance of Just-In-Time (JIT) compilation in TensorFlow for the relatively straightforward use-case of training Multi-Layer Perceptrons (MLPs). By employing performance analysis, we aim to develop an understanding of when JIT compilation may be beneficial for performance, which could then be used to enable or disable JIT compilation in future program executions. We then suggest additional performance analysis techniques that could be employed in order to extend and automate the process for optimizing TensorFlow, in particular with respect to JIT compilation.