Optimizing Performance by Reducing Redundancy in Looping through PyTorch Tensors

Samet_Copur: <p>Dear PyTorch Community,</p>
<p>I’m currently working on a project where I need to populate a tensor <strong>ws_expanded</strong> based on certain conditions using a nested loop structure. However, I’ve noticed that reconstructing this loop each time incurs a significant computational cost. Here’s the relevant portion of the code for context:</p>
<p>ws_expanded = torch.empty_like(y_rules, device=y_rules.device, dtype=y_rules.dtype)<br>
index = 0</p>
<p>for col, rules in enumerate(rule_paths):<br>
for rule in rules:<br>
mask = y_rules[:, col] == rule<br>
ws_expanded[mask, col] = ws[index][0]<br>
index += 1</p>
<p>As you can see, the nested loops iterate over <strong>rule_paths</strong> and <strong>rules</strong> to populate <strong>ws_expanded</strong> based on certain conditions. However, as the size of the tensors increases, reconstructing this loop becomes prohibitively expensive.</p>
<p>I’m exploring ways to optimize this process. Specifically, I’m wondering if there’s a way to assign the weights (<strong>ws</strong>) to <strong>ws_expanded</strong> permanently using pointers in PyTorch, thus eliminating the need to reconstruct the loop every time.</p>
<p>Could you please advise on the best approach to handle this situation? Any insights or alternative strategies would be greatly appreciated.</p>
<p>Thank you in advance for your assistance.</p>

KFrank: <p>Hi Samet!</p>
<aside class="quote no-group" data-username="Samet_Copur" data-post="1" data-topic="201396" 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/samet_copur/48/69462_2.png" class="avatar"> Samet_Copur:</div>
<blockquote>
<p>Specifically, I’m wondering if there’s a way to assign the weights (<strong>ws</strong>) to <strong>ws_expanded</strong> permanently using pointers in PyTorch, thus eliminating the need to reconstruct the loop every time.</p>
</blockquote>
</aside>
<p>I’m guessing you mean something like the following:</p>
<p>You have some values in <code>ws</code> and you assign them into <code>ws_expanded</code> with<br>
a loop.  If you could have done this with “pointers,” then if the values in <code>ws</code><br>
were to change, the corresponding values in <code>ws_expanded</code> would also<br>
change <em>without</em> rerunning the loop (because the pointers would point to<br>
the new values).</p>
<p>If this is what you want, you can do something conceptually similar by running<br>
your loop once to pre-compute the appropriate indices and use those indices<br>
to assign values of <code>ws</code> into <code>ws_expanded</code>.</p>
<p>When the values of <code>ws</code> change, you do have to reassigned them into<br>
<code>ws_expanded</code>, but you don’t have to recompute the indices, so you don’t<br>
have to rerun the loop.</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; print (torch.__version__)
2.2.2
&gt;&gt;&gt;
&gt;&gt;&gt; ws_expanded = torch.zeros (2, 3, 4)
&gt;&gt;&gt; ws = torch.arange (4.).reshape (2, 2) + 1
&gt;&gt;&gt;
&gt;&gt;&gt; # pre-compute indices somehow (perhaps using a loop)
&gt;&gt;&gt; ind0 = torch.tensor ([[1, 1], [1, 1]])   # note that the index-tensors have the same shape as ws
&gt;&gt;&gt; ind1 = torch.tensor ([[0, 0], [2, 2]])
&gt;&gt;&gt; ind2 = torch.tensor ([[0, 1], [2, 3]])
&gt;&gt;&gt;
&gt;&gt;&gt; ws_expanded[ind0, ind1, ind2] = ws       # use indices to assign the original ws values into ws_expanded
&gt;&gt;&gt;
&gt;&gt;&gt; ws
tensor([[1., 2.],
        [3., 4.]])
&gt;&gt;&gt; ws_expanded
tensor([[[0., 0., 0., 0.],
         [0., 0., 0., 0.],
         [0., 0., 0., 0.]],

        [[1., 2., 0., 0.],
         [0., 0., 0., 0.],
         [0., 0., 3., 4.]]])
&gt;&gt;&gt;
&gt;&gt;&gt; ws += 100                                # new values for ws
&gt;&gt;&gt;
&gt;&gt;&gt; ws_expanded[ind0, ind1, ind2] = ws       # reuse indices with the new ws values (no loop)
&gt;&gt;&gt;
&gt;&gt;&gt; ws
tensor([[101., 102.],
        [103., 104.]])
&gt;&gt;&gt; ws_expanded
tensor([[[  0.,   0.,   0.,   0.],
         [  0.,   0.,   0.,   0.],
         [  0.,   0.,   0.,   0.]],

        [[101., 102.,   0.,   0.],
         [  0.,   0.,   0.,   0.],
         [  0.,   0., 103., 104.]]])
</code></pre>
<p>Best.</p>
<p>K. Frank</p>