Parallelize the application of multiple CNNs to multiple images

ShahRukh: <p>Hello!</p>
<p>I’m testing out a rather unusual method that requires the application of <code>n</code> different CNNs to <code>n</code> different images. It is a one-to-one mapping where <code>cnn_1</code> is applied to <code>img_1</code>, <code>cnn_2</code> is applied to <code>img_2</code> and so on. Currently I loop through each of the <code>n</code> networks and images to perform the forward pass, which gets rather slow for large <code>n</code>. I was just wondering if there is a way to parallelize this operation so that all <code>n</code> forwards passes can happen simultaneously?</p>

LeviViana: <p>I hope you have enough GPUs &amp; memory to scale your system well as <code>n</code> increases !</p>
<p>Otherwise, you can use the <code>torch.distributed.launch</code> module. Take a look at <a href="https://discuss.pytorch.org/t/how-to-use-multigpu-pytorch/46290/2">this snippet</a>, it might give you a better idea on how you can easily parallelize and even distribute compute using pytorch.</p>

ShahRukh: <p>Thanks <a class="mention" href="/u/leviviana">@LeviViana</a>! If my understanding of  <code>torch.distributed.launch</code> is correct, it can be used to distribute compute over multiple GPUs? However, I’d like to parallelize this operation on a single GPU. Ideally, it would be akin to <code>torch.bmm</code>, with the only difference being that instead of a matrix multiplication being batched the CNN forward pass is batched.</p>

LeviViana: <p>Here is a snippet explaining how you can achieve this using grouped convolutions:</p>
<pre><code class="lang-python">import torch

image_1 = torch.rand(1, 3, 50, 50)
image_2 = torch.rand(1, 3, 50, 50)

conv_weight_1 = torch.rand([4, 3, 3, 3]) # in_channels=3, out_channels=4, kernel_size=(3, 3)
conv_weight_2 = torch.rand([4, 3, 3, 3]) # in_channels=3, out_channels=4, kernel_size=(3, 3)

conv_bias_1 = torch.rand([4]) # in_channels=3, out_channels=4, kernel_size=(3, 3)
conv_bias_2 = torch.rand([4]) # in_channels=3, out_channels=4, kernel_size=(3, 3)

# 1st case -&gt; not fused convolutions:

res_1 = torch.nn.functional.conv2d(image_1, conv_weight_1, conv_bias_1)
res_2 = torch.nn.functional.conv2d(image_2, conv_weight_2, conv_bias_2)

res_not_fused = torch.cat((res_1, res_2), dim=0)

# 2nd case -&gt; fused convolutions:

conv_weight = torch.cat((conv_weight_1, conv_weight_2), dim=0)
conv_bias = torch.cat((conv_bias_1, conv_bias_2))

batch = torch.cat((image_1,image_2), dim=1)

res_fused = torch.nn.functional.conv2d(batch, conv_weight, conv_bias, groups=2)
res_fused = res_fused.view(2, 4, 48, 48)

res_fused.allclose(res_not_fused) # &lt;- they are actually equal
</code></pre>
<p>I needed to use <a href="https://pytorch.org/docs/master/torch.html#torch.allclose" rel="nofollow noopener">torch.allclose</a> instead of <code>torch.equal</code> because there is a very small numerical difference between the two methods.</p>

ShahRukh: <p>Apologies, I just saw this! Thank you so much for the snippet! <img src="https://discuss.pytorch.org/images/emoji/apple/slightly_smiling_face.png?v=9" title=":slightly_smiling_face:" class="emoji" alt=":slightly_smiling_face:"></p>