PyTorch中从1D张量批量收集值的高效实现方案（规避expand内存开销）

你观察得很到位！虽然expand本身是视图操作不会立刻复制内存，但当处理超大尺寸的list_values时，后续操作确实可能触发隐式内存复制，而且完全没必要多这一步操作——我们有更直接、内存友好的方式实现需求。

最简洁的方案：直接张量索引 #

PyTorch原生支持用多维索引张量直接索引一维张量，会自动对齐批量维度，代码超简洁：

import torch

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

运行后result就是你要的tensor([[2, 3], [3, 4]])，全程不需要对list_values做任何维度扩展，完全基于原张量的视图索引，内存开销可以忽略不计。

如果你偏好gather：无需手动expand #

要是你习惯用gather实现，也不用先expand，只需要给一维张量增加一个维度，利用PyTorch的广播机制自动匹配批量维度：

import torch

list_values = torch.tensor([1, 2, 3, 4, 5, 6])
indices = torch.tensor([[1, 2], [2, 3]])
# 仅增加维度，不用expand，gather会自动广播匹配
result = torch.gather(list_values.unsqueeze(0), 1, indices)

这里list_values.unsqueeze(0)把原张量变成(1,6)的形状，gather在维度1上取索引时，会自动把这个形状广播到和indices的(2,2)兼容，避免了手动expand的冗余步骤，同样不会额外占用内存。

补充：为什么原方法的expand没必要？ #

PyTorch的expand只是创建一个“虚拟”的扩展视图，不会立刻复制数据，但如果后续操作需要连续内存（比如某些CUDA操作），就会触发隐式的内存复制，这对超大张量来说是不必要的开销。而上面两种方法都是直接基于原张量的索引操作，完全绕开了这个潜在问题。

内容的提问来源于stack exchange，提问作者XPenguen