An MPI-Based Python Framework for Distributed Training with Keras

Recent progress in machine learning has enabled deep neural networks (DNNs) to advance the state of the art in a wide range of problem domains, from computer vision to high energy physics [3] [4]. As the applicability of DNNs has broadened, there have been efforts to develop userfriendly tools for building them. Software packages such as Keras [5] and TFLearn [6] facilitate the construction and training of deep neural networks, offering a flexible interface for combining common model components and configuring the optimization process. Large model sizes and long training times have motivated the development of distributed training algorithms for DNNs [7] [8]. These algorithms work by splitting the training task across multiple concurrent processes, which can be threads on a single machine or jobs spread across the nodes of a cluster. The speedup provided by distributed algorithms is relevant when fast training is critical, such as when iterating on model choice during development, or when retraining a model on new data in a production environment. Despite the rise of convenient model-building software packages such as Keras, there are few tools for interfacing these packages with distributed training algorithms. In this paper we introduce a lightweight Python framework, mpi learn, that provides a straightforward means of training Keras models in a distributed fashion. The framework is built on the Message Processing Interface (MPI) protocol [10] and can operate on personal machines, multi-GPU servers, and large supercomputing sites alike.