PyTorch中如何高效将张量内整数替换为对应子张量？

直接利用PyTorch的批量索引+张量重塑就能以最小计算量完成需求，完全避开循环或低效操作，具体步骤如下：

实现代码 #

import torch

# 输入张量
a = torch.tensor([[1,2,3],[2,1,3]])
# 建立映射表：索引0对应原元素1的替换值，索引1对应原元素2，索引2对应原元素3
map_tensor = torch.tensor([[1,2,3], [4,5,6], [7,8,9]])

# 1. 将原张量元素转为0基索引（匹配映射表的索引规则）
idx = a - 1
# 2. 批量提取映射表中的对应行，得到形状为(2, 3, 3)的中间张量
selected = map_tensor[idx]
# 3. 重塑为目标形状(2, 9)
result = selected.flatten(start_dim=1)

print(result)
# 输出：
# tensor([[1, 2, 3, 4, 5, 6, 7, 8, 9],
#         [4, 5, 6, 1, 2, 3, 7, 8, 9]])

方法优势 #

• 计算量极小：所有操作都是PyTorch底层优化的张量批量操作，无Python层面循环，GPU上可完全并行处理，适合大规模张量场景。
• 避免低效方案：无需使用numpy.vectorize（本质是循环，速度慢），也不需要复杂的vmap操作——你的vmap尝试失败是因为没必要逐个元素调用.item()，批量索引本身就是更高效的实现方式。

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