Vectorized 6 loops or convert to Pytorch operations for efficient computation

username: jpainam
cooked: <p>Hi.<br>
I would like to rewrite this code in a Pytorch-like and vectorized form for efficient computation.</p>
<pre><code class="lang-python">import torch

num_faces = 10
max_m = 5
C2 = torch.rand(num_faces,max_m,max_m,max_m)
C3 = torch.rand(num_faces,max_m,max_m,max_m)


# This is the naive slow implementation with 6 loops, but is correct
D_abc_correct = torch.zeros(num_faces,max_m,max_m,max_m)
for a in range(max_m):
    for b in range(max_m):
        for c in range(max_m):
            for i2 in range(a+1):
                for j2 in range(b+1):
                    for k2 in range(c+1):
                        D_abc_correct[:,a,b,c] += C2[:,i2,j2,k2]*C3[:,a-i2,b-j2,c-k2]
</code></pre>
<p>A slightly optimized version with 3 loops, also correct</p>
<pre><code class="lang-python"># 
D_abc = torch.zeros(num_faces,max_m,max_m,max_m)
for a in range(max_m):
    for b in range(max_m):
        for c in range(max_m):
            D_abc[:,a,b,c] = (C2[:,:(a+1),:(b+1),:(c+1)]*(C3[:,:(a+1),:(b+1),:(c+1)]).flip([1,2,3])).sum([1,2,3])

</code></pre>
<p>Any changes after that lead to some <code>mismatch shape</code></p>
<pre><code class="lang-python"># Check the results

print(torch.allclose(D_abc, D_abc_correct))
# True
print(D_abc.shape)
#torch.Size([10, 5, 5, 5])
</code></pre>
<p>Thanks</p>

username: jpainam
cooked: <p>On StackOverflow <a href="https://stackoverflow.com/questions/76640108/vectorized-6-loops-or-convert-to-pytorch-numpy-operations-for-efficient-computat" rel="noopener nofollow ugc">python - Vectorized 6 loops or convert to Pytorch/numpy operations for efficient computation - Stack Overflow</a></p>

username: Osama_Rais
cooked: <p>It seems like a simplified minimal problem in 1D is as follows</p>
<pre><code class="lang-auto">max_m = 10

c2 = torch.rand(max_m)
c3 = torch.rand(max_m)
d =torch.zeros(max_m)

for a in range(max_m):
    d[a] = (c2[:(a+1)]*(c3[:a+1].flip([0]))).sum(0)
</code></pre>
<p>It seems like you want to perform an inner product over the two tensors, but with varying lengths.</p>
<p>I couldn’t solve it yet but am putting this here in case it is useful as a minimal example for others.</p>

username: KFrank
cooked: <p>Hi Jean Paul!</p>
<aside class="quote no-group" data-username="jpainam" data-post="1" data-topic="183666" 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/jpainam/48/5870_2.png" class="avatar"> jpainam:</div>
<blockquote>
<p>I would like to rewrite this code in a Pytorch-like and vectorized form for efficient computation.</p>
</blockquote>
</aside>
<p>Here’s the best I could come up with for a loop-free version:</p>
<p>At the cost of materializing a large, nine-dimensional “contraction-kernel”<br>
tensor, you can replace the six for-loops with a single <code>einsum()</code> call.</p>
<p>Here is a script that implements this:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

_ = torch.manual_seed (2023)

num_faces = 10
max_m = 5
C2 = torch.rand(num_faces,max_m,max_m,max_m)
C3 = torch.rand(num_faces,max_m,max_m,max_m)

# This is the naive slow implementation with 6 loops, but is correct
D_abc_correct = torch.zeros(num_faces,max_m,max_m,max_m)
for a in range(max_m):
    for b in range(max_m):
        for c in range(max_m):
            for i2 in range(a+1):
                for j2 in range(b+1):
                    for k2 in range(c+1):
                        D_abc_correct[:,a,b,c] += C2[:,i2,j2,k2]*C3[:,a-i2,b-j2,c-k2]

# set up "contraction kernel"
inds = torch.arange (max_m)
kernel = ((inds.unsqueeze (0) + inds.unsqueeze (1)).unsqueeze (0) == inds.unsqueeze (-1).unsqueeze (-1)).float()

print ('check one slice of kernel:')
print (kernel[2])   # check one slice to see what it looks like

kernelProd = torch.einsum ('ail, bjm, ckn -&gt; abcijklmn', kernel, kernel, kernel)   # "outer-product" kernel
print ('kernelProd is big, but rather sparse:')
print ('numel():     ', kernelProd.numel())                                        # quite big
print ('num-zero:    ', (kernelProd == 0.0).sum().item())                          # with lots of zeros
print ('frac-nonzero:', (kernelProd != 0.0).sum().item() / kernelProd.numel())     # so rather sparse
kernelProd = kernelProd.unsqueeze (0).expand (num_faces, -1, -1, -1, -1, -1, -1, -1, -1, -1)   # expand() view with num_faces dimension

