python - 以窗口方式重新排列 PyTorch 张量

Question

所以我有一个张量

A = torch.tensor([[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]])

我希望在滑动窗口中重新排列它，也就是说，

f(A) = torch.tensor([[1,2,5,6],[3,4,7,8],[9,10,13,14],[11,12,15,16]])

通过将每个张量A视为 2×2 小窗口，然后将“窗口”排列在 2×2 矩阵中。此操作类似于ArrayFlattenMathematica，但我找不到在 PyTorch 中执行此操作的方法。欢迎任何帮助。

标签： pythonpytorch

Answer 1

我想不出一个巧妙的方法来做到这一点，但你可以通过一些选择切片和连接来实现它：

A_ = A.reshape(4,2,2)
torch.cat([torch.cat([*A_[:2]],1), torch.cat([*A_[2:]],1)],0)

---

替代方法：

A.unfold(1,2,2).unfold(0,2,2).reshape(*A.shape).index_select(1, torch.LongTensor([0,2,1,3]))

torch.cat([*A_],1).reshape(*A.shape).index_select(0, torch.LongTensor([0,2,1,3]))