将每个输入元素映射到 pytorch 中的数组

Question

我有一个 (A,B) 张量，我正在寻找一种高效的方法来将该张量中的每个值映射到一个数组，以创建一个大小为 (A,B,N) 的新张量。这是一段功能代码，展示了我正在尝试做的事情。

A, B, N = 3, 4, 5
my_old_tensor = torch.ones((A,B), dtype=torch.float32)
my_new_tensor = torch.zeros((A, B, N), dtype=torch.float32)
for val in range(N):
    my_new_tensor[:,:,val] = (val - my_old_tensor)/2

我的代码目前很慢，我认为是 for 循环的问题。是否有一种更符合 pytorch 性能的方法来消除 for 循环？我试过这样的东西

x = torch.arange(0, N, 1, dtype=torch.float32)
my_new_tensor = (x - my_old_tensor)/2

但这给出了“运行时错误：张量 a (5) 的大小必须与非单维 1 处的张量 b (4) 的大小匹配”

如有任何帮助，我们将不胜感激！

Answer 1

使用unsqueeze到broadcastmy_old_tensor:

my_new_tensor = (torch.arange(N, dtype=torch.float32) - my_old_tensor.unsqueeze(-1))/2

将每个输入元素映射到 pytorch 中的数组

Map each input element to array in pytorch

python

numpy

pytorch