做图像patch块提取的时候，很多人的第一反应是写双层循环，挨个切片再torch.stack。我第一次做滑动窗口裁剪时也是这么干的，直到后来在项目里发现torch.nn.Unfold可以一步把整张图的patch全部取出来，速度快、代码短，而且能直接接在神经网络里做可微操作。简单理解，Unfold就是PyTorch对im2col的标准封装，输入一张N,C,H,W的图像，输出变成N, C*kH*kW, L这种“把每个滑动窗口摊平成列”的张量，全程不需要自己管理索引。这篇文章会从原理、参数公式、实操代码、踩坑记录、扩展应用几个角度把Unfold讲透，适合正在写自定义卷积、滑动窗口网络、图像复原和patch分类相关代码的人。

1. Unfold的工作逻辑：从patch循环到im2col

1.1 写循环切patch的痛点

先看一段最直觉的patch提取代码，假设输入是x，形状是(N, C, H, W)，想提取大小为k、步长为s的所有块：

def extract_patches_by_loop(x, k=3, s=1): n, c, h, w = x.shape patches = [] for i in range(0, h - k + 1, s): for j in range(0, w - k + 1, s): patch = x[:, :, i:i+k, j:j+k] patches.append(patch) return torch.stack(patches, dim=-1)

这段代码看起来很好懂，但实际用起来有几个问题：第一，Python循环在patch数量多的时候相当慢，一张512x512的图用3x3窗口、步长1会产生约26万个patch，循环一次要等很久；第二，每个patch都涉及一次切片和一次stack，会产生大量中间张量，显存容易浪费；第三，pyTorch的自动求导虽然能穿过这个流程，但因为中间操作太多，反向传播的效率也不理想。torch.nn.Unfold把这些全部下沉到底层实现，并且把窗口提取统一成一次张量运算，遇到batch数据时优势更明显。

1.2 im2col的直观理解

img2col这个名字来自“image to column”，把图像转成列的集合。它的核心逻辑可以类比成：把一本大书的一页纸剪成若干小方块，再把每个小方块里的内容按顺序抄写成一列清单，所有小方块抄完后，整张纸就变成一张“列清单”。Unfold做的就是这件事：每个滑动窗口的内容被拉直成向量，然后按窗口所在位置排列成列。

这种做法的意义在于，原本卷积操作需要在多个空间位置上重复做乘加运算，如果先把所有窗口变成列，那么卷积就变成了一个矩阵乘法，底层可以调用高度优化的矩阵运算库。这也是为什么很多深度学习框架在实现卷积时，会在底层使用im2col配合通用矩阵乘法，性能远高于逐位置循环。torch.nn.Unfold把这个过程暴露给了开发者，等于给了大家一个“可微分的可编程滑动窗口提取器”。

1.3 输入输出张量的维度变化

假设输入x的形状是(N, C, H, W)，传入kernel_size=(kH, kW)，不填充、步长为s，Unfold的输出是：

patches = unfold(x) patches.shape = (N, C * kH * kW, L)

其中L表示总共有多少个窗口，等于输出特征图空间位置的个数。举个例子：输入(1, 3, 4, 4)，窗口大小3x3，步长1，无padding，这时能在4x4图上放(4-3+1)^2 = 4个3x3窗口。C * kH * kW = 3 * 3 * 3 = 27，所以输出形状是(1, 27, 4)。

第一次看到这个形状可能会有点奇怪：为什么第二维包含了C, kH, kW三个维度的乘积？因为每个窗口的所有通道、所有行、所有列都被拉成了一个一维向量，这样一来，第二维就是每个窗口内部元素的完整特征集合，第三维就是窗口的位置编号。理解了这个布局，后面的索引和重塑就会顺很多。

1.4 Unfold和卷积的本质关系

Unfold和卷积的关系非常密切。从计算流程来看，卷积本质上可以拆成两步：第一步用Unfold把输入变成patch矩阵，第二步对每个patch做线性映射（也就是权重矩阵的乘法）再reshape。换句话说，F.conv2d的很多底层实现就是Unfold加矩阵乘法。写成一个伪代码就是：

