函数功能
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