Numpy/Torch中如何对批量向量按批量索引进行重索引？

当处理批量向量时（二维数组v的每一行v[i,:]代表一个独立向量），需要用对应行的索引矩阵IX[i,:]对每行向量做重索引，Python循环的方式效率极低，这里用Numpy和PyTorch的原生高级索引实现高效、可读的批量重索引。

Numpy 实现 #

你找到的方案就是Numpy官方推荐的规范实现，核心是利用高级索引，通过构造行索引矩阵和列索引矩阵配对，实现逐行的重索引：

import numpy as np

# 输入批量向量和对应行的索引矩阵
v = np.array([[10, 20, 30], [40, 50, 60], [70, 80, 90]])
IX = np.array([[2, 1, 0], [0, 2, 1], [1, 0, 2]])

# 构造行索引：利用np.newaxis扩展维度，让行索引可以和IX广播匹配
row_indices = np.arange(v.shape[0])[:, np.newaxis]

# 高级索引：row_indices对应每行的位置，IX对应该行内的列索引
new_v = v[row_indices, IX]

print(new_v)
# 输出：
# [[30 20 10]
#  [40 60 50]
#  [80 70 90]]

原理说明 #

np.arange(v.shape[0])[:, np.newaxis]生成形状为(N,1)的行索引数组，和形状为(N,M)的IX矩阵广播后，两者会配对成(N,M)的索引对(row_idx, col_idx)，Numpy会根据这些索引对直接从v中提取对应元素，完全避免Python循环，效率和底层C实现一致。

对比你之前尝试的v.ravel()[ (IX + range(v.shape[0]) ).ravel() ].reshape(N,-1)，这种方法需要手动计算扁平化后的偏移量，不仅可读性差，还容易在维度变化时出错，而高级索引的方式更直观，代码维护性更强。

PyTorch 实现 #

PyTorch的索引逻辑和Numpy完全一致，直接复用相同的思路即可：

import torch

# 输入批量张量和对应行的索引张量
v = torch.tensor([[10, 20, 30], [40, 50, 60], [70, 80, 90]])
IX = torch.tensor([[2, 1, 0], [0, 2, 1], [1, 0, 2]])

# 构造行索引，用unsqueeze扩展维度实现广播
row_indices = torch.arange(v.shape[0]).unsqueeze(1)

# 高级索引实现批量重索引
new_v = v[row_indices, IX]

print(new_v)
# 输出：
# tensor([[30, 20, 10],
#         [40, 60, 50],
#         [80, 70, 90]])

注意事项 #

• 确保row_indices和IX的维度匹配，通过np.newaxis（Numpy）或unsqueeze（PyTorch）让行索引的维度兼容，实现广播。
• 高级索引返回的数组/张量形状和IX的形状一致，无需手动reshape，更简洁。

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