<p dir="auto">Hi,</p>
<p dir="auto">The example <a href="https://github.com/fchollet/keras/blob/master/examples/cifar10_cnn.py">here </a> claims to go to 0.65 loss in 25 epochs. I tried out the code with both the theano and tensorflow backend and with the default hyper-parameters as is but it does not seem to go below 2.30 validation (and 0.1 accuracy) in 25 epochs.</p>
<p dir="auto">Is there anything that I am doing wrong and can anyone else reproduce this? (At least a friend of mine could)</p>
<p dir="auto">The goal is to implement data-parallel multi-GPU training with gradient-averaging properly in Keras (at least explicitly for TensorFlow backend).</p>
<p dir="auto">In this issue I'd like to discuss a particular approach which tries to fix problems of current solutions. Since Keras seems not to be designed for data-parallel SGD, I'm trying to find ways how to modify or adapt the Keras code, while keeping the API philosophy. Since this problem is quite important for  many people, including our team at <a class="user-mention notranslate" data-hovercard-type="organization" data-hovercard-url="/orgs/rossumai/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/rossumai">@rossumai</a>, I'd like to ask for advice. Any feedback is really welcome.</p>
<h2 dir="auto">Quick outline of the data-parallel SGD algorithm</h2>
<p dir="auto">We use N identical model replicas (towers) to train on slices on a mini-batch. Model parameters are placed on a parameter server device (PS), CPU or one of the GPUs, computations are made on N worker devices. A minibatch of inputs is split into N sub-batches, distributed to each worker which computes the forward and backwards pass, the resulting gradients are sent to the PS device, averaged and used to update the weights, which are then copies back to the workers.</p>
<h2 dir="auto">Previous experiments</h2>
<p dir="auto">As a baseline I checked the <a href="https://www.tensorflow.org/tutorials/deep_cnn#training_a_model_using_multiple_gpu_cards" rel="nofollow">TensorFlow CIFAR 10 multi-GPU tutorial</a>. It worked as expected for 1-4 GPUs (TensorFlow 1.2.1, CUDA 8.0, GTX 1070).</p>
<p dir="auto">I tried the approach of <a href="https://github.com/kuza55/keras-extras">kuza55/keras-extras</a>, discussed earlier in other issues (<a class="issue-link js-issue-link" data-error-text="Failed to load title" data-id="149949934" data-permission-text="Title is private" data-url="https://github.com/keras-team/keras/issues/2436" data-hovercard-type="issue" data-hovercard-url="/keras-team/keras/issues/2436/hovercard" href="https://github.com/keras-team/keras/issues/2436">#2436</a>) and blog post <a href="https://medium.com/@kuza55/transparent-multi-gpu-training-on-tensorflow-with-keras-8b0016fd9012" rel="nofollow">Transparent Multi-GPU Training on TensorFlow with Keras</a>, adapting MNIST MLP and CIFAR10 Keras examples (Keras 2.0.6, TensorFlow 1.2.1, CUDA 8.0). In practice using more than one GPU lead to decrease of performance. Between 2 and 4 GPUs there was a 2x speedup, however.</p>
<p dir="auto"><a href="https://gist.github.com/bzamecnik/390d3802b31ce766fbf6fd6c9fd682d3">https://gist.github.com/bzamecnik/390d3802b31ce766fbf6fd6c9fd682d3</a></p>
<h2 dir="auto">Problems with kuza55/keras-extras</h2>
<p dir="auto">After examining the graph in TensorBoard I discovered a problem in this approach: gradients are not computed in parallel on each device, but in whole on the parameter service device. Indeed each worker computes only the predictions which are distributed to PS and concatenated. The loss + gradients are computed there. Another potential problem is that the whole mini-batch is fed to each device which only takes it's slice. We waste our precious IO bandwidth.</p>
<h2 dir="auto">Proposed fixes</h2>
<ul dir="auto">
<li>gradients should be computed in each tower separately, then averaged on PS device</li>
<li>only tower sub-batch slices of input/labels transferred to each tower (not full batch)</li>
<li>should we use queue for providing inputs asynchronously?</li>
<li>are the model parameters properly placed on the PS device and shared?</li>
<li>in addition for a correct parallel SGD implementation we should incorporate corrections outlined in the <a href="https://arxiv.org/abs/1706.02677" rel="nofollow">Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour</a> paper:
<ul dir="auto">
<li>scaled learning rate with warm-up</li>
<li>momentum correction</li>
<li>batch normalization correction</li>
<li>gradient aggregation
<ul dir="auto">
<li>for averaging we can put all the normalization terms inside the loss and then reduce by sum</li>
<li>"Normalize the per-worker loss by total minibatch size kn, not per-worker size n."</li>
</ul>
</li>
<li>random shuffling: "Use a single random shuffling of the training data (per epoch) that is divided amongst all k workers."</li>
<li>gradients aggregatin should be done in parallel with back-propagation
<ul dir="auto">
<li>"As soon as the gradient for a layer is computed, it is aggregated across workers, while gradient computation for the next layer continues."</li>
</ul>
</li>
</ul>
</li>
</ul>
<h2 dir="auto">Proposed implementation</h2>
<h3 dir="auto">Parallel losses and gradients - DataParallelOptimizer</h3>
<p dir="auto">Since the Keras API is at present not directly suitable for data-parallel SGD computation, in the first step of making a working prototype we can make different implementations of Optimizer and Model, let's say DataParallelOptimizer and DataParallelModel.</p>
<p dir="auto">We need to compute loss and gradients in parallel. Tensor for loss is created by Keras within <code class="notranslate">Model.compile()</code> and stored as <code class="notranslate">Model.total_loss</code>. Gradients are computed in <code class="notranslate">Optimizer.get_gradients()</code> which is called in lazily created functions <code class="notranslate">Model.{train,test,predict}_function()</code> (called from <code class="notranslate">fit()</code>, etc.). This function accepts single loss tensor. Various optimizers then compute updates based on the gradients. The problem is a single loss tensor (which can be placed on one device) passed to  <code class="notranslate">Optimizer.get_updates()</code>.</p>
<p dir="auto">So far the only way I see it so change <code class="notranslate">Model.total_loss</code> from a single tensor into a list of tensors, each of them able to be placed on a different device.<br>
<code class="notranslate">DataParallelOptimizer</code> wrapper class can derive from <code class="notranslate">Optimizer</code> and override <code class="notranslate">get_gradients()</code> to accept loss as a list of tensors and average them. This would be the place of synchronization of the GPU workers. The <code class="notranslate">get_updates()</code> function (implemented any of the wrapped <code class="notranslate">Optimizer</code>) just calls <code class="notranslate">get_gradients()</code>. Note that thanks to the option <code class="notranslate">collocate_gradients_with_ops=True</code> in TF implementation of <code class="notranslate">K.gradients()</code> the gradient ops would automatically be placed on the same device as the loss ops even though it's called outside compile() and the device scope. (TODO: issue link)</p>
<h3 dir="auto">Model replication and feeding data - DataParallelModel</h3>
<p dir="auto">We need a Model which contains the replicas (towers) and provides the list of losses to the <code class="notranslate">DataParallelOptimizer</code>. We would adapt the code in the <code class="notranslate">make_parallel()</code> function from kuza55/keras-extras. The <code class="notranslate">DataParallelModel</code> would take via the constructor an instance of the basic model. Similar like in <code class="notranslate">make_parallel()</code> it would make N replicas of this model placed on different devices. We could try to set TF variable reuse after the first replica. Also we make the merged model, which concatenates the outputs, and use it for the actual training. Better then slicing the duplicate inputs we can pass the sub-batch inputs separately and route them to each replicate directly.</p>
<p dir="auto">Then it would override <code class="notranslate">compile()</code> which has to call <code class="notranslate">compile()</code> on each replica (and the merged model) - in order to place losses and gradients - and gather <code class="notranslate">total_loss</code> operation from all replicas. In <code class="notranslate">compile()</code> we also wrap the provided optimizer with <code class="notranslate">DataParallelOptimizer</code> and inject both the total_loss list and the DataParallelOptimizer instance to the merged model. The rest of the methods in <code class="notranslate">DataParallelModel</code> will be proxied to the merged model.</p>
<p dir="auto">In case we want to avoid slicing the inputs we could change the inputs in {train,test,predict}_function() and perform the slice in <code class="notranslate">*_on_batch()</code> functions.</p>
<h2 dir="auto">Code</h2>
<p dir="auto"><del><a href="https://gist.github.com/bzamecnik/92607207af912ae53dd2aa557631b977">https://gist.github.com/bzamecnik/92607207af912ae53dd2aa557631b977</a></del></p>
<p dir="auto"><a href="https://github.com/rossumai/keras-multi-gpu">https://github.com/rossumai/keras-multi-gpu</a></p>
<p dir="auto">I have prepared an implementation of DataParallelOptimizer and I'm working on DataParallelModel. The mechanism of the latter is not as clear at the moment. In the first stage I'd like to make a working prototype, then make experiments to show that the model produces correct results and that we obtain benefit from scaling to multiple GPUs. Next I wish to make the API cleaner. So far I think the code might be separate from Keras, since it will depend on TensorFlow explicitly and I'm not sure about Theano support.</p>
<p dir="auto">If you read this rather longer I'd like to kindly ask for advice if you think this approach is feasible or you see any problems with that. Any feedback is really welcome. Thanks!</p>
<p dir="auto">/cc <a class="user-mention notranslate" data-hovercard-type="organization" data-hovercard-url="/orgs/rossumai/hovercard" data-octo-click="hovercard-link-click" data-octo-dimensions="link_type:self" href="https://github.com/rossumai">@rossumai</a></p>