# unfold_x: (N, C*kH*kW, L) unfold_x = torch.nn.Unfold(kernel_size=k)(x) # weight: (C_out, C_in*kH*kW), bias: (C_out,) out = (weight @ unfold_x + bias.unsqueeze(1)).view(N, C_out, H_out, W_out)

之所以要强调这层关系，是因为当你需要自定义一个卷积变体时，完全可以用Unfold作为基础算子，自己控制权重作用方式。比如想在每个窗口内加入位置权重、通道权重，或者想在不同patch之间共享不同参数，直接用torch.nn.Conv2d就不够灵活了，但自己通过Unfold做往往只需要几十行代码。这也是我在很多自定义网络结构里优先考虑Unfold的原因。

2. Unfold参数选型：kernel_size、stride、padding、dilation的计算与避坑

2.1 四个参数各管什么

Unfold的初始化签名和卷积非常像，常用参数有kernel_size、dilation、padding、stride。它们分别控制窗口的尺寸、窗口内元素间隔、边缘补零长度和窗口移动步长，语义与Conv2d完全一致：

这四个参数组合起来决定窗口覆盖范围。dilation很多同学容易忽略，但它非常有用。比如想把一个3x3的窗口变成“间隔一点的3x3”，实际感受野达到5x5，可以设置dilation=2，这样不需要真正放大kernel尺寸也能增大感受野，内存占用还更小。

2.2 输出patch数量的公式推导

Unfold输出张量的第三个维度L，对应窗口个数，和卷积输出的空间尺寸公式一模一样：

H_out = floor((H + 2 * padding_h - dilation_h * (kernel_h - 1) - 1) / stride_h + 1) W_out = floor((W + 2 * padding_w - dilation_w * (kernel_w - 1) - 1) / stride_w + 1) L = H_out * W_out

推导逻辑很简单：窗口在水平方向上能停留的位置数等于W_out，垂直方向上能停留的位置数等于H_out，两者相乘就是全部窗口数。之所以要单独解释，是因为很多人直接用H_out和W_out当作patch坐标矩阵的行列数，一旦理解错了，后面重建原图时索引很容易出错。

以大图H=W=5、窗口3x3、步长2、padding1为例，代入公式：

H_out = floor((5 + 2*1 - 3 - 1) / 2 + 1) = floor(4/2 + 1) = 3 W_out = 3 L = 9

也就是说，原图尺寸没变但通过padding补了一圈0后，窗口可以在水平方向取3个位置：起始坐标分别为0、2、4。这正好对应卷积中padding=1保持空间尺寸的情况，直观而且好记。

2.3 用一个具体的数值例子验证

为了彻底把公式和形状对应起来，我建议在本地跑一下这个最简单的验证代码：

import torch x = torch.arange(1 * 1 * 4 * 4, dtype=torch.float32).view(1, 1, 4, 4) print(x) # tensor([[[[ 0., 1., 2., 3.], # [ 4., 5., 6., 7.], # [ 8., 9., 10., 11.], # [12., 13., 14., 15.]]]]) unfold = torch.nn.Unfold(kernel_size=3, dilation=1, padding=0, stride=1) patches = unfold(x) print(patches.shape) # torch.Size([1, 9, 4])

此时kH=3, kW=3，单通道，所以第二维是9。四个窗口分别对应原图中的左上、右上、左下、右下四个3x3块。第一个窗口的展开结果是0,1,2,4,5,6,8,9,10，正好是左上角那个3x3方块按行读取。遇到这类问题，用固定的连续数值构造输入，一眼就能看出行列读取顺序是否正确。

2.4 实际项目中的参数选择经验

在实际项目中，参数选择更多取决于后续任务需求。如果目的是做滑动窗口推理，通常希望有少量重叠来避免窗口边界的信息割裂，这时可以在stride小于kernel_size，比如窗口8x8、步长4。如果想要在通道维做池化前的局部特征提取，一般取padding=(kernel-1)//2、stride=1，这样输出窗口数量和原图尺寸一致，方便后续按像素位置对齐。

