torch.nn.functional.unfold与torch.nn.functional.fold函数

函数功能

torch.nn.functional.unfold是PyTorch中的一个函数，用于将输入张量展开为一个二维张量。
展开的方式是将输入张量的每个滑动窗口视为一个列向量，并将这些列向量按顺序排列成一个二维张量。
torch.nn.functional.fold常与torch.nn.functional.unfold搭配使用，是unfold的逆过程，参数基本一致。

参数含义

torch.nn.functional.unfold(input, kernel_size, dilation=1, padding=0, stride=1）
torch.nn.functional.fold(input, output_size, kernel_size, dilation=1, padding=0, tride=1)

input表示输入张量，即unfold函数操作的对象，注意必须是4D的张量；
kernel_size表示滑动窗口的大小；
dilation表示邻域内元素的膨胀步幅，默认为1；
padding表示在输入的两侧进行填充的长度，通常情况下是采用全0填补；
stride表示滑动窗口的步长。
output_size表示输出张量的尺寸（H，W）而非输出张量的shape→（B, C, H, W）

示例

为了便于理解和记忆该函数，这里使用unfold及fold函数手动实现平均池化，代码如下：

#平均池化
def mean_pool(input, ph = 2, pw = 2, padding = 0, stride = 2):
    (b, c, h, w) = input.shape
    
    nh = 1 + int((h + 2 * padding - ph) / stride) 
    nw = 1 + int((w + 2 * padding - pw) / stride)
    
    input_unfold = torch.nn.functional.unfold(input, (ph, pw), padding = padding, stride = stride)
    input_unfold = input_unfold.transpose(1, 2).view(b, nh*nw, c, -1)
    out_mean_pool = torch.mean(input_unfold, dim=3, keepdim=False).transpose(1,2)
    output = torch.nn.functional.fold(out_mean_pool, (nh, nw), (1, 1))
    
    assert (output.size() == (b, c, nh, nw))
    return output