Pytorch, retrieving values from a 3D tensor using several indices. Most computationally efficient solution

Question

This is another question about retrieving values from a 3D tensor, using a list of indices.

In this case, I have a 3d tensor, for example

b = [[[4, 20], [1, -1]], [[1, 2], [8, -1]], [[92, 4], [23, -1]]]
tensor_b = torch.tensor(b)
tensor_b

tensor([[[ 4, 20],
         [ 1, -1]],

        [[ 1,  2],
         [ 8, -1]],

        [[92,  4],
         [23, -1]]])

In this case, I have a list of 3D indices. So

indices = [
           [[1, 0, 1], [2, 0, 1]],
           [[1, 1, 1], [0, 0, 0]],
           [[2, 1, 0], [0, 1, 0]]
]

Each triple is an index for tensor-b. The desired result is

[[2, 4], [-1, 4], [23, 1]]

Potential Approach

Like in the last question, the first solution that comes to mind is a nested for loop, but there is probably a more computationally efficient solution using pytorch function.

And like in the last question, perhaps reshape would be needed to get the desired shape for the last solution.

So a desired solution could be [2, 4, -1, 4, 23, 1], which can come from a flattened list of indices [ [1, 0, 1], [2, 0, 1], [1, 1, 1], [0, 0, 0], [2, 1, 0], [0, 1, 0] ]

But I am not aware of any pytorch functions so far which allow for a list of 3D indices. I have been looking at gather and index_select.

user_2790047 · Answer

You can use advanced indexing specifically integer array indexing

tensor_b = torch.tensor([[[4, 20], [1, -1]], [[1, 2], [8, -1]], [[92, 4], [23, -1]]])
indices = torch.tensor([
           [[1, 0, 1], [2, 0, 1]],
           [[1, 1, 1], [0, 0, 0]],
           [[2, 1, 0], [0, 1, 0]]
])

result = tensor_b[indices[:, :, 0], indices[:, :, 1], indices[:, :, 2]]

results in

tensor([[ 2,  4],
        [-1,  4],
        [23,  1]])

Pytorch, retrieving values from a 3D tensor using several indices. Most computationally efficient solution

1 Answers