另外，dilation在需要扩大感受野而不增加参数量的场景下很实用。我做过一个分割任务，希望每个窗口能看到更广的上下文，但显存又有限，于是把kernel_size=3, dilation=2，实际感受野相当于5x5，patch数量不变，内存开销没有明显增长。不过使用dilation后要注意，窗口内的元素不再连续，做可视化时如果直接reshape成二维图像，会看到间隔排列的像素，不要以为是自己代码写错了。

3. 实操：用Unfold提取patch、查看patch、拼回原图

3.1 最简提取与维度打印

先写一个完整可运行的流程，直接从张量构造开始：

import torch import torch.nn.functional as F # 构造一个单通道4x4输入 x = torch.arange(1 * 1 * 4 * 4, dtype=torch.float32).view(1, 1, 4, 4) # 实例化Unfold unfold = torch.nn.Unfold(kernel_size=3, dilation=1, padding=0, stride=1) # 提取patch patches = unfold(x) print(patches.shape) # torch.Size([1, 9, 4])

这里patches的第二维是C*kH*kW=9，第三维是4。打印第一个窗口的内容可以写：

print(patches[0, :, 0]) # tensor([0., 1., 2., 4., 5., 6., 8., 9., 10.])

这串数字和原图左上角3x3块完全一致。从实际操作的角度讲，强烈建议在新项目里先用这种arange构造的可控输入验证一遍，再换成真实图像，否则后面出了问题很难定位。

3.2 怎么判断patch提取得对不对

最直接的方法是把patches按(C, kH, kW, L)重新view，再取某一列变成二维矩阵看：

C, kH, kW = 1, 3, 3 patches_view = patches[0].view(C, kH, kW, -1) # (1, 3, 3, L) first_patch = patches_view[0, :, :, 0] print(first_patch) # tensor([[0., 1., 2.], # [4., 5., 6.], # [8., 9., 10.]])

这样做的好处是，可以直接用肉眼看懂每个patch的空间结构。遇到多通道输入时，patches_view的第0维是通道，第1、2维是窗口的行列，第3维是窗口编号，排列顺序非常清晰。如果你发现第一个patch不符合预期，大概率是stride、padding或者输入维度的排列出了问题。常见的错误是把(N, H, W, C)格式的输入直接传给Unfold，结果窗口切片完全错位，因为Unfold默认输入是(N, C, H, W)，如果是NHWC需要先permute。

3.3 把patch重新拼回原图

Unfold把图切开，自然也需要逆操作把patch拼回去。PyTorch有torch.nn.Fold，但它默认处理方式是重叠区域求和，如果窗口之间有重叠，直接Fold会得到“叠加”的结果，而不是原图。在做可视化和重建时，我通常自己写一个带计数的重建逻辑，把重叠区域做平均：

def patches_to_image(patches, C, H, W, kH, kW, stride=1, padding=0): H_out = (H + 2 * padding - kH) // stride + 1 W_out = (W + 2 * padding - kW) // stride + 1 L = patches.shape[-1] assert H_out * W_out == L, "patch数量与输出尺寸不匹配" img = torch.zeros(C, H, W, dtype=patches.dtype) count = torch.zeros(C, H, W, dtype=patches.dtype) for idx in range(L): h_idx = idx // W_out w_idx = idx % W_out top = h_idx * stride - padding left = w_idx * stride - padding patch = patches[:, :, idx].view(C, kH, kW) # 计算有效区域，防止越界 h_start = max(0, top) w_start = max(0, left) h_end = min(H, top + kH) w_end = min(W, left + kW) ph_start = h_start - top pw_start = w_start - left ph_end = ph_start + (h_end - h_start) pw_end = pw_start + (w_end - w_start) img[:, h_start:h_end, w_start:w_end] += patch[:, ph_start:ph_end, pw_start:pw_end] count[:, h_start:h_end, w_start:w_end] += 1 img = img / count.clamp(min=1) return img

