How to use matrix manipulation to make this faster?

Charlie_Li: <p>Hi, I hope to use some kind of matrix manipulation to make this code faster:</p>
<pre><code class="lang-auto">for i in range(dl_dw.size(2)):
            for j in range(dl_dw.size(3)):
                dl_dx[:, :, i//self.scale_factor, j//self.scale_factor] += dl_dw[:, :, i, j]
</code></pre>
<p>Is it possible? Thanks!</p>

AlphaBetaGamma96: <p>When sharing code can you share a minimal reproducible example, not just the function?</p>
<p>So share an example input and output for this function to show the expected behavior.</p>

Charlie_Li: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/6/b695698bb163600c0c0a8744c642534f97401fc6.png" data-download-href="https://discuss.pytorch.org/uploads/default/b695698bb163600c0c0a8744c642534f97401fc6" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/b/6/b695698bb163600c0c0a8744c642534f97401fc6.png" alt="image" data-base62-sha1="q3d0RHUCqabcz0oIpe3Y7YG7dVs" width="671" height="500" data-dominant-color="FEEFEF"><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">image</span><span class="informations">753×561 13.5 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><br>
Hi, thanks for the reply! The idea is to reduce the size of matrix by summing the submatrix in the original matrix. In this example, the scale factor is 2, and each matrix is a 4D tensor, but we only need to consider the last two dimensions</p>

Matias_Vasquez: <pre><code class="lang-python">a = torch.tensor([[1.,2,5,6], [3,4,7,8], [1,1,2,2], [1,1,3,4]], requires_grad=True).view(1,1,4,4)
</code></pre>
<p>You could do one of the following</p>
<ul>
<li>Use a convolutional layer with custom kernel</li>
</ul>
<pre><code class="lang-python">conv = torch.nn.Conv2d(1, 1, 2, 2, bias=False)
kernel = torch.tensor([[1., 1.],
                       [1., 1.]]).unsqueeze(0).unsqueeze(0)
with torch.no_grad():
    conv.weight = nn.Parameter(kernel)

print(conv(a))
# Output:
#tensor([[[[10., 26.],
#          [ 4., 11.]]]], grad_fn=&lt;ConvolutionBackward0&gt;)
</code></pre>
<ul>
<li>Use average pooling and multiply by the number of pixels inside the sliding window</li>
</ul>
<pre><code class="lang-python">avg = torch.nn.AvgPool2d(2,2)
print(avg(a)*4)
# Output:
#tensor([[[[10., 26.],
#          [ 4., 11.]]]], grad_fn=&lt;MulBackward0&gt;)
</code></pre>
<p>Hope this helps <img src="https://discuss.pytorch.org/images/emoji/apple/smile.png?v=12" title=":smile:" class="emoji" alt=":smile:" loading="lazy" width="20" height="20"></p>

Charlie_Li: <p>Thanks for the solution! Yes, this would work, I also wonder if there’s a way to achieve this without using the nn modules but only the tensor operations (without loops) <img src="https://discuss.pytorch.org/images/emoji/apple/blush.png?v=12" title=":blush:" class="emoji" alt=":blush:" loading="lazy" width="20" height="20"></p>

Matias_Vasquez: <p>You could do something like this</p>
<pre><code class="lang-python">a = torch.tensor([[[[1.,2,5,6], [3,4,7,8], [1,1,2,2], [1,1,3,4]]]], requires_grad=True)
B, C, H, W = a.shape

kernel = 2
out_h = H // kernel
out_w = W // kernel

rows = torch.arange(0, H, kernel).repeat(out_h)
cols = torch.arange(0, W, kernel).repeat_interleave(out_w)

x0y0 = a[..., cols+0, rows+0]
x0y1 = a[..., cols+0, rows+1]
x1y0 = a[..., cols+1, rows+0]
x1y1 = a[..., cols+1, rows+1]

out = (x0y0 + x0y1 + x1y0 + x1y1).reshape(B, C, out_h, out_w)
print(out)
</code></pre>
<p>If you want a different kernel size, then you need to add the necessary <code>xMyN=a[..., cols+M, rows+N]</code> and add them all up.</p>
<p>But this workaround is very prone to errors. I would rather use any of the other solutions given, as they are more robust and will not break that easily.</p>