本文介绍一种无需显式循环、兼容 cpu/gpu 的优雅方法：将映射字典转为索引张量，利用 pytorch 的高级索引机制直接完成批量映射。
本文介绍一种无需显式循环、兼容 cpu/gpu 的优雅方法：将映射字典转为索引张量，利用 pytorch 的高级索引机制直接完成批量映射。
在深度学习实践中，常需将离散标签张量（如类别索引）映射为对应的权重、嵌入向量或概率值。例如，给定一个仅含 0 和 1 的张量 t = torch.tensor([1, 0, 0, 1]) 和映射字典 weights = {0: 0.1, 1: 0.9}，目标是生成 new_t = tensor([0.9, 0.1, 0.1, 0.9])。传统 for 循环或 .apply_()（已弃用且不支持 GPU）效率低且不优雅。
推荐方案：字典 → 张量 → 高级索引
核心思想是将字典键值对转换为以键为索引、值为内容的一维张量，再利用 PyTorch 的整数张量索引（Integer Tensor Indexing）实现 O(1) 每元素映射：
import torch
t = torch.tensor([1, 0, 0, 1], dtype=torch.long) # 确保索引类型为 long
weights_dict = {0: 0.1, 1: 0.9}
# 构建映射张量：索引 i 对应 weights_dict[i]，要求键为连续非负整数
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/autograd
new_t = weights_tensor[t]
print(new_t) # tensor([0.9000, 0.1000, 0.1000, 0.9000])
✅ 优势说明：
- 零循环：全程向量化操作，性能最优；
-
设备无关：
weights_tensor和t可同置于 CUDA，索引自动在 GPU 上执行； -
可微分：若
weights_tensor是 requires_grad=True 的参数，该操作支持反向传播； -
通用性强：适用于任意离散整数键（如类别 ID、token ID），只需确保键范围紧凑（否则可用
torch.nn.Embedding替代）。
⚠️ 注意事项：
- 输入张量
t必须为torch.long或torch.int64类型，否则索引报错； - 字典键应为从
0开始的连续非负整数；若键稀疏（如{3: 0.2, 7: 0.8}），建议改用torch.nn.Embedding(num_embeddings=8, embedding_dim=1)并初始化对应行； - 若需映射字符串或负数键，应先通过
torch.unique(..., return_inverse=True)构建整数编码映射表。
该方法体现了 PyTorch “以张量为中心”的设计哲学——将逻辑关系转化为张量结构与运算，是构建高效、可扩展数据预处理流程的关键实践。