D_abc_B = torch.einsum ('eijk, eabcijklmn, elmn -&gt; eabc', C2, kernelProd, C3)

# check einsum result
print ('torch.allclose (D_abc_correct, D_abc_B):', torch.allclose (D_abc_correct, D_abc_B))
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.0.0
check one slice of kernel:
tensor([[0., 0., 1., 0., 0.],
        [0., 1., 0., 0., 0.],
        [1., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.],
        [0., 0., 0., 0., 0.]])
kernelProd is big, but rather sparse:
numel():      1953125
num-zero:     1949750
frac-nonzero: 0.001728
torch.allclose (D_abc_correct, D_abc_B): True
</code></pre>
<p>Note, you can reuse <code>kernelProd</code>, so even if you have to run this<br>
computation multiple times for differing values of the input tensors,<br>
<code>C2</code> and <code>C3</code>, you only have to construct <code>kernelProd</code> once.</p>
<p>Best.</p>
<p>K. Frank</p>

username: Osama_Rais
cooked: <p>It seems like a further optimization can be performed using convolutions.</p>
<p>The issue is now vectorizing the operation over num_faces since performing convolutions over a set of pairs of tensors doesn’t seem to be supported.</p>
<p>Here is an implementation of the minimal problem in 1D:</p>
<pre><code class="lang-auto">import torch
import torchaudio
torch.__version__

max_m = 10

c2 = torch.rand(max_m)
c3 = torch.rand(max_m)
d_two =torch.zeros(max_m)
d_one =torch.zeros(max_m)

# Naive implementation with two loops
d_two = torch.zeros(max_m)
for a in range(max_m):
    for i in range(a+1):
        d_two[a] += c2[i]*c3[a-i]
        
# Better implementation with innermost loop vectorized
for a in range(max_m):
    d_one[a] = (c2[:(a+1)]*(c3[:a+1].flip([0]))).sum(0)
    
# No loops, with convolution operation
d_conv = torchaudio.functional.convolve(c2,c3,'full')[0:max_m]

print(torch.allclose(d_two,d_one))
print(torch.allclose(d_two,d_conv))
</code></pre>

username: Osama_Rais
cooked: <p>This is quite an elegant solution to the problem. Unfortunately it seems that it runs much slower than even the 6 loop implementation:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/c/7c525cf8d5c1d3f6028af4e49e094d30a1ff57c4.png" data-download-href="https://discuss.pytorch.org/uploads/default/7c525cf8d5c1d3f6028af4e49e094d30a1ff57c4" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/c/7c525cf8d5c1d3f6028af4e49e094d30a1ff57c4_2_690x434.png" alt="image" data-base62-sha1="hJNFsajBfiMx7fYO5IBtKIsr1Xe" width="690" height="434" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/c/7c525cf8d5c1d3f6028af4e49e094d30a1ff57c4_2_690x434.png, https://discuss.pytorch.org/uploads/default/optimized/3X/7/c/7c525cf8d5c1d3f6028af4e49e094d30a1ff57c4_2_1035x651.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/7/c/7c525cf8d5c1d3f6028af4e49e094d30a1ff57c4_2_1380x868.png 2x" data-dominant-color="EDEEED"><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">1500×944 120 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>Perhaps <a class="mention" href="/u/jpainam">@jpainam</a>  might want to try running on their GPU to compare results?</p>

username: eqy
cooked: <p>I’m curious if there is some overhead associated with the first time <code>torch.einsum</code> is used and whether a warmup would be necessary. Does <code>%%timeit</code> do a warmup iteration?</p>

username: Osama_Rais
cooked: <p><a class="mention" href="/u/kfrank">@KFrank</a> provided a great implementation. I am putting in order all the snippets, with my optimization at the end, which runs fastest.</p>
<pre><code class="lang-auto">import torch
print (torch.__version__)
_ = torch.manual_seed (2023)

num_faces = 1000
max_m = 5
C2 = torch.rand(num_faces,max_m,max_m,max_m)
C3 = torch.rand(num_faces,max_m,max_m,max_m)

# set up "contraction kernel"
inds = torch.arange (max_m)
kernel = ((inds.unsqueeze (0) + inds.unsqueeze (1)).unsqueeze (0) == inds.unsqueeze (-1).unsqueeze (-1)).float()
kernelProd = torch.einsum ('ail, bjm, ckn -&gt; abcijklmn', kernel, kernel, kernel)
kernelProd = kernelProd.unsqueeze (0).expand (num_faces, -1, -1, -1, -1, -1, -1, -1, -1, -1)   # expand() view with num_faces dimension





