如何在PyTorch中将张量沿特定的新维度重复？

Question

如果我有一个形状为[M, N]的张量A，我想将该张量重复K次，使结果B的形状为[M，K，N]，并且每个切片B [:，k，:] 应具有与A相同的数据。哪种方法是最佳实践，而不需要使用循环语句？K 可能在其他维度上。 torch.repeat_interleave()和tensor.repeat()似乎无法工作。或者我使用它们的方式有误。

- River

---

或许。如果您能提供您正在使用的实际代码，会更有帮助。 - DeadChex

Answer 1

tensor.repeat可以满足您的需求，但您需要先插入一个单位维度。为此，我们可以使用tensor.unsqueeze或tensor.reshape。由于unsqueeze是专门定义用于插入一个单位维度的，因此我们将使用它。

B = A.unsqueeze(1).repeat(1, K, 1)

代码说明 A.unsqueeze（1） 将 A 从 [M，N] 变为 [M，1，N]，.repeat（1，K，1） 沿第二维度重复张量 K 次。

- jodag

Answer 2

Einops 提供了 repeat 函数

import einops
einops.repeat(x, 'm n -> m k n', k=K)

repeat 可以任意添加轴，并可以同时对现有轴进行重新排序。

- Alleo

Answer 3

tensor.expand 可能比 tensor.repeat 更好，因为根据这里的说法：“扩展张量不会分配新内存，而只是在现有张量上创建一个新视图，在该视图中将大小为一的维度通过设置步幅为零扩展到更大的大小。”

但是，请注意：“扩展张量的多个元素可能引用单个内存位置。因此，原地操作（特别是向量化的操作）可能导致不正确的行为。如果需要写入张量，请先克隆它们。”

M = N = K = 3
A = torch.arange(0, M * N).reshape((M, N))
B = A.unsqueeze(1).expand(M, K, N)
B

'''
tensor([[[0, 1, 2],
         [0, 1, 2],
         [0, 1, 2]],

        [[3, 4, 5],
         [3, 4, 5],
         [3, 4, 5]],

        [[6, 7, 8],
         [6, 7, 8],
         [6, 7, 8]]])
'''

- fatpanda2049

Answer 4

明确重复值会快速产生巨大的内存成本。在大多数情况下，您可以通过使用广播来隐式保留值。因此，可以使用A[:, None, :]并获得新形状A.shape==(M, 1, N)。

有一种情况我同意重复值的使用，那就是在以下步骤中进行原地操作。 由于numpy和torch在实现上有所不同，我喜欢使用通用的(A * torch.ones(K, 1, 1)))，然后再进行转置。

- 0-_-0

Answer 5

除了@Alleo提供的答案外，您可以使用以下Einops函数。

einops.repeat(example_tensor, 'b h w -> (repeat b) h w', repeat=b)

其中b是您希望张量重复的次数，h、w是张量的附加维度。

示例：

example_tensor.shape -> torch.Size([1, 40, 50]) 
repeated_tensor = einops.repeat(example_tensor, 'b h w -> (repeat b) h w', repeat=8)
repeated_tensor.shape -> torch.Size([8, 40, 50]) 

这里有更多的例子 - https://einops.rocks/api/repeat/

- Ashish Alex