SparkNet: Training Deep Networks in Spark

TL;DR

<3-5 sentence high-level summary> SparkNet tackles the challenge of distributed deep-network training within batch-processing frameworks by proposing a minimal-communication SGD scheme that works in Spark. It integrates Caffe with Spark via a Net API and a lightweight tensor library to enable out-of-the-box training on existing architectures. Theoretical limits and EC2 benchmarks show meaningful speedups even under bandwidth constraints, validating the approach for large-scale datasets like ImageNet. The work highlights the practicality of co-locating data processing and model training in Spark, enabling end-to-end data-to-model pipelines.

Abstract

Training deep networks is a time-consuming process, with networks for object recognition often requiring multiple days to train. For this reason, leveraging the resources of a cluster to speed up training is an important area of work. However, widely-popular batch-processing computational frameworks like MapReduce and Spark were not designed to support the asynchronous and communication-intensive workloads of existing distributed deep learning systems. We introduce SparkNet, a framework for training deep networks in Spark. Our implementation includes a convenient interface for reading data from Spark RDDs, a Scala interface to the Caffe deep learning framework, and a lightweight multi-dimensional tensor library. Using a simple parallelization scheme for stochastic gradient descent, SparkNet scales well with the cluster size and tolerates very high-latency communication. Furthermore, it is easy to deploy and use with no parameter tuning, and it is compatible with existing Caffe models. We quantify the dependence of the speedup obtained by SparkNet on the number of machines, the communication frequency, and the cluster's communication overhead, and we benchmark our system's performance on the ImageNet dataset.

Figures (7)

• Figure 1: This figure depicts the SparkNet architecture.
• Figure 2: Computational models for different parallelization schemes.
• Figure 3: This figure shows the speedup $\tau M_a(b,\tau,K)/N_a(b)$ given by SparkNet's parallelization scheme relative to training on a single machine to obtain an accuracy of $a=20\%$. Each grid square corresponds to a different choice of $K$ and $\tau$. We show the speedup in the zero communication overhead setting. This experiment uses a modified version of AlexNet on a subset of ImageNet ($100$ classes each with approximately $1000$ images). Note that these numbers are dataset specific. Nevertheless, the trends they capture are of interest.
• Figure 4: This figure shows the speedups obtained by the naive parallelization scheme and by SparkNet as a function of the cluster's communication overhead (normalized so that $C(b)=1$). We consider $K=5$. The data for this plot applies to training a modified version of AlexNet on a subset of ImageNet (approximately $1000$ images for each of the first $100$ classes). The speedup obtained by the naive parallelization scheme is $C(b)/(C(b)/K + S)$. The speedup obtained by SparkNet is $N_a(b) C(b) / [(\tau C(b) + S) M_a(b,K,\tau)]$ for a specific value of $\tau$. The numerator is the time required by serial SGD to achieve an accuracy of $a$, and the denominator is the time required by SparkNet to achieve the same accuracy (see Equation \ref{['eq:naive_time']} and Equation \ref{['eq:sparknet_time']}). For the optimal value of $\tau$, the speedup is $\max_{\tau} N_a(b) C(b) / [(\tau C(b) + S) M_a(b,K,\tau)]$. To plot the SparkNet speedup curve, we maximize over the set of values $\tau \in \{1, 2, 5, 10, 25, 100, 500, 1000, 2500\}$ and use the values $M_a(b,K,\tau)$ and $N_a(b)$ from the experiments in the fifth row of Figure \ref{['fig:performance_heatmap2']}. In our experiments, we have $S \approx 20s$ and $C(b) \approx 2s$.
• Figure 5: This figure shows the performance of SparkNet on a $3$-node, $5$-node, and $10$-node cluster, where each node has $1$ GPU. In these experiments, we use $\tau=50$. The baseline was obtained by running Caffe on a single GPU with no communication. The experiments are performed on ImageNet using AlexNet.