pytorch 从一维张量收集

Question

I have a long list of values stored in a 1d tensor from which I like to gather values. Given a 2d tensor where the second dimension denotes a set of indices, I would like to compute a new 2d tensor where the second dimension are all the values gathered given by the indices. Basically, I have a batch of indices and I use those indices to obtain a batch of values from the 1d-tensor.

I am new to pytorch but I managed to compute exactly that using pytorch's gather function

list_values = torch.tensor([1, 2, 3, 4, 5, 6])
list_values = list_values.unsqueeze(0)
list_values = list_values.expand((2, 6))
indices = torch.tensor([[1, 2], [2, 3]])
result = torch.gather(list_values, 1, indices)

This works perfectly fince and gives the correct result. However, If I am not mistaken, the expand operation is quite expensive in terms of memory as the number of elements in "list_value" grows.

Is there a better solution?

Answer 1

Nope, take a look at the documentation :

Expanding a tensor does not allocate new memory, but only creates a new view on the existing tensor where a dimension of size one is expanded to a larger size by setting the stride to 0. Any dimension of size 1 can be expanded to an arbitrary value without allocating new memory.

You solution is completely fine