Maximize tensorflow multi gpu performance
<p>I was wondering if anybody could advise on how to get peak performance out of tensorflow in a 4 GPU setting.</p>

<p>As a test I created two of the same network (18 ish layer residual network with small filter banks (ranging from 16-128) on 32x32 inputs. Batch size 512, 128 per GPU.). One in MXNet and one I have modelled off of <a href="https://github.com/tensorflow/models/tree/master/inception">the inception example</a>.</p>

<p>My MXNet network can train at around 7k examples a second where tensorflow is only capable of 4.2k with dummy data and 3.7 with real data.</p>

<p>(when running on 1 GPU the numbers are 1.2k examples a second vs 2.1k)</p>

<p>In my experiment I have a few questions in hopes to speed things up.</p>

<ol>
<li><p>GPU utilization seems quite low when training. I noticed that in the tensorflow white paper there is support for running multiple streams on the same GPU. Is this possible in the public release?</p></li>
<li><p>Is there anyway to perform multiple train operations in one execution of <code>session.run()</code>? Or have async execution? This would allow for weight updates to be done at the same time as the next batches forward pass? I have tried using 2 threads (both system and with <code>QueueRunners</code>'s), but this only resulted in a slowdown. MXNet is able to increase speeds by running weight updates on the CPU so that the gpu's can be used for the next batch.</p></li>
<li><p>Will the new distributed run time get around some of these issues by letting me run more than one worker on a single machine?</p></li>
<li><p>Is there something else that can be done?</p></li>
</ol>

<p>I know there are a number of similar questions here on stack overflow, but though my searching I couldn't find a solution to my problems that I have not already tried.</p>

<p><strong>Edit:</strong></p>

<p>I did a little bit of CUDA profiling to see what the expensive kernels were. According to my run, 21.4% of the time is spent inside:</p>

<pre><code>void Eigen::internal::EigenMetaKernel_NonVectorizable&lt;Eigen::TensorEvaluator
&lt;Eigen::TensorAssignOp&lt;Eigen::TensorMap&lt;Eigen::Tensor&lt;float, int=4, int=1, long&gt;, int=16&gt;,
Eigen::TensorPaddingOp&lt;Eigen::array&lt;std::pair&lt;int, int&gt;,
unsigned long=4&gt; const, Eigen::TensorMap&lt;Eigen::Tensor&lt;float const,
int=4, int=1, long&gt;, int=16&gt; const &gt; const &gt; const, Eigen::GpuDevice&gt;, long&gt;(float, int=4)
</code></pre>

<p>and 20.0% of the time were spent in </p>

<pre><code>void Eigen::internal::EigenMetaKernel_NonVectorizable&lt;Eigen::TensorEvaluator
&lt;Eigen::TensorAssignOp&lt;Eigen::TensorMap&lt;Eigen::Tensor&lt;float, int=4, int=1, long&gt;, int=16&gt;,
Eigen::TensorBroadcastingOp&lt;Eigen::array&lt;int, unsigned long=4&gt;
const, Eigen::TensorMap&lt;Eigen::Tensor&lt;float const, int=4, int=1, long&gt;,
int=16&gt; const &gt; const &gt; const, Eigen::GpuDevice&gt;, long&gt;(float, int=4)
</code></pre>

<p>Off of the Signature I am not exactly sure what these are doing. Do these make sense?</p>

<p>In addition to this, the analysis reports low kernel concurrency, 0%, as expected.
And Low compute utilization 34.9% (granted this includes start-up time and a little bit of python in train loop. Around 32 seconds total out of 91. This comes out to around 50% utilization inside tensorflow.)</p>

<p><strong>Edit 2:</strong></p>

<p>I have attached a copy of the trimmed down <a href="https://gist.github.com/lukemetz/2072e9f7b3f2b9325e25">source code</a>. In general though I am more concerned about question 1-3 and don't want to take too much of ever bodies time.</p>

<p>In addition I am running on tensorflow built from: f07234db2f7b316b08f7df25417245274b63342a</p>

<p><strong>Edit 3:</strong></p>

<p>Updated to the most recent tensorflow (63409bd23facad471973b110df998782c0e19c06) same code, default data format (NHWC) and that seemed to speed this up a lot.
On fake data 6.7k-6.8k (thermal dependence I think?) examples a second 4gpu. 1gpu -- 2.0k examples a second.
Real data performance is around 4.9k examples a second for 4gpu. 1gpu -- 1.7k examples a second.</p>

<p><strong>Edit 4:</strong></p>

<p>In addition I tried out switching data formats to BCHW. I made the conversion modelled off of <a href="https://github.com/soumith/convnet-benchmarks/blob/master/tensorflow/benchmark_googlenet.py">Soumith's benchmarks</a>. The convolution parts were indeed faster, but batch norm appears to be messing everything up. With a <a href="https://gist.github.com/lukemetz/2072e9f7b3f2b9325e25#file-batchnorm_function-py">naive implementation</a> (fixing axis, and making weights [1,C,1,1] instead of [C,]) I am only able to get 1.2k examples a second on 4 gpu (fake data). Where as with a transpose before and after the batch norm op I am able to get 6.2k examples a second (fake data). Still slower than the NHWC data_format.</p>