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

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可以按任意顺序添加任意数量的轴，并同时重新排列现有轴。