How to quickly cut slices from cuda tensor wrt to another tensors values

Question

I have a torch cuda tensor A of shape (100, 80000, 4), and another cuda tensor B of shape (1, 80000, 1). What I want to do, is for each element i in the second dimension (from 0 to 79999), take the value of tensor B (which will be from 0 to 99 and will point out to which value in the first dimension of A to take.

An additional problem is that for this element of B (B[0, i, 0]), I want to take a slice from A that is A[lower_bound:upper_bound, i:i+1, :].

To sum it up, my tensor B has the indices of the centers of slices that I would like to cut from A. I am wondering, whether there is a way, to do what I am doing with the below code faster (eg. using cuda)

A  # tensor of shape (100, 80000, 4)
B  # tensor of shape (1, 80000, 1) which has 
k = 3  # width of the slice to take (or half of it to be exact)

A_list = []
for i in range(80000):
    lower_bound = max(0, B[0, i, 0]-k)
    upper_bound = min(100, B[0, i, 0]+k+1)
    A_mean = A[lower_bound:upper_bound, i:i+1, :].mean(0, keepdim=True)
    A_list.append(A_mean)

A = torch.cat(A_list , dim=1)

a[:, b.reshape(-1), :] will give you ith element wanted. — Amir
k is the width of the slice that I want to take from A in the first dimension. In this case, 3. I will update the code — Jakub Chłędowski

Amir Amir · Accepted Answer · 2021-01-19T20:23:51

Something similar to this can work if k = 1. (requires torch>1.7)

a = torch.rand((10, 20, 4))
b = torch.randint(10, (20, 1))
b2 = torch.cat((b-1, b, b+1), dim=1)
b2 = torch.minimum(9*torch.ones_like(b2), b2)
b2 = torch.maximum(0*torch.ones_like(b2), b2)
a[:, b2, :]

and then reshape to get the right size

How to quickly cut slices from cuda tensor wrt to another tensors values

1 Answers