Tensor advanced slicing/ indexing

80Ultraman: <p>I am working on a vision problem that requires shift in preprocessing.</p>
<p><strong>img</strong> is 3d tensor shaped [N, m, m].<br>
<strong>shiftmap</strong> is an int tensor containing starting index, shaped [N, 2].<br>
<strong>K</strong> is a constant.</p>
<p>I did the shift in a for-loop:</p>
<pre><code class="lang-auto">import torch

ret = torch.zeros(N, K, K)
for i in range(N):
    xx = shiftmap[i, 0]
    yy = shiftmap[i , 1]
    ret[i, ...] = img[i, xx: xx+K, yy: yy+K]

</code></pre>
<p>Is there an advanced Tensor slicing operation that can do the shift in O(1) instead of a for-loop?</p>
<p>Thanks a ton!</p>

KFrank: <p>Hi Ultraman!</p>
<aside class="quote no-group" data-username="80Ultraman" data-post="1" data-topic="170336" 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/80ultraman/48/43680_2.png" class="avatar"> 80Ultraman:</div>
<blockquote>
<p>Is there an advanced Tensor slicing operation that can do the shift</p>
</blockquote>
</aside>
<p>Yes.  One scheme would be to create index-tensor “templates” for the <code>xx</code><br>
and <code>yy</code> indices of <code>img</code> and index into those templates with <code>shiftmap</code> to get<br>
the actual <code>xx</code> and <code>yy</code> indices.  Then index into <code>img</code> itself with tensor indexing<br>
(together with a rather trivial tensor index for the batch (<code>N</code>) dimension) to get<br>
the final shifted result.</p>
<p>Here is an illustration:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'1.13.0'
&gt;&gt;&gt;
&gt;&gt;&gt; _ = torch.manual_seed (2023)
&gt;&gt;&gt;
&gt;&gt;&gt; N = 5
&gt;&gt;&gt; m = 10
&gt;&gt;&gt;
&gt;&gt;&gt; K = 3
&gt;&gt;&gt;
&gt;&gt;&gt; L = m - K + 1   # limit of allowed shiftmap (plus 1)
&gt;&gt;&gt;
&gt;&gt;&gt; img = torch.randn (N, m, m)            # sample img
&gt;&gt;&gt; shiftmap = torch.randint (L, (N, 2))   # sample shiftmap
&gt;&gt;&gt;
&gt;&gt;&gt; # with for-loop
&gt;&gt;&gt; ret = torch.zeros(N, K, K)
&gt;&gt;&gt; for i in range(N):
...     xx = shiftmap[i, 0]
...     yy = shiftmap[i , 1]
...     ret[i, ...] = img[i, xx: xx+K, yy: yy+K]
...
&gt;&gt;&gt; # loop-free version
&gt;&gt;&gt; ind0 = torch.arange (N)[:, None, None].repeat (1, K, K)
&gt;&gt;&gt; ind2 = torch.arange (m).repeat (m + 1).reshape (m, m + 1)[0:L, 0:K].unsqueeze (1).expand (L, K, K)
&gt;&gt;&gt; ind1 = ind2.transpose (1, 2)
&gt;&gt;&gt; retB = img[ind0, ind1[shiftmap[:, 0]], ind2[shiftmap[:, 1]]]
&gt;&gt;&gt;
&gt;&gt;&gt; torch.equal (retB, ret)
True
</code></pre>
<aside class="quote no-group">
<blockquote>
<p>in O(1) instead of a for-loop?</p>
</blockquote>
</aside>
<p>Technically speaking, you can’t do this in O(1) time (even with a loop-free<br>
algorithm).  No matter how you slice it you have to create and populate your<br>
result tensor of shape <code>[N, K, K]</code>, so the best your algorithm can be is<br>
O (N * K**2).  (Of course, pushing those operations out of python loops down<br>
into pytorch tensor operations will offer a huge speed-up.)</p>
<p>Best.</p>
<p>K. Frank</p>

tianhaowu0122: <pre><code class="lang-auto">img_flat = img.view(N, -1)

patch_index = (torch.arange(K)[None, :] + m * torch.arange(K)[:, None]).view(1, -1)

shift_patch_index = m * shiftmap[:, 0] + shiftmap[:, 1]

patch_index = patch_index + shift_patch_index[:, None]

ret_ = img_flat.gather(dim=1, index=patch_index).view(N, K, K)
</code></pre>
<p>This might be faster than direct indexing. I believe you could apply unfold to even speedup more.</p>