PyTorch中索引子集相关操作的向量化实现及for循环移除方法

只要list_of_subset_indices中每个子列表的长度一致，就可以完全通过框架自带的高级索引实现向量化，移除Python层的for循环，运算效率会大幅提升。如果子列表长度不一致，也有对应的优化方案。

场景1：所有子集索引长度相同（绝大多数场景） #

原理是利用张量的高级索引规则，一次性提取所有位置的子集，再批量求乘积：

PyTorch 实现 #

import torch

# 把嵌套索引列表转为整形张量，形状为 [N, k]，N为result第一维长度，k为每个子集的固定长度
indices = torch.tensor(list_of_subset_indices, dtype=torch.long)
# 高级索引批量提取对应元素，输出形状为 [N, k, feat_dim]
selected_tensor = other_tensor[indices, :]
# 沿子集维度求乘积，输出形状为 [N, feat_dim]，就是最终的result
result = selected_tensor.prod(dim=1)

NumPy 实现 #

逻辑和PyTorch完全一致，高级索引规则通用：

import numpy as np

indices = np.array(list_of_subset_indices, dtype=np.int64)
selected_tensor = other_tensor[indices, :]
result = selected_tensor.prod(axis=1)

如果你要严格对齐原代码中range(1, result.shape[0])的逻辑（不对i=0的位置赋值），仅需将最后一步改为result[1:] = selected_tensor.prod(dim=1)即可，i=0的位置会保留原值。

场景2：不同位置的子集索引长度不一致 #

如果每个子列表的长度不相等，无法直接构造规整的索引张量，可以选择两种优化方案：

• 方案1：给短的子集补索引，补的位置取other_tensor中对应值为1的元素（1不影响乘积结果），补全到和最长子集长度一致后，仍可使用上述向量化逻辑
• 方案2：使用torch.vmap（PyTorch场景）或者numba.jit（NumPy场景）包装索引+求积的逻辑，框架会自动完成并行化处理，性能远高于手写Python层for循环

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