如何向 Torch 张量添加一个维度，并在新维度中重复前两个维度的值

Question

我有一个形状为 [a,b] 的张量。如何将其重塑为 [a,b,c]，以便沿新维度的所有值重复形状 [a,b] 的原始张量的 i,j 条目的值？

例如：a = 2，b= 3

tensor([[1,2,3],  
       [4,5,6]]) 

现在用 c = 4 重塑以获得形状 [2,3,4]

结果：

tensor([[[1,1,1,1],
         [2,2,2,2],
         [3,3,3,3]],

         [4,4,4,4],
         [5,5,5,5],
         [6,6,6,6]]])

Answer 1

我想最快的方法是使用

einops

库。它通过使变换清晰可见来提供大量张量操作。这里：

from einops import repeat
X = torch.tensor([[1, 2, 3], [4, 5, 6]])
Y = repeat(X, "a b -> a b c", c=4)

输出：

tensor([[[1,1,1,1],
         [2,2,2,2],
         [3,3,3,3]],

         [4,4,4,4],
         [5,5,5,5],
         [6,6,6,6]]])