torch.gather与torch.index_select的区别及场景选型咨询

针对你从形状为(max_len, batch_size, embedding_dim)的位置嵌入中，在max_len维度选取特定索引得到(new_indices, batch_size, embedding_dim)结果的场景，torch.index_select是更合适的选择，原因如下：

1. 语法简洁，无需额外索引变形 #

index_select专门设计用于在单个指定维度上选取一组统一的索引，完全匹配你的需求：

• 你只需要传入目标维度（这里是dim=0，对应max_len）和一维的索引张量（形状为(new_indices,)），不需要对索引做任何形状扩展或广播操作。
• 而torch.gather是更通用的索引工具，它要求索引张量的形状必须和输入张量的形状在非目标维度上完全匹配。要实现相同效果，你需要先把索引扩展成(new_indices, batch_size, embedding_dim)的形状，额外增加代码复杂度。

2. 语义明确，可读性更强 #

index_select的命名和功能高度一致，其他开发者看到代码就能立刻明白你是在某个维度上批量选取固定的索引集合；而gather的语义更宽泛，常用于“每个元素位置选取不同索引”的场景，用在这里会让代码意图变得模糊。

3. 性能更高效 #

index_select是针对单维度批量索引场景做了专门优化的算子，在这种所有非目标维度共享同一组索引的情况下，它的执行效率比gather更高——后者需要处理更通用的索引逻辑，会带来不必要的计算开销。

代码对比 #

使用torch.index_select（推荐）

import torch

# 模拟输入的位置嵌入
max_len = 10
batch_size = 3
embedding_dim = 5
pos_emb = torch.randn(max_len, batch_size, embedding_dim)

# 要选取的索引（比如第2、5、7个位置）
target_indices = torch.tensor([2, 5, 7])

# 直接在dim=0上选取索引
result = torch.index_select(pos_emb, dim=0, index=target_indices)
print(result.shape)  # 输出: torch.Size([3, 3, 5]) → 符合预期的(new_indices, batch_size, embedding_dim)

使用torch.gather（冗余且麻烦）

# 需要先将索引扩展为和输入同形状的张量
expanded_indices = target_indices.unsqueeze(1).unsqueeze(2).expand(-1, batch_size, embedding_dim)
result = torch.gather(pos_emb, dim=0, index=expanded_indices)
print(result.shape)  # 同样得到torch.Size([3, 3, 5])，但多了不必要的形状扩展步骤

什么时候用torch.gather？ #

如果你的需求是不同的batch元素或不同的embedding维度需要选取不同的max_len索引（比如索引形状为(batch_size, embedding_dim)），此时gather才是更合适的工具——它支持为每个位置指定独立的索引。

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