D_abc_correct = torch.zeros(num_faces,max_m,max_m,max_m)
for a in range(max_m):
    for b in range(max_m):
        for c in range(max_m):
            for i2 in range(a+1):
                for j2 in range(b+1):
                    for k2 in range(c+1):
                        D_abc_correct[:,a,b,c] += C2[:,i2,j2,k2]*C3[:,a-i2,b-j2,c-k2]



D_abc = torch.zeros(num_faces,max_m,max_m,max_m)
for a in range(max_m):
    for b in range(max_m):
        for c in range(max_m):
            D_abc[:,a,b,c] = (C2[:,:(a+1),:(b+1),:(c+1)]*(C3[:,:(a+1),:(b+1),:(c+1)]).flip([1,2,3])).sum([1,2,3])




D_abc_B = torch.einsum ('eijk, eabcijklmn, elmn -&gt; eabc', C2, kernelProd, C3)



D_abc_C = torch.einsum ('eijk, ail, bjm, ckn, elmn -&gt; eabc', C2, kernel, kernel, kernel, C3)



print(torch.allclose (D_abc_correct, D_abc))
print(torch.allclose (D_abc_correct, D_abc_B))
print(torch.allclose (D_abc_correct, D_abc_C))

</code></pre>
<p>All the results are correct:</p>
<pre><code class="lang-auto">2.0.1
True
True
True
</code></pre>
<p>Timing the different versions:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/a/da31989474a02d0487c1a0986bc130a2e37df938.png" data-download-href="https://discuss.pytorch.org/uploads/default/da31989474a02d0487c1a0986bc130a2e37df938" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/a/da31989474a02d0487c1a0986bc130a2e37df938_2_690x326.png" alt="image" data-base62-sha1="v8elxUfAzVbLwePzlJENDdH1MSs" width="690" height="326" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/a/da31989474a02d0487c1a0986bc130a2e37df938_2_690x326.png, https://discuss.pytorch.org/uploads/default/optimized/3X/d/a/da31989474a02d0487c1a0986bc130a2e37df938_2_1035x489.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/a/da31989474a02d0487c1a0986bc130a2e37df938_2_1380x652.png 2x" data-dominant-color="F1F1F1"><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">2182×1034 244 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>

username: Osama_Rais
cooked: <p>It doesn’t seem to be the case. I put an optimized version of <a class="mention" href="/u/kfrank">@KFrank</a> 's code, which now runs faster than the three loop implementation.</p>
<p>I swapped the order when running '%%timeit ’ to run KFrank’s code and its optimized versions just to make sure, the results are the same</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/4/e/4e90810030d4fbafdda7811bf3755e9764a08b9d.png" data-download-href="https://discuss.pytorch.org/uploads/default/4e90810030d4fbafdda7811bf3755e9764a08b9d" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/4/e/4e90810030d4fbafdda7811bf3755e9764a08b9d_2_690x122.png" alt="image" data-base62-sha1="bd0Q8oeSjjREgYPtFtbagrh6LGR" width="690" height="122" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/4/e/4e90810030d4fbafdda7811bf3755e9764a08b9d_2_690x122.png, https://discuss.pytorch.org/uploads/default/optimized/3X/4/e/4e90810030d4fbafdda7811bf3755e9764a08b9d_2_1035x183.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/4/e/4e90810030d4fbafdda7811bf3755e9764a08b9d_2_1380x244.png 2x" data-dominant-color="F0EFEF"><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">2188×388 45.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></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/b/abcef5449eab4f7fcd491924dab053c2c7f0005d.png" data-download-href="https://discuss.pytorch.org/uploads/default/abcef5449eab4f7fcd491924dab053c2c7f0005d" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/a/b/abcef5449eab4f7fcd491924dab053c2c7f0005d_2_690x117.png" alt="image" data-base62-sha1="ovT3i3xLqGxgbLHQhVDyWiO4WWF" width="690" height="117" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/a/b/abcef5449eab4f7fcd491924dab053c2c7f0005d_2_690x117.png, https://discuss.pytorch.org/uploads/default/optimized/3X/a/b/abcef5449eab4f7fcd491924dab053c2c7f0005d_2_1035x175.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/a/b/abcef5449eab4f7fcd491924dab053c2c7f0005d_2_1380x234.png 2x" data-dominant-color="F3F2F2"><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">2218×378 45.2 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>

username: KFrank
cooked: <p>Hi Osama!</p>
<aside class="quote no-group" data-username="Osama_Rais" data-post="8" data-topic="183666" 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/osama_rais/48/61565_2.png" class="avatar"> Osama_Rais:</div>
<blockquote>
<pre><code class="lang-auto">D_abc_C = torch.einsum ('eijk, ail, bjm, ckn, elmn -&gt; eabc', C2, kernel, kernel, kernel, C3)
</code></pre>
</blockquote>
</aside>
<p>Yes, this is much better.  Materializing <code>kernelProd</code> is wasteful and<br>
unnecessary.</p>
<p>Best.</p>
<p>K. Frank</p>