Source-linked AI summary
One weird trick for parallelizing convolutional neural networks
Alex Krizhevsky
TL;DR
Training convolutional neural networks across multiple GPUs requires a parallelization strategy suited to their mixed layer structure. The paper combines data parallelism in convolutional layers with model parallelism in fully-connected layers, and reports good but not quite linear scaling for the evaluated model. Larger batch sizes can reduce accuracy, although the loss depends on the dataset.
Problem
CNNs contain convolutional and fully-connected layers with different computation, parameter, and representation properties, motivating a parallelization strategy matched to this structure.
Method
The method uses data parallelism heavily in convolutional layers and model parallelism in fully-connected layers across K workers.
Results
The parallelization scheme scales well for the evaluated model, although its scaling is not quite linear.
Takeaways & Limitations
The scheme works well on existing model architectures, with potentially better scaling from architectures adapted for multi-GPU training.
Takeaways & Limitations
Larger batch sizes incur dataset-dependent accuracy loss, and dense matrix multiplications near the output limit scaling for the evaluated network.
Abstract
from arXiv · showhide
I present a new way to parallelize the training of convolutional neural networks across multiple GPUs. The method scales significantly better than all alternatives when applied to modern convolutional neural networks.
1 Introduction
The paper introduces a new way to parallelize convolutional-neural-network training with stochastic gradient descent. It presents two variants: one exactly simulates synchronous SGD, while the other approximates it but works better in practice.
- The note introduces a new method for parallelizing convolutional-neural-network training with stochastic gradient descent.
- The first algorithmic variant perfectly simulates synchronous SGD on one core.
- The second variant introduces an approximation and works better in practice despite no longer exactly simulating SGD.
2 Existing approaches
CNN training can be parallelized across model or data dimensions, with synchronization requirements and efficiency determined by model architecture and batch size. The paper motivates combining these forms of parallelism rather than treating either as universally superior.
- CNN training offers model parallelism, which assigns different model parts to workers, and data parallelism, which assigns different examples.
- Model parallelism synchronizes neuron activities, whereas data parallelism synchronizes parameters or parameter gradients.
- Neither parallelization scheme is universally better; their relative use should reflect the model architecture.
- Model parallelism is efficient with high computation per neuron activity, while data parallelism is efficient with high computation per weight.
- Increasing batch size can improve data-parallel efficiency, but very large batches can reduce SGD convergence rate and final-solution quality.
3 Some observations
Modern CNNs divide into computationally heavy convolutional layers with few parameters and parameter-heavy fully-connected layers with little computation. This motivates using data parallelism for convolutional layers and model parallelism for fully-connected layers.
- Convolutional layers contain about 90-95% of computation, about 5% of parameters, and large representations.
- Fully-connected layers contain about 5-10% of computation, about 95% of parameters, and small representations.
- The paper proposes heavily using data parallelism in convolutional layers and model parallelism in fully-connected layers.
- Figure 1 depicts K workers training a network with three convolutional layers and two fully-connected layers.
4 The proposed algorithm
The proposed algorithm combines data parallelism in convolutional layers with model parallelism in fully-connected layers, using three communication schemes. Scheme (c) is identified as most efficient, while schemes (b) and (c) can be equivalent to synchronous SGD or use variable batch sizes for faster convergence.
- Hybrid parallelization: The method uses data parallelism for convolutional layers and switches to model parallelism for fully-connected layer activities across K workers.Each worker processes a distinct batch of 128 examples before exchanging last-stage convolutional activities.
- Communication schemes: Scheme (a) assembles a batch of 128K examples at every worker, pausing useful work and increasing memory demands.Large batches may nevertheless be processed more efficiently by GPUs.
- Communication schemes: Scheme (b) hides much of communication by broadcasting successive workers’ convolutional activities while fully-connected computation proceeds.The hidden communication can overlap with fully-connected computation, which is the network’s most significant communication.
- Communication schemes: Scheme (c) keeps communication-to-computation ratio constant in K by distributing activity exchanges across many workers rather than relying on one sender.This is described as a major advantage for large K and makes scheme (c) the most efficient in the backward pass as well.
- Weight synchronization: After backpropagation, workers synchronize convolutional-layer gradients by accumulating and broadcasting disjoint 1/Kth gradient-matrix portions.The method assigns each worker one gradient portion, gathers corresponding portions, and broadcasts the accumulated results.
- Variable batch size: Schemes (b) and (c) can exactly match synchronous SGD with batch size 128K, or use batch sizes 128 and 128K in different layers.The variable-batch variant is not pure SGD but leads to faster convergence and better minima as the effective batch reaches the thousands.
5 Experiments
Experiments examine batch-size accuracy costs and multi-GPU training performance on an ILSVRC 2012 model, using adjusted optimization settings and scheme (b). The results show reduced accuracy costs with variable batch sizes, but non-linear scaling constrained by communication and architecture.
- Experimental setup: The experiments use a 1.2-million-image, 1000-category ILSVRC 2012 dataset and a minor single-tower variation of the winning 2012 model.The model has 0.2% more parameters and 2.4% fewer connections than the original two-tower model.
- Parallel training: Scheme (b) replaces standard training passes with a multi-GPU propagation procedure whose communication includes one-to-all broadcast/reduction.The reported scaling limitations identify this communication pattern as increasingly costly.
- Optimization settings: When batch size is multiplied by k, the practical heuristic multiplies the learning rate by k, while keeping momentum at µ = 0.9.The author reports that this heuristic worked best for the batch sizes considered, despite differing from the theoretical prescription.
- Results: The main result is that larger batch sizes incur an accuracy cost, but the variable batch size trick can greatly reduce it.Table 1 also reports that the parallelization scheme scales fairly well for the evaluated model, though not linearly.
- Scaling constraints: Dense matrix multiplications near the output spend more time communicating than computing for 4096×4096 matrices on the stated hardware.The author expects better scaling from larger matrices or restricted connectivity in the last two hidden layers.
- Scaling constraints: Scaling from 4 to 8 GPUs is particularly impaired because the machine cannot provide simultaneous full-speed communication among all eight GPUs.Simultaneous full-speed communication is available only among certain subsets of four GPUs.
6 Comparisons to other work on parallel convolutional neural network training
Table 1 compares the paper’s results favorably with published alternatives, including a prior multi-GPU approach that used model and data parallelism uniformly across layers.
- Table 1 compares the results favorably to published alternatives.
- Yadan et al. parallelized convolutional neural-network training with model parallelism and data parallelism, using the same parallelism form in every layer.
- 2.2x speedup was achieved on 4 GPUs relative to a 1-GPU implementation requiring 226.8 hours for 90 epochs.The baseline used an NVIDIA GeForce Titan.
7 Other work on parallel neural network training
Prior work parallelized neural network training through spatial model partitioning or asynchronous SGD, while this paper’s setup combines data parallelism in convolutional layers with model parallelism in fully-connected layers.
- Coates et al. distributed workers spatially across locally-connected image regions, communicating neuron activations near region edges.Applying this scheme to convolutional networks would additionally require synchronizing convolutional weights.
- The work is positioned as extending prior work on parallel neural network training.The supplied passage identifies Coates et al. as extending Dean et al.’s work, which itself builds on Niu et al.’s asynchronous SGD.
- Each GPU performs 4096 × 512 × 2 FLOPs per sample and receives 4096 floats in the reported communication example.The passage gives 2.09µs for computation at 2 TFLOPs/sec and 2.73µs for receiving data at 6GB/sec.
- The paper’s experiments use data parallelism in convolutional layers and model parallelism in fully-connected layers.Table 1 reports effective batch sizes separately for convolutional and fully-connected layers, with training time measured in hours.
- Dean et al. introduced this form of model parallelism for locally-connected networks and the asynchronous SGD variant later used by Paine et al.That line of work builds on Niu et al., who introduced asynchronous SGD and demonstrated it for models with sparse gradients.
8 Conclusion
The proposed scheme is presented as effective on existing convolutional architectures, with further scaling expected from architectures designed for multi-GPU settings. The conclusion identifies restricted connectivity, alternative communication schemes, and reduced effective batch sizes as directions, while noting batch-size accuracy costs.
- 8 Conclusion: The scheme works well on existing model architectures without adaptation to the multi-GPU setting.The authors expect better scaling from architectures more suited to multi-GPU training.
- 8 Conclusion: Restricted connectivity in upper layers could replace dense connectivity or limit fully-connected communication to small linear projections.These architectural changes are proposed as the algorithm scales past 8 GPUs.
- 8 Conclusion: The conclusion proposes switching from scheme (b) to scheme (c), or using a hybrid between them.
- 8 Conclusion: Restricted model parallelism in convolutional layers could reduce effective batch size, as in the two-column network of Krizhevsky et al.
- 8 Conclusion: Bigger batch sizes can reduce accuracy, with the magnitude depending on the dataset and generally being smaller for larger, more varied datasets.