How does convolution scale so efficiently with batches?

marcman411: <p>I believe that conv2D is implemented with some cudnn tricks. But how does it scale so well with batches?</p>
<p>When I run a batch size of 1:</p>
<pre><code class="lang-auto">bs = 1
h = 128
w = 256
in_channels = 512
out_channels = 1024
kernel_size = (3,3)
stride = (1,1)
padding = (1,1)
dilation = (1,1)
cuda = True

test = Test(
    in_channels=in_channels,
    out_channels=out_channels,
    kernel_size=kernel_size,
    stride=stride,
    padding=padding,
    dilation=dilation
)

x = torch.ones(bs, in_channels, h, w).float()
weights = torch.ones(out_channels, in_channels, kernel_size[0], kernel_size[1])
bias = 100*torch.ones(out_channels)

print("Input:")
print(x.shape)

if cuda:
    test = test.cuda()
    x = x.cuda()
    weights = weights.cuda()
    bias = bias.cuda()

s = time.time()
pytorch = F.conv2d(x, 
    weights, 
    bias,
    stride=stride,
    padding=padding,
    dilation=dilation)
pytorch_time = time.time() - s

print('PyTorch:')
print(pytorch.shape)
print('PyTorch Time: ', pytorch_time)
</code></pre>
<p>I get:</p>
<pre><code class="lang-auto">&gt;&gt; Input:
&gt;&gt; torch.Size([1, 512, 128, 256])
&gt;&gt; PyTorch:
&gt;&gt; torch.Size([1, 1024, 128, 256])
&gt;&gt; PyTorch Time:  0.004613637924194336

</code></pre>
<p>Running the same code above with a batch size of 5 gives:</p>
<pre><code class="lang-auto">&gt;&gt; Input:
&gt;&gt; torch.Size([5, 512, 128, 256])
&gt;&gt; PyTorch:
&gt;&gt; torch.Size([5, 1024, 128, 256])
&gt;&gt; PyTorch Time:  0.0049724578857421875
</code></pre>
<p>The wall-clock time is negligibly different. Looking at the code in the THCUNN library (granted, it’s not cudnn), it appears that the batches are processed in a for-loop (<a href="https://github.com/pytorch/pytorch/blob/master/aten/src/THCUNN/generic/SpatialFullDilatedConvolution.cu#L138" rel="nofollow noopener">see here</a>). This is the same way Caffe handles it too.</p>
<p>I can’t find any code to suggest that multiple batches are processed (im2col+gemm) together. Everything I find seems to suggest each batch is handled sequentially. Given this, <strong>how is there no linear scaling to the running time?</strong></p>

albanD: <p>The thing is that the cuda API is asynchronous so the only thing you measure here is the time to launch the job, not how long it takes to run on the GPU. You need to synchronise:</p>
<pre><code class="lang-auto"># Make sure nothing was still running on the GPU
torch.cuda.synchronize()
s = time.time()
pytorch = F.conv2d(x, 
    weights, 
    bias,
    stride=stride,
    padding=padding,
    dilation=dilation)
# Wait for the job to actually be done
torch.cuda.synchronize()
pytorch_time = time.time() - s
</code></pre>

marcman411: <p>Ahh, thank you, that makes sense.</p>
<p>Now the real question is: how does PyTorch launch the jobs so efficiently? For the learning experience, I’ve re-implemented convolution with my own version of <code>im2col</code>. Using mine, the kernel launch time is an order of magnitude slower, but the actual running time (after synchronizing) is only 1.25x slower. Is there some funky stuff going on under the hood to launch the job?</p>

albanD: <p>There is nothing special, just that the whole convolution is going directly going to cpp and then launches a single cudnn kernel.<br>
Your version with im2col most certainly launch few kernels which result in an increased launch time.</p>

marcman411: <p>Is there somewhere I can find more detail about this? You’re right about launching a few kernels. Each batch launches a kernel for the <code>im2col</code> computation and a kernel for the gemm (done through cuBLAS). However, when I profile the loop, I see that most of the launch time is taken up by the initial <code>im2col</code> call (i.e. in the first batch). Synchronizing with <code>cudaDeviceSynchronize()</code> after each call produces these times:</p>
<pre><code class="lang-auto">batch: 0
im2col: 0.000030s
gemm: 0.000026s
batch: 1
im2col: 0.000012s
gemm: 0.000008s
batch: 2
im2col: 0.000008s
gemm: 0.000007s
batch: 3
im2col: 0.000008s
gemm: 0.000007s
batch: 4
im2col: 0.000008s
gemm: 0.000007s
batch: 5
im2col: 0.000008s
gemm: 0.000007s
</code></pre>
<p>I’d presumed the speed up for the later batches was due to caching the columns (for <code>im2col</code>) and the weights (for gemm), but I don’t know how I’d do this with a single kernel. Or this this something that cudnn makes possible? As in, does <a href="https://docs.nvidia.com/deeplearning/sdk/cudnn-developer-guide/index.html#cudnnConvolutionForward" rel="nofollow noopener"><code>cudnnConvolutionForward</code></a> implement the entire operation on a single kernel?</p>

albanD: <p>The very first batch is always slower as many things are lazily initialized in cuda. So the first one will actually potentially do extra initialization work. Hopefully you should be able to ignore these timings for any real workload.</p>
<p>I’m not sure about the cudnn internals. You should be able to use <code>nvvp</code> to see that. But I think yes it does everything in a single kernel (but that may vary depending on the algorithm that was selected).</p>

marcman411: <p>I think you’re right. I was trying my implementation with varying kernel sizes and strides. I found with a standard (3,3) kernel and unit strides, I’m about 30% off the speed of the cudnn convolution used by Pytorch. However, for non-square kernels and non-unit strides I am roughly equal. In general, though, I’m finding square kernels with unit strides are unbeatable.</p>
<p>I think it heavily depends on the algorithm. My guess is that the unit stride, (3,3) convolution is particularly optimized as it’s simple and very common. I think for that operation, an FFT convolution is faster than <code>im2col</code> so the selector is likely choosing that algorithm, which I am pretty sure can be implemented on a single kernel.</p>
<p>Thanks for the help making sense of this!</p>

Bjorn_Lindqvist: <p>You may glean some insight from this research: <a href="https://arxiv.org/abs/2110.03901" class="inline-onebox" rel="noopener nofollow ugc">[2110.03901] Characterizing and Demystifying the Implicit Convolution Algorithm on Commercial Matrix-Multiplication Accelerators</a> cuDNN uses the <code>CUDNN_CONVOLUTION_FWD_ALGO_IMPLICIT_GEMM</code> algorithm to avoid unfolding the whole input tensor before doing the matrix multiplication. It saves huge amounts of DRAM bandwidth and will thus always be faster than the “explicit” im2col algorithm.</p>