本文介绍如何在不显式循环的前提下，将 pytorch 张量中每个元素按预定义映射关系（如字典）快速替换为对应值，核心方法是将映射字典转为索引张量后直接索引。
本文介绍如何在不显式循环的前提下，将 pytorch 张量中每个元素按预定义映射关系（如字典）快速替换为对应值，核心方法是将映射字典转为索引张量后直接索引。
在深度学习与数据预处理中，常需根据离散标签（如类别索引 0、1、2…）批量映射为对应权重、嵌入向量或概率值。若使用 Python 字典配合 for 循环逐元素查表，不仅低效，还无法利用 GPU 加速。幸运的是，PyTorch 提供了张量索引（tensor indexing）这一向量化操作，可优雅、高效地完成此类映射。
✅ 推荐方案：字典 → 映射张量 → 索引访问
关键思路是：将字典 weights = {0: 0.1, 1: 0.9} 转换为一个按键顺序排列的张量，其中索引即原字典的 key，值即对应的 value。例如，key 最小值为 0、最大值为 1，则构造长度为 2 的张量 torch.tensor([0.1, 0.9])；此时对输入张量 t 中每个元素 t[i]，直接以它为下标访问该张量即可：
import torch
t = torch.tensor([1, 0, 0, 1], dtype=torch.long) # 注意：索引需为 long 类型
weights_dict = {0: 0.1, 1: 0.9}
# 构建映射张量：按 key 从小到大填充，确保 key 是连续非负整数
max_key = max(weights_dict.keys())
weights_tensor = torch.zeros(max_key + 1)
for k, v in weights_dict.items():
weights_tensor[k] = v
# 向量化映射（支持 CPU/GPU）
new_t = weights_tensor[t]
print(new_t)
# tensor([0.9000, 0.1000, 0.1000, 0.9000])
✅ 优势：完全向量化、零 Python 循环、天然支持 GPU（只要
t和weights_tensor在同一设备上）；时间复杂度 O(n)，远优于 O(n) 的纯 Python 查表。
⚠️ 注意事项
-
索引类型必须为
torch.long：t若为浮点型（如torch.float32），需先.long()转换，否则报错IndexError: tensors used as indices must be long or byte tensors。 -
key 需为非负整数且尽量连续：若字典 key 为
{1: 0.1, 3: 0.9}，则需填充中间空位（如weights_tensor[0]=0, [1]=0.1, [2]=0, [3]=0.9），或先做偏移映射（如t_shifted = t - min_key）。 -
GPU 兼容性：只要
t和weights_tensor均位于 CUDA 设备（如t.cuda(), weights_tensor.cuda()），索引操作自动在 GPU 上执行，无需额外修改。
? 扩展：泛化函数封装
为提升复用性，可封装为通用映射函数：
def map_tensor_by_dict(t: torch.Tensor, mapping: dict, device=None) -> torch.Tensor:
keys = list(mapping.keys())
if not all(isinstance(k, (int, np.integer)) and k >= 0 for k in keys):
raise ValueError("Keys must be non-negative integers")
max_k = max(keys)
mapping_tensor = torch.zeros(max_k + 1, dtype=torch.float32)
for k, v in mapping.items():
mapping_tensor[k] = float(v)
if device:
mapping_tensor = mapping_tensor.to(device)
t = t.to(device)
return mapping_tensor[t.long()]
# 使用示例
t = torch.tensor([1, 0, 0, 1])
new_t = map_tensor_by_dict(t, {0: 0.1, 1: 0.9})
该方法简洁、高效、可扩展，是 PyTorch 中替代字典查表的首选实践。
