Pytorch: Most computationally and memory efficient way to make a series of concatenations from extracting tensor rows?

user_3259896 0
Viewed 76

Say that this is my sample tensor

sample = torch.tensor(
    [[2, 7, 3, 1, 1],
        [9, 5, 8, 2, 5],
        [0, 4, 0, 1, 4],
        [5, 4, 9, 0, 0]]
)

I want to have a new tensor, which will consist of concatenations of 2 rows from the sample tensor.

So I have a tensor which contains pairs of the row numbers that I want concatenated into a single row for the new tensor

cat_indices = torch.tensor([[0, 1], [1, 2], [0, 2], [2, 3]])

The current method I am using is this

torch.cat((sample[cat_indices[:,0]], sample[cat_indices[:,1]]), dim=1)

Which gives the desired result

tensor([[2, 7, 3, 1, 1, 9, 5, 8, 2, 5],
        [9, 5, 8, 2, 5, 0, 4, 0, 1, 4],
        [2, 7, 3, 1, 1, 0, 4, 0, 1, 4],
        [0, 4, 0, 1, 4, 5, 4, 9, 0, 0]])

Is this the most memory and computationally efficient method of doing this? I am not sure because I am making two calls to cat_indices, and then I am doing a concatenation operation.

I feel that there should be a way to do this via some sort of view. Perhaps advanced indexing. I've tried things like sample[cat_indices[:,0], cat_indices[:,1]] or sample[cat_indices[0], cat_indices[1]] but I can't make the view come out right.

pytorch
1 Answers

What you have should be pretty fast. An alternative is

sample[cat_indices].reshape(cat_indices.shape[0],-1)

You would have to benchmark the performance on your machine though to see which is better.

user_27900470
Related