Tensor advanced slicing/ indexing

1

I am working on a vision problem that requires shift in preprocessing.

img is 3d tensor shaped [N, m, m].
shiftmap is an int tensor containing starting index, shaped [N, 2].
K is a constant.

I did the shift in a for-loop:

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]

Is there an advanced Tensor slicing operation that can do the shift in O(1) instead of a for-loop?

Thanks a ton!

2

Hi Ultraman!

Yes. One scheme would be to create index-tensor “templates” for the xx
and yy indices of img and index into those templates with shiftmap to get
the actual xx and yy indices. Then index into img itself with tensor indexing
(together with a rather trivial tensor index for the batch (N) dimension) to get
the final shifted result.

Here is an illustration:

>>> import torch
>>> torch.__version__
'1.13.0'
>>>
>>> _ = torch.manual_seed (2023)
>>>
>>> N = 5
>>> m = 10
>>>
>>> K = 3
>>>
>>> L = m - K + 1   # limit of allowed shiftmap (plus 1)
>>>
>>> img = torch.randn (N, m, m)            # sample img
>>> shiftmap = torch.randint (L, (N, 2))   # sample shiftmap
>>>
>>> # with for-loop
>>> 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]
...
>>> # loop-free version
>>> ind0 = torch.arange (N)[:, None, None].repeat (1, K, K)
>>> ind2 = torch.arange (m).repeat (m + 1).reshape (m, m + 1)[0:L, 0:K].unsqueeze (1).expand (L, K, K)
>>> ind1 = ind2.transpose (1, 2)
>>> retB = img[ind0, ind1[shiftmap[:, 0]], ind2[shiftmap[:, 1]]]
>>>
>>> torch.equal (retB, ret)
True

Technically speaking, you can’t do this in O(1) time (even with a loop-free
algorithm). No matter how you slice it you have to create and populate your
result tensor of shape [N, K, K], so the best your algorithm can be is
O (N * K**2). (Of course, pushing those operations out of python loops down
into pytorch tensor operations will offer a huge speed-up.)

Best.

K. Frank

2 Likes