给定几个块矩阵，如何得到整体的大矩阵

Question

简单来说，我们有一个四维的张量 y，它的形状是 ( B // s2, D2 // s1, s1, s2)。这里的 y[i,j,...] 表示一个形状为 (s1,s2) 的矩阵。这些小矩阵是用来构建一个大的整体矩阵，整体矩阵的形状是 (D2, B)。总共有 (B//s2) * (D2 //s1) 这样的小矩阵。我们假设所有的数字都是整数。我知道怎么用循环来实现这个操作：

# y shape ( B // s2, D2 // s1, s1, s2)
result = torch.zeros(D2, B)
for i in range(D2 // s1):
    for j in range(B // s2):
         result[i * s1: (i + 1) * s1, j * s2: (j + 1) * s2] = y[j,i, ...]

我知道这个赋值操作可以并行进行。我们能不能用 pytorch 的内置函数来去掉这两个循环呢？

Answer 1

这被称为折叠操作，nn.Fold 和 F.fold 是为这个目的而设计的。如果你查看文档，会看到这样一段说明：

这个操作将一组滑动的局部块组合成一个大的张量。
想象一下一个包含滑动局部块的输入张量，比如图像的补丁，形状是 (N,C×∏(kernel_size),L)，其中：

• N 是批次维度，
• C×∏(kernel_size) 是一个块内的值的数量（一个块有 ∏(kernel_size) 个空间位置，每个位置包含一个 C 通道的向量），
• L 是块的总数量。

这和 Unfold 的输出形状完全一样。这个操作通过对重叠的值进行求和，将这些局部块组合成一个形状为 (N,C,output_size[0],output_size[1],…) 的大输出张量。

在你的情况下，输入张量的形状是 (B//s2,D2//s1,s1,s2)。要达到上述规格，你有 L = B//s2 * D2//s1，而卷积核的大小是 ∏(kernel_size) = s1 * s2。由于这个函数期望 L 在最后一个位置，所以在扁平化之前你需要进行一些维度调整：

y_ = y.permute(2,3,1,0).reshape(s1*s2,-1)

现在 y_ 的形状是 (s1*s2, B//s2*D2//s1)。最后你可以应用折叠操作：

F.fold(y_, output_size=(D2,B), kernel_size=(s1,s2), stride=(s1,s2))