重塑 PyTorch 中张量的维度

Question

有一个形状为[b, nt*nh*nw, dim]的张量。 nt、nh 和 nw 的值在手。 如何将此张量重塑为 [b, dim, nt, nh, nw] 的形式？ 例如，如何将 [2, 3x2x4, 512] 重塑为 [2,512,3,2,4]？

Answer 1

这完全取决于您在 memory 中的数据布局。

但是，假设nt 、 nh和nw在您的基础数据张量中的顺序正确，那么您可以通过排列和重塑您的张量来做到这一点。

首先使用torch.transpose或torch.permute交换尺寸以将dim作为第二轴。 然后使用torch.view或torch.reshape将张量重塑为所需的形状：

>>> x.transpose(1,2).view(b, dim, nt, nh, nw)