如何在 PyTorch 中将值从 3D 张量映射到 1D 张量？

Question

我遇到了 Pytorch 问题，需要一些帮助：

我有两个张量：

• 具有从 0 到 n 的整数值的 3D 张量（形状：i、j、j）
• 一维张量（形状：n）

我需要创建一个与第一个张量 (i, j, j) 形状相同的新张量，但其中每个值都被第二个张量中的相应值替换。就像使用第一个张量中的值作为第二个张量的索引一样。

这是一个简单的例子：

import torch

# My tensors
big_tensor = torch.randint(0, 256, (10, 25, 25))
small_tensor = torch.rand(256)

# What I'm trying to do
result = magic_function(big_tensor, small_tensor)

# How it should work
print(big_tensor[0, 0, 0])    # Let's say this outputs 42
print(small_tensor[42])       # This might output 0.7853
print(result[0, 0, 0])        # This should also be 0.7853

我正在寻找一种高性能的方法来做到这一点，最好不使用循环，因为两个张量都可能很大。我可以使用高效的 Pytorch 操作或方法吗？

任何帮助都会很棒 - 提前致谢！

Answer 1

这应该有效：

small_tensor[big_tensor]

注意

big_tensor

的类型必须是long/int。