这个函数里最关键的是top和left的计算。窗口编号idx换算成行号和列号之后，要减去padding才能得到窗口在原图中的起始位置。如果加了padding，窗口的某些区域实际上落在原图外面，对应部分要舍弃，否则会越界报错。最后用计数做平均，能让边界处因为补0产生的暗边被淡掉。实测下来，只要参数设置正确，重建出的图和原图基本一致，误差只有浮点精度级别。

3.4 封装成可复用的工具函数

实际项目中我会把Unfold相关的操作封装成一个类，方便在多个地方复用：

class PatchExtractor: def __init__(self, kernel_size, stride=1, padding=0, dilation=1): self.unfold = torch.nn.Unfold( kernel_size=kernel_size, dilation=dilation, padding=padding, stride=stride ) self.kernel_size = kernel_size if isinstance(kernel_size, tuple) else (kernel_size, kernel_size) self.stride = stride self.padding = padding def __call__(self, x): return self.unfold(x) def reconstruct(self, patches, x_shape): n, c, h, w = x_shape kH, kW = self.kernel_size return patches_to_image( patches, c, h, w, kH, kW, stride=self.stride, padding=self.padding )

这样一个类既能提取patch，又能重建原图，用起来很顺手。我在做图像修复和超分辨率实验时，经常需要把图切成patch输入模型，再把模型输出的patch拼回原图，这种封装能节省很多重复代码。需要注意的是reconstruct函数里传入的x_shape是原始输入的形状，如果中间经过了通道数变化，要对应调整。

4. 常见问题与踩坑记录：维度、顺序、padding和显存

4.1 问题速查表

先列一个我在调试Unfold过程中遇到的高频问题速查表，方便你直接对照：

4.2 灰度图也要扩成4维

torch.nn.Unfold要求输入必须至少有(N, C, H, W)四个维度。很多人在处理单张灰度图时直接传形状为(H, W)或(C, H, W)的张量，结果报错expected 4D input。解决办法很简单，用x.unsqueeze(0)补batch维，或者用x.unsqueeze(1)补通道维。如果灰度图的形状是(H, W)，需要先view(1, 1, H, W)，再交给Unfold。这个问题出现的频率比想象中高得多，建议在封装工具函数里自动检查维度。

4.3 patch的行排列顺序是通道优先

Unfold输出第二维的组织顺序是：先遍历通道，再遍历窗口的行，最后遍历窗口的列。也就是说，patches[0, :, l]并不是“第l个窗口的空间展开”，而是“第0通道的第0行第0列、第0通道第0行第1列……”。这在做自定义矩阵乘法时很关键。比如想手写一个卷积核：权重形状是(C_out, C_in, kH, kW)，要把它应用到Unfold输出上，需要先weight.view(C_out, C_in*kH*kW)再和patches做矩阵乘，顺序才匹配。如果直接用weight.view(C_out, kH*kW*C_in)，顺序就会错。

4.4 padding产生的0会让重建图偏暗

使用padding之后，Unfold输出的边缘patch会包含补上的0像素。如果直接用Fold恢复原图，重叠区域会被求和，零值虽然不影响求和，但如果之后做归一化时除以固定窗口数量L，没有重叠的位置会偏暗。解决办法就是我前面给的计数平均法：统计每个输出像素被多少个patch真实覆盖，再除以计数，这样重叠区域取平均，边缘区域也不会被0污染。如果你只需要保留内部区域，也可以提前把padding区域裁掉。

4.5 显存开销要提前估算

Unfold带来的最大隐性成本是显存。算一笔账：输入(1, 3, 256, 256)，窗口7x7，步长1，无padding，输出形状是：

行数 = C * kH * kW = 3 * 49 = 147 L = (256 - 7 + 1)^2 = 250 * 250 = 62500 元素总量 = 147 * 62500 = 9,187,500 float32字节 = 约36.75 MB

这还只是单张图、单向一层的量。如果batch增大、层数变多或者pytorch需要保存梯度用于反向传播，内存占用会成倍增长。所以在设计网络时，如果窗口大、步长小，最好估算一下中间张量大小，必要时把大图分成子区域处理，或者使用更大步长减少patch数量。

