Question

如果我的张量A的形状为[M, N]，我想将张量重复K次，以便结果B的形状为[M, K, N] 每个切片B[:, k, :]的数据应与A相同。这是没有for循环的最佳实践。 K可能在其他维度上。

torch.repeat_interleave()和tensor.repeat()似乎不起作用。否则我使用的方式有误。

Answer 1

tensor.repeat应该适合您的需求，但您需要先插入一个单一尺寸。为此，我们可以使用tensor.reshape或tensor.unsqueeze。由于unsqueeze是专门为插入单一尺寸而定义的，因此我们将使用它。

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

代码说明 A.unsqueeze(1)将A从[M, N]变成[M, 1, N]，.repeat(1, K, 1)重复张量K次维度。

Answer 2

Einops提供重复功能

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

repeat可以按任意顺序添加任意数量的轴，并同时重新排列现有轴。

Answer 3

添加到@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/

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

3 个答案: