如何在不遍历张量的情况下将torch.Tensor的每个元素映射为字典中对应值（支持GPU张量）

嘿，这个需求我之前做项目的时候也碰到过！完全不用写循环遍历，PyTorch里有好几种高效的方法，而且完美支持GPU张量，给你分享几个最实用的：

方法1：张量索引法（最简洁高效） #

这是我日常最常用的方案，不管CPU还是GPU张量都能直接用，核心思路是把字典的权重转换成有序张量，再用原张量做索引直接取值：

import torch

# 示例张量（如果是GPU张量，直接改成 t = torch.Tensor([1,0,0,1]).long().cuda() 即可）
t = torch.Tensor([1, 0, 0, 1]).long()  # 注意要转成整数类型，索引操作需要整数输入
weights = {0: 0.1, 1: 0.9}

# 按字典键的顺序创建权重张量，确保位置和键一一对应
weight_tensor = torch.tensor([weights[0], weights[1]])
# 直接用原张量索引，一步到位得到映射结果
new_t = weight_tensor[t]

这个方法之所以好用，是因为它是完全向量化的操作，PyTorch会自动批量处理所有元素，不管张量多大速度都拉满，GPU上也能完美运行——毕竟这是PyTorch原生的张量操作，会自动利用CUDA加速，完全不用操心遍历的问题。

方法2：torch.gather 实现（原理类似，适合扩展场景） #

如果之后你遇到更复杂的索引场景，torch.gather 也是个不错的选择，同样支持GPU：

import torch

t = torch.Tensor([1, 0, 0, 1]).long()
weights = {0: 0.1, 1: 0.9}

weight_tensor = torch.tensor([weights[0], weights[1]]).unsqueeze(0)
# 用gather按索引取值，最后squeeze去掉多余维度
new_t = torch.gather(weight_tensor, 0, t).squeeze()

不过说实话，你的场景里方法1已经足够简洁了，这个可以作为扩展知识备用。

小提醒 #

• 一定要确保原张量是整数类型（比如long或int），PyTorch的索引操作不接受float类型的张量，所以如果你的t是float类型，记得用t.long()转一下。
• 如果你的张量在GPU上，只需要把weight_tensor也移到GPU上就行（比如weight_tensor = weight_tensor.cuda()），或者直接在创建时就指定设备，操作逻辑完全不变。

备注：内容来源于stack exchange，提问作者onthebox