PyTorch实现张量索引映射：无需循环的高效方法

在PyTorch中，你可以直接利用张量索引的特性完成这个映射，完全不需要循环，这是最简洁高效的线性实现方式：

直接通过b[a]就能得到目标张量，原理是将张量a中的每个元素作为索引，自动取张量b对应位置的值，完成逐元素的映射。

完整代码示例 #

import torch

a = torch.tensor([0, 2, 5, 1, 0, 0, 4, 3, 0, 0, 0, 2, 2, 1])
b = torch.tensor([1.5, 2.7, 1.8, 7.0, 3.9, 10.0])

new_tensor = b[a]
print(new_tensor)

运行后输出结果：

tensor([ 1.5000,  1.8000, 10.0000,  2.7000,  1.5000,  1.5000,  3.9000,  7.0000,
         1.5000,  1.5000,  1.5000,  1.8000,  1.8000,  2.7000])

这个方法基于PyTorch底层优化的张量操作，性能远优于循环实现。

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