Batched Conv2d with filters "grouped" in batch dimension

username: Tycho_van_der_Oudera
cooked: <p>Hi pytorch community,</p>
<p>For a research project, I am trying to apply different sets of filters to elements in the batch dimension.<br>
That is, I want to group the batch into <code>G</code> different sets and apply different filters to the batch dim group, in parallel. The extreme case where the number of groups <code>G=N</code> equals the batch size would mean we have a separate filter for each batch element.</p>
<p>input dimension: <code>N x C x H x W</code><br>
filter dimension: <code>G x C_out x C_in x H_filter x W_filter</code><br>
output dimension: <code>N x C_out x H_out x W_out</code></p>
<p>I believe this can be done by temporarily treating batch elements as channels and using the <code>groups=</code> argument of Conv2d to do so. However, it does not yield the correct result.</p>
<p>The code below seems to work for <code>G=1</code> (normal convolution with same filter applied to all batch elements) and <code>G=N</code> (different filter for each element in the batch). However, it does not seem to work for some <code>G</code> in between 1 and N.</p>
<pre><code class="lang-python">N, C, H, W = 128, 3, 28, 28
CC = 64

for groups in (1, 2, N):
    # dummy data
    x = torch.randn(N, C, H, W)
    w = torch.randn(N, CC, C, 7, 7)

    sub_size = N // groups

    x = torch.randn(N, C, H, W)  
    w = torch.randn(groups, CC, C, 7, 7) 

    # desired computation
    out = []
    for i in range(groups):
        sub_x = x[i*sub_size:i*sub_size + sub_size]
        sub_w = w[i]

        sub_out = F.conv2d(sub_x, sub_w)

        out.append(sub_out)
    out = torch.cat(out)

    # first attempt exploiting groups to perform the computation in parallel
    w_fast = w.view(groups*CC, C, 7, 7)
    x_fast = x.view(sub_size, groups*C, H, W)

    out_grouped = F.conv2d(x_fast, w_fast, groups=groups)
    out_grouped = out_grouped.view(N, CC, *out_grouped.shape[-2:])

    # should be the same:
    print((out - out_grouped).abs().max())
&gt;&gt; tensor(0.) # &lt;-- good
&gt;&gt; tensor(91.9094) # &lt;-- bad, G=2 fails.
&gt;&gt; tensor(4.5776e-05) # &lt;-- good
</code></pre>
<p>What am I missing here?</p>
<p>Many thanks in advance!</p>
<p>Tycho</p>

username: KFrank
cooked: <p>Hi Tycho!</p>
<aside class="quote no-group" data-username="Tycho_van_der_Oudera" data-post="1" data-topic="166861" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/tycho_van_der_oudera/48/55154_2.png" class="avatar"> Tycho_van_der_Oudera:</div>
<blockquote>
<p>I believe this can be done by temporarily treating batch elements as channels and using the <code>groups=</code> argument of Conv2d to do so.</p>
</blockquote>
</aside>
<p>This is exactly correct – you may merge the batch and channels dimensions<br>
together and then use <code>conv2d()</code> with its <code>groups</code> argument.</p>
<p>However, in order to interleave the grouped batch elements and the channels<br>
properly, you will need to insert a <code>transpose()</code> into the manipulations you use<br>
to create <code>x_fast</code>.  Note <code>x_fastB</code> and <code>out_groupedB</code> in this tweaked version<br>
of the code you posted:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

_ = torch.manual_seed (2022)

N, C, H, W = 128, 3, 28, 28
CC = 64

for groups in (1, 2, 16, 64, N):   # add a couple more group sizes
    sub_size = N // groups
    
    x = torch.randn(N, C, H, W)  
    w = torch.randn(groups, CC, C, 7, 7) 
    
    # desired computation
    out = []
    for i in range(groups):
        sub_x = x[i*sub_size:i*sub_size + sub_size]
        sub_w = w[i]
        
        sub_out = torch.nn.functional.conv2d (sub_x, sub_w)
        
        out.append(sub_out)
    out = torch.cat(out)
    
    # first attempt exploiting groups to perform the computation in parallel
    w_fast = w.view(groups*CC, C, 7, 7)
    x_fast = x.view(sub_size, groups*C, H, W)
    
    # need transpose() to "interleave" the input groups with one another properly
    x_fastB = x.reshape (groups, sub_size, C, *x.shape[-2:]).transpose (0, 1).reshape (sub_size, groups * C, *x.shape[-2:])
    
    out_grouped = torch.nn.functional.conv2d (x_fast, w_fast, groups = groups)
    out_grouped = out_grouped.view(N, CC, *out_grouped.shape[-2:])
    
    out_groupedB = torch.nn.functional.conv2d (x_fastB, w_fast, groups = groups)
    # and need transpose() to "deinterleave" the output groups
    out_groupedB = out_groupedB.reshape (sub_size, groups, CC, *out_groupedB.shape[-2:]).transpose (0, 1).reshape (N, CC, *out_groupedB.shape[-2:])
    
    # should be the same:
    print ('groups:', groups)
    print ('check out_grouped:', (out - out_grouped).abs().max())
    print ('check out_groupedB:', (out - out_groupedB).abs().max())
</code></pre>
<p>Here is its output:</p>
<pre><code class="lang-plaintext">1.13.0
groups: 1
check out_grouped: tensor(0.)
check out_groupedB: tensor(0.)
groups: 2
check out_grouped: tensor(88.9461)
check out_groupedB: tensor(0.)
groups: 16
check out_grouped: tensor(91.0194)
check out_groupedB: tensor(4.5776e-05)
groups: 64
check out_grouped: tensor(92.4197)
check out_groupedB: tensor(3.8147e-05)
groups: 128
check out_grouped: tensor(3.8147e-05)
check out_groupedB: tensor(3.8147e-05)
</code></pre>
<p>Best.</p>
<p>K. Frank</p>