5. Unfold的扩展玩法：窗口自注意力、可微滑动窗口与手写卷积

5.1 把Unfold接入窗口自注意力

视觉Transformer里经常把图像划分成固定大小的窗口，在每个窗口内做自注意力。用Unfold实现这个切分过程非常自然。思路是先用Unfold提取窗口，再把输出形状reshape成(N, L, C*kH*kW)，这样就得到了类似token序列的输入。

def window_partition(x, window_size): n, c, h, w = x.shape unfold = torch.nn.Unfold(kernel_size=window_size, stride=window_size) windows = unfold(x) # (N, C*wh*ww, L) windows = windows.permute(0, 2, 1) # (N, L, C*wh*ww) return windows

拿到窗口token之后，可以继续加位置编码、做多头注意力，最后用Fold把结果拼回原尺寸。相比手动reshape，这种做法的好处是对H和W是否整除窗口大小没有硬性要求，配合padding可以处理任意尺寸，边界逻辑也更统一。

5.2 可微的滑动窗口处理流程

在图像修复、去噪、超分辨率等任务里，最常见的流程是：切patch、逐patch推理、拼回全图。如果整个过程都使用Unfold和Fold，那么整条链路是可微的，可以直接端到端训练。这一点在实际项目中意义很大：你不需要在patch边界处做额外处理，模型可以学习到重叠区域的一致性。需要注意的是，torch.nn.Fold在重叠区域默认求和，如果你的模型输出需要加权融合，建议在Fold之前对每个patch乘以一个可学习的权重矩阵，或者自己实现带计数的加权Fold逻辑。

5.3 手写一个卷积来验证理解

为了验证自己对Unfold布局的理解是否正确，可以写一个简化版卷积，和官方Conv2d做对比：

import torch.nn as nn def conv_via_unfold(x, weight, bias=None, stride=1, padding=0): n, c, h, w = x.shape c_out, c_in, kh, kw = weight.shape unfold = torch.nn.Unfold(kernel_size=(kh, kw), padding=padding, stride=stride) patches = unfold(x) # (N, c_in*kh*kw, L) weight_flat = weight.view(c_out, c_in * kh * kw) out = weight_flat @ patches # (N, C_out, L) if bias is not None: out = out + bias.view(1, -1, 1) h_out = (h + 2 * padding - kh) // stride + 1 w_out = (w + 2 * padding - kw) // stride + 1 return out.view(n, c_out, h_out, w_out)

用随机张量对拍一下，输出和F.conv2d在数值上几乎完全一致，只有极小的浮点误差。这个实验做完之后，你对Unfold的维度排列和卷积底层逻辑的理解会立刻上一个台阶。我建议每个想深入PyTorch底层的人都跑一遍这种对拍测试，花不了几分钟，但能消除很多模糊感。

5.4 性能优化与使用建议

最后给你一个性能上的建议：Unfold适合中等尺寸的窗口提取，如果窗口特别大或者L特别大，中间矩阵会占掉大量显存。我遇到过一个场景，输入1024x1024、窗口16x16、步长8，Unfold输出的中间张量直接撑爆了显存，最后改成把图像切成四块分别处理才解决。另外，如果只是想在数据预处理阶段提取patch，并不需要梯度，可以加上with torch.no_grad():，同时把结果转换到CPU或者磁盘存储，避免GPU显存和训练模型叠加。Unfold是一个看起来简单、用起来很爽、但细节容易出错的操作，我的建议是在自己的项目里先用小张量把索引关系验证一遍，再上全尺寸数据，这样可以省去大量排查时间。

我在实际使用中养成了一个习惯：每次构建新的Unfold流程时，都会先打印patches.shape，再随机取其中一列和原图对应区域做切片对比，确认stride和padding都算对了才继续往下写。有一次我把padding方向理解反了，导致训练loss怎么都降不下去，浪费了一整天，最后是用这份对照结果定位到问题的。希望你看完这篇之后，能少走我走过的这些弯路。