Masking the intermediate 5D Conv2D output

username: abhinav1kumar
cooked: <p>Hi PyTorch Team,<br>
I have an input tensor of shape (B, C_in, H, W), conv (C_in, C_out, K, K, Pad= K//2) and a mask of shape (B, C_in, H, W) . I wanted to multiply the intermediate 5D output with my mask before carrying out the final summation.</p>
<p>The pictures of the usual and convolution with masking is as follows:<br>
The usual convolution operates as follows:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/0/60994f2ed1f8c34ee6e83741e2e87fca0a1b4655.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/60994f2ed1f8c34ee6e83741e2e87fca0a1b4655" title="conv"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/0/60994f2ed1f8c34ee6e83741e2e87fca0a1b4655_2_690x388.jpeg" alt="conv" data-base62-sha1="dMyiEqPL9jAWQiOcToQ4fMdFhAh" width="690" height="388" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/0/60994f2ed1f8c34ee6e83741e2e87fca0a1b4655_2_690x388.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/6/0/60994f2ed1f8c34ee6e83741e2e87fca0a1b4655_2_1035x582.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/6/0/60994f2ed1f8c34ee6e83741e2e87fca0a1b4655.jpeg 2x" data-dominant-color="ECE5E1"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">conv</span><span class="informations">1200×675 71.8 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>However, I want to mask out the intermediate 5D output with the mask matrix before the summation across the channels as shown in the following updated figure<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/e/0/e04339b4988f8da74d3d8b946a37a1d83994b7b8.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/e04339b4988f8da74d3d8b946a37a1d83994b7b8" title="conv_2"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/e/0/e04339b4988f8da74d3d8b946a37a1d83994b7b8_2_690x388.jpeg" alt="conv_2" data-base62-sha1="vZUYXgST6obCOJ2cVlfHlZJOHUs" width="690" height="388" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/e/0/e04339b4988f8da74d3d8b946a37a1d83994b7b8_2_690x388.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/e/0/e04339b4988f8da74d3d8b946a37a1d83994b7b8_2_1035x582.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/e/0/e04339b4988f8da74d3d8b946a37a1d83994b7b8.jpeg 2x" data-dominant-color="EAE2DF"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">conv_2</span><span class="informations">1200×675 78.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>Figures’ Courtesy: <a href="https://developpaper.com/how-to-customize-convolution-kernel-weight-parameters-by-pytorch/" rel="noopener nofollow ugc">developpaper.com</a></p>
<p>I tried getting the intermediate 5D tensor first. However, I am not able to extract the intermediate 5D tensor of shape (B, C_out, C_in, H, W). Following are the things I tried:</p>
<ul>
<li><a class="mention" href="/u/ptrblck">@ptrblck</a> <a href="https://discuss.pytorch.org/t/convolution-operation-without-the-final-summation/56466/4">solution</a>. However, this does not give me the 5D tensor on my end</li>
</ul>
<pre><code class="lang-auto">in_channels = 5
conv = nn.Conv2d(in_channels, 20, 3, groups= in_channels, padding= 1)
x = torch.randn(1, in_channels, 64, 64)
output = conv(x)
print(output.shape)
# torch.Size([1, 20, 64, 64])
</code></pre>
<ul>
<li>I also tried unfolding it</li>
</ul>
<pre><code class="lang-auto">inp = torch.randn(1, in_channels, 10, 12)
w   = torch.randn(2, in_channels, 3, 3)
inp_unf = torch.nn.functional.unfold(inp, (3, 3))
print(inp_unf.shape)
out_unf = inp_unf.transpose(1, 2).matmul(w.view(w.size(0), -1).t()).transpose(1, 2)
print(out_unf.shape)
out = out_unf.view(1, 2, 8, 10)
# torch.Size([1, 2, 80])
</code></pre>
<p>Since I could not find any working solution on the forums, I decided to post here. It would be great if you could answer this question. In case it is a duplicate post, would you mind pointing me to the correct answer.</p>

username: ptrblck
cooked: <p><a href="https://discuss.pytorch.org/t/visualizing-the-outputs-of-kernels-instead-of-the-outputs-of-filters/137266/4">This post</a> of using the grouped conv approach and reducing the output manually afterwards might be helpful.</p>

username: abhinav1kumar
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Thank you for the pointer. I implemented the masked convolution using the <code>unfold</code> function as shown below. However, as you might have guessed, if the convolution has <code>group</code> argument other than <code>1</code>, this code would not work.</p>
<p>I had a related question. Can I extend the following code based on <code>unfold</code> for conv2d with <code>group</code> arguments as well?</p>
<pre><code class="lang-auto">def mask_convolution(input, kernel_wt, padding= 0, stride= 1, mask= None, mask_patches= False, debug= False):
    if debug:
        print("\nInside function ...")

    # kernel_wt with shape Cout x Cin x K x K
    out_channels, in_channels, K, K = kernel_wt.shape
    batch, _, H, W                  = input.shape
    Hout = (H + 2*padding - 2*(K//2))//stride
    Wout = (W + 2*padding - 2*(K//2))//stride

    norm_kernel_wt        = kernel_wt

    inp_unf               = torch.nn.functional.unfold(inp, (K, K), padding= padding, stride= stride)  # B x Cin*K*K X Hout*Wout
    inp_unf_interm        = inp_unf.transpose(1, 2).unsqueeze(1)  # B x 1 x Hout*Wout x Cin*K*K

    norm_kernel_wt_interm = norm_kernel_wt.view(out_channels, -1).unsqueeze(0).unsqueeze(2) # 1 x Cout x 1 x Cin*K*K

    out_unf = torch.mul(inp_unf_interm, norm_kernel_wt_interm).reshape(batch, out_channels, inp_unf_interm.shape[2], in_channels, -1) # B x Cout x Hout*Wout x Cin x K*K
    output  = torch.sum(out_unf, dim= 4).transpose(2, 3).reshape(batch, out_channels, in_channels, Hout, Wout) # B x Cout x Cin x Hout x Wout
    if debug:
        print("out_unf.shape before sum", out_unf.shape)
        print("output.shape before sum", output.shape)

    if mask_patches:
        output_final  = output * (mask.unsqueeze(1))
    else:
        output_final  = output

    output_final = torch.sum(output_final, dim= 2) # B x Cout x Hout x Wout
    if debug:
        print("output_final shape", output_final.shape)

    return output_final
</code></pre>

username: ptrblck
cooked: <p>The idea of the linked post is to use the grouped convolution in order to avoid the reduction in the <code>in_channels</code> dimension so that you can manipulate the activations manually before applying the sum.<br>
But yes, <code>unfold</code> should also work.</p>

username: abhinav1kumar
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> I followed your <a href="https://discuss.pytorch.org/t/visualizing-the-outputs-of-kernels-instead-of-the-outputs-of-filters/137266/4">post</a> to get the following.</p>
<p>However, because of the loops, the solution is slow when we have multiple groups. How should I get the 5D output <strong>without the loops</strong> for the case when the <code>mask_convolution</code> function itself has groups &gt; 1</p>
<pre><code class="lang-python">def mask_convolution(input, kernel_wt, padding= 0, stride= 1, \
        groups= 1, mask= None, mask_patches= False, debug= False):

    output       = cast_to_cpu_cuda_tensor(torch.zeros((batch, out_channels, Hout, Wout)), reference_tensor= input)
    index_input  = torch.arange(groups+1)*(in_channels//groups)
    index_output = torch.arange(groups+1)*(out_channels//groups)

    for i in range(groups):
        inp_group             = input[:, index_input[i]:index_input[i+1]]
        norm_kernel_wt_interm = norm_kernel_wt[index_output[i]:index_output[i+1]]\
                                        .reshape(out_channels_grp * in_channels_grp, 1, K, K)
        
        # B x Cout*Cin x Hout x Wout
        output_grp            = F.conv2d(inp_group, norm_kernel_wt_interm, padding= padding, stride= stride, groups= in_channels_grp)
        output_grp            = output_grp.view(batch, in_channels_grp, out_channels_grp, Hout, Wout)\
                                        .permute(0, 2, 1, 3, 4)\
                                        .reshape(batch, out_channels_grp, in_channels_grp, Hout, Wout)
        if mask_patches:
             output_final  = output_grp * (mask.unsqueeze(1))
        else:
             output_final  = output_grp

        output_final = torch.sum(output_final, dim= 2) # B x Cout x Hout x Wout

        output[:, index_output[i]:index_output[i+1]] = output_div_by_energy

    return output

</code></pre>

username: ptrblck
cooked: <p>I’m not sure why the <code>groups</code> loop is needed as my code snippet doesn’t use it.</p>

username: abhinav1kumar
cooked: <p>Thank you for your reply.</p>
<p>I wanted to get the intermediate 5D output for the original convolution with groups&gt;1.</p>
<p>e.g. Let us use an example similar to your <a href="https://discuss.pytorch.org/t/visualizing-the-outputs-of-kernels-instead-of-the-outputs-of-filters/137266/3">post</a>. Consider an input activation in the shape <code>[batch_size=1, channels=6, height=24, width=24]</code> and a standard conv layer with a <code>weight</code> tensor in the shape <code>[out_channels=6, in_channels=2, height=3, width=3]</code></p>
<p>We want to get the intermediate 5D outputs of the result of <code>F.conv2d(input, weight, </code><strong>groups=3</strong><code>, stride=1, bias= None, padding= 1)</code> . So, I kept the group loop to loop over the three slices of <code>input</code> and <code>output</code> in accordance with your <a href="https://discuss.pytorch.org/uploads/default/original/3X/a/d/adf32e382273800dbda282c878eba5db1628ea89.jpeg">figure</a>. With mask = identity, the output produced by my function is same as the <code>F.conv2d(input, weight,</code><strong>groups=3</strong><code>, stride=1)</code>. However, my function is orders of magnitude slower to run because of the loop.</p>
<p>Can we get the intermediate 5D output with <strong>groups=3</strong> without using a loop?</p>