如果你碰过 Vision Transformer 的前处理、写过图像修复里的相似块匹配，或者只是想把一张图像按固定大小、固定步长切成一堆窗口，多半会有这样的一瞬间：循环写到一半，突然意识到 PyTorch 里好像有个函数天生就是干这个的。对，就是torch.nn.Unfold（以及功能等价的torch.nn.functional.unfold）。它在 PyTorch 里待了很多年，却经常被忽略，大多数人只在看卷积实现源码时才偶然翻到它。

这篇东西我想一次性把这个函数讲透：它输出的张量到底按什么顺序排列、怎么把结果还原成我们习惯的"一堆 patch"、为什么它能用来做 Transformer 里的 patch embedding、以及配合Fold使用时重叠区域的累加问题。我还会给出几个实际项目里验证过的替代方案和避坑经验。如果你正在做图像分类、目标检测、图像复原或者任何涉及"局部窗口"特征提取的任务，这篇应该对你有用。

1. Unfold 背后的滑动窗口语义：先搞懂它到底在做什么

Unfold的数学名字叫 im2col（image to column），但比传统 im2col 更简洁一点：它在输入特征图上按kernel_size滑动窗口，把每个窗口内的元素"拍平"成一个列向量，然后把所有窗口的列向量拼起来。就这么简单，没有任何可学习参数，也没有任何花哨的变换。

1.1 用 5×5 小图看懂 Unfold 的输出

与其背文档，不如直接跑一个最小例子。我生成一张 5×5 的单通道图像，值从 0 到 24，然后用一个 3×3 的窗口、步长 1、无 padding 去提取 patch：

import torch import torch.nn as nn import torch.nn.functional as F x = torch.arange(25.0).reshape(1, 1, 5, 5) unfold = nn.Unfold(kernel_size=3, dilation=1, padding=0, stride=1) out = unfold(x) print(out.shape) # torch.Size([1, 9, 9]) print(out[0, :, 0].reshape(3, 3)) # 取出第一个窗口的内容

输出是(1, 9, 9)，这三个数字的含义我后面马上解释。先看第一个窗口，也就是输出张量中第 0 列的内容：

tensor([[ 0., 1., 2.], [ 5., 6., 7.], [10., 11., 12.]])

这正是原图左上角 3×3 的区域。因为图像值是arange(25)按行填充的，所以左上角自然就是 0 到 12 这一块。第二个窗口呢？向右移动一格，对应原图第 0 到 2 行、第 1 到 3 列，内容应该是：

tensor([[ 1., 2., 3.], [ 6., 7., 8.], [11., 12., 13.]])

以此类推。窗口滑到最右边后回到下一行继续滑，最终 3×3 的滑动范围一共产生 9 个窗口。这就是输出第三个维度是 9 的原因，它表示窗口数量。

1.2 窗口数量 L 的计算公式与 padding/dilation 的作用

窗口数量的计算和卷积输出尺寸的公式其实一模一样。对高度方向，窗口能滑过的位置数是：

L_h = floor((H + 2 * padding - dilation * (kernel_size - 1) - 1) / stride) + 1

宽度方向同理，最终的窗口总数L = L_h * L_w。还是用刚才的例子验证：H=5, kernel_size=3, padding=0, dilation=1, stride=1，代入得到：

L_h = (5 + 0 - 1 * 2 - 1) / 1 + 1 = 3

所以L = 3 * 3 = 9。换成padding=1的话，等效输入尺寸变成了 7×7，L_h变成 5，窗口数直接跳到 25。

dilation的作用是扩大窗口的"视野"而保持kernel_size参数不变。比如kernel_size=3, dilation=2时，实际覆盖范围是 5×5，但只有 4 个采样点。这个参数在空洞卷积类任务里用得多，纯 patch 提取的场景一般不会碰它，但理解它对读懂公式有好处。

1.3 和 Conv2d 的内在联系：Unfold 就是去掉权重的卷积

很多人第一次看到 Unfold 时，其实已经见过它了——它就在卷积的底层计算里。Conv2d的前向传播可以拆成两步：

1. 用 Unfold 把输入特征图变换成(N, C * kh * kw, L)的矩阵；
2. 把卷积权重 reshape 成(out_channels, C * kh * kw)，和上面的矩阵做矩阵乘法，再加上偏置。

换句话说，Unfold 输出了卷积层在"加权求和"之前的所有候选窗口数据。之所以把窗口拍平成向量、把通道数乘到窗口内部维度上，是为了让后续的mm或addmm能直接跑在高性能矩阵乘法库上。理解这一层关系后，你就明白为什么 Unfold 的输出是(N, C * kh * kw, L)而不是更"符合人类直觉"的(N, L, C, kh, kw)了。

所以如果你期望的是一个(batch, num_patches, height, width, channels)这样的窗口序列，Unfold 的原始输出不够直观，但它的信息量和那个形式完全等价，只需要一步 reshape 加 permute。

2. 输出维度里的排列密码：N、C、kh*kw、L 到底怎么解读

这是最容易把人绕晕的一节。Unfold返回的每个 batch 元素是一个形状为(C * kh * kw, L)的二维矩阵，我把它拆开讲。

2.1 通道 C 为什么排在空间位置前面

先看一个多通道例子。假设输入是(1, 3, 5, 5)，同样用 3×3 窗口、步长 1：

x = torch.randn(1, 3, 5, 5) out = F.unfold(x, kernel_size=3, stride=1) print(out.shape) # torch.Size([1, 27, 9])

中间维度从 9 变成了 27，也就是3 * 3 * 3。这个 27 的内部顺序是：先枚举通道（channel），再枚举窗口内的行，最后枚举窗口内的列。用公式表示，第c个通道、窗口内第kh行、第kw列对应的展平索引是：

idx = c * (kh * kw) + kh * kw + kw

也就是说，连续的 9 个数属于同一个通道的同一个窗口，而不是把三个通道的同一行拼在一起。这个排列方式和卷积权重是一致的：卷积核权重的原始形状是(out_channels, in_channels, kh, kw)，flatten 后自然就是"通道优先、空间在后"。

如果你想确认某个窗口内的原始 patch 内容，可以用这样的方式还原：

C, kh, kw = 3, 3, 3 patch_idx = 0 patch = out[0, :, patch_idx].reshape(C, kh, kw) # patch.shape == (3, 3, 3)，即 3 个通道各一个 3×3 块

注意reshape是按行填充的，这正好符合 C-major 的存储顺序。我最初用这个函数时犯过一个错：直接对out做permute(0, 2, 1)就当成(N, L, C, kh, kw)用，结果把通道和空间的顺序彻底搞反了，后续网络怎么训都不收敛。后来统一改成先 reshape 再 permute，问题就消失了。

2.2 把 Unfold 输出还原成"窗口序列"的通用姿势

实际代码里我们大多数时候想要的是"第 i 个 patch 长什么样"，也就是原始的二维信息。标准还原写法如下：

def unfold_to_patches(x, kernel_size, stride): b, c, h, w = x.shape kh = kw = kernel_size unfold_out = F.unfold(x, kernel_size=kernel_size, stride=stride) # unfold_out: (N, C * kh * kw, L) unfold_out = unfold_out.reshape(b, c, kh, kw, -1) unfold_out = unfold_out.permute(0, 4, 1, 2, 3) # 返回形状: (N, L, C, kh, kw) return unfold_out

permute的作用是把最后那个L维度挪到 batch 维度之后，得到(N, L, C, kh, kw)。这个形状才是符合我们直觉的"每个样本有 L 个 patch，每个 patch 是 C 通道的 kh×kw 小块"。

如果后续要接 Transformer 风格的序列模型，通常还要把空间维度和通道维度合并成向量，也就是再看一眼reshape：

patches = unfold_to_patches(x, kernel_size=16, stride=8) b, l, c, ph, pw = patches.shape patch_vectors = patches.reshape(b, l, c * ph * pw) # 现在 patch_vectors 每一行就是一个 patch 的展平向量

到这一步，Unfold 的使命基本完成。

2.3 直接对中间维度做 Linear 时的权重对齐

有时候我们不需要还原成(N, L, C, kh, kw)，而是想直接在 Unfold 的原始输出上加一个线性层，最典型的就是 patch embedding。这时候要保证线性层的输入维度是C * kh * kw，并且它是作用在第二个维度（C * kh * kw）上的。

# unfold_out: (N, C*kh*kw, L) linear = nn.Linear(C * kh * kw, embed_dim) # 注意：Linear 默认作用在最后一个维度上，所以要先把 L 挪到最后 unfold_out = unfold_out.permute(0, 2, 1) # (N, L, C*kh*kw) embedded = linear(unfold_out) # (N, L, embed_dim)

这个写法在数学上和Conv2d(in_channels, out_channels, kernel_size, stride)完全等价，前提是 stride 等参数一致、而且你用的是逐 patch 的线性变换而不是跨位置的卷积。实际工程里，如果你目标就是得到(N, L, embed_dim)，直接用Conv2d会更快，因为卷积底层做了更多优化、不需要把窗口数据全部展开到显存里；Unfold + Linear 的优势在于灵活——线性层之后你还可以对每个 patch 做任意非线性操作，不局限于逐点卷积的形式。

3. 实战：图像 patch 块提取的工程写法与 Transformer patch embedding

理论说完了，来点直接能抄的代码。我会对比三种做法：朴素切片、Unfold、以及无重叠场景下的轻量 reshape。顺便展示一个图像修复任务里的候选块匹配例子。

3.1 朴素切片写法与 Unfold 写法对比

大多数人第一反应是用双层循环切片。假设要提取的所有 patch 之间无重叠、patch 大小为p，朴素写法长这样：

def extract_patches_slice(x, p): b, c, h, w = x.shape patches = [] for i in range(0, h - p + 1, p): for j in range(0, w - p + 1, p): patch = x[:, :, i:i+p, j:j+p] patches.append(patch) return torch.stack(patches, dim=1) # (B, L, C, p, p)

这个写法的问题有两个。一是慢：Python 循环在 patch 数量多的时候开销很夸张；二是它隐式要求h和w都能被p整除，否则会漏掉右侧和下侧的区域。更麻烦的是，torch.stack会再拷贝一次数据，显存和带宽都被白白浪费。

换成 Unfold 之后，同样的功能只需要两行：

def extract_patches_unfold(x, p, stride): b, c, h, w = x.shape pad_h = (p - h % p) % p pad_w = (p - w % p) % p x_pad = F.pad(x, (0, pad_w, 0, pad_h)) out = F.unfold(x_pad, kernel_size=p, stride=stride) # out: (B, C * p * p, L) out = out.reshape(b, c, p, p, -1).permute(0, 4, 1, 2, 3) return out

我在一张 512×512 的 RGB 图上做过简单对比，patch 大小 16×16、无重叠，次数取 100 次。朴素切片跑了接近 4 秒，Unfold 版本只用了不到 0.1 秒，差距接近两个数量级。这个差距在 batch 维度变大后会进一步拉大，因为 Unfold 本质上是纯矩阵运算，能充分发挥 GPU 的并行能力。

3.2 无重叠 patch 提取就是 ViT 的 patch embedding

Vision Transformer 里最经典的一个操作就是把图像切成不重叠的 16×16 patch，然后线性投影到指定维度。很多人第一次看源码时，看到的是Conv2d(in_ch, embed_dim, kernel_size=patch, stride=patch)，心想这不是卷积吗？其实它就是 patch embedding 的卷积实现。

等价的 Unfold 写法是这样的：

class PatchEmbedUnfold(nn.Module): def __init__(self, in_ch=3, patch_size=16, embed_dim=768): super().__init__() self.patch_size = patch_size self.proj = nn.Linear(in_ch * patch_size * patch_size, embed_dim) def forward(self, x): b, c, h, w = x.shape p = self.patch_size out = F.unfold(x, kernel_size=p, stride=p) # (B, C*p*p, L) out = out.permute(0, 2, 1) # (B, L, C*p*p) return self.proj(out) # (B, L, embed_dim)

两者在数学上严格等价。区别在于感受野和组织方式：卷积版本把 patch 内元素和投影矩阵融合成了一个巨大的卷积核，Unfold 版本则先把 patch 切出来再做矩阵乘法。如果你只是想要标准 ViT 的输入，直接用卷积版本效率最高；如果你需要保留每个 patch 本身做后续处理（比如做 token 之间的空间结构建模、局部注意力、或者把 patch 重新拼回图像），Unfold 版本会更直观，因为它把"切块"和"投影"明确分开了。

3.3 图像修复场景里"候选块匹配"的用法

再分享一个更实际的例子。图像修复里有一个经典做法：破损区域周围的像素块，需要在图像其他位置找最相似的完整块来做替换或插值。这个需求用 Unfold 写特别爽，因为你要做的就是从全图提取所有候选 patch，然后算相似度。我用一个简化版说明思路：

# x: (1, 3, H, W) patch_size = 5 # 提取全图所有 5×5 候选块 candidates = F.unfold(x, kernel_size=patch_size, stride=1) # candidates: (1, 3*5*5, L)，L 是所有可提取窗口的数量 # 目标块：假设我们要匹配的是坐标 (10, 10) 处的 5×5 块 target = x[:, :, 10:15, 10:15].reshape(1, -1, 1) # (1, 75, 1) # 计算每个候选块与目标块的欧氏距离 dist = ((candidates - target) ** 2).sum(dim=1) # (1, L) best_idx = dist.argmin(dim=1).item()

这里有个关键点：candidates的输出维度把空间索引完全放在了L里，所以argmin得到的索引可以反推出原图坐标。反过来，如果你想把找出来的相似块"放回"原图某个位置，就用Fold。这个模式在做纹理合成、视频插帧、光流估计的 patch 匹配模块里都出现过。Unfold + Fold 的组合本质上就是"选址-采集-回填"的完整流程，比手写偏移循环要稳得多。

4. 反向操作 Fold：重叠区域的累加陷阱与归一化修复

有 Unfold 就有 Fold。torch.nn.Fold做的事情刚好反过来：把(N, C * kh * kw, L)的 patch 集合按对应位置放回一张输出特征图。听起来像 Unfold 的逆过程，但它的默认行为很容易让第一次用的人翻车。

4.1 Fold 到底做了什么：一次小实验看清累加

我先给一个 3×3 的单通道图像，用 2×2 窗口、步长 1 做 Unfold 再立刻 Fold 回去，看看会发生什么：

x = torch.arange(9.0).reshape(1, 1, 3, 3) u = F.unfold(x, kernel_size=2, stride=1) # (1, 4, 4) r = F.fold(u, output_size=(3, 3), kernel_size=2, stride=1) print(r)

输出结果中，四角像素只被一个窗口覆盖，所以数值不变；四个边上的像素被两个窗口覆盖，数值变成原来的两倍；中心像素被四个窗口覆盖，数值变成原来的四倍。具体到16这个中心值，就是原本的4被累加了 4 次。

换句话说，Fold在重叠区域执行的是sum，不关心一个位置到底被多少个窗口覆盖。如果你拿着一个本来不重叠的 patch 集合直接掉进有重叠的窗口参数里，回填后图像会"变亮"或者数值整体偏大，这在做图像重建任务时是非常隐蔽的 bug。

4.2 重叠区域归一化：让 Fold 真正变成 Unfold 的逆过程

要让 Fold 在重叠情况下还原出正确的逐像素结果，标准做法是同时算一个"覆盖次数矩阵"。把所有位置都当成 1，对这个全 1 张量做同样的 Unfold 再 Fold，就能知道每个位置被多少个窗口覆盖，然后用它做除法：

def unfold_fold_restore(patches, output_size, kernel_size, stride): ones = torch.ones(1, 1, output_size[0], output_size[1], device=patches.device) count = F.fold( F.unfold(ones, kernel_size=kernel_size, stride=stride), output_size=output_size, kernel_size=kernel_size, stride=stride ) restored = F.fold( patches, output_size=output_size, kernel_size=kernel_size, stride=stride ) return restored / count

count矩阵的形状和输出特征图一样，每个位置的值就是覆盖次数。这样除完之后，每个像素就被还原成了所有覆盖它的窗口的平均贡献，而不是简单求和。

我之前在做一个图像复原 Demo 时用这个技巧做过"滑动窗口去噪再拼回原图"的实验：把有重叠的 patch 做处理后 Fold 回原图，不归一化时拼接缝特别明显，一旦用覆盖次数归一化，拼接缝立刻消失了。这个操作在医学影像重建、超分辨率 tile 合并、去块效应后处理里都能用上。

4.3 梯度在 Unfold/Fold 之间的传播

Unfold 和 Fold 都支持自动求导，它们本质上是 scatter 和 scatter_add 的封装。Unfold 的正向是从原图"取"数据，反向就是把梯度按照同样的位置关系"写回"原图；Fold 的正向是往输出位置"累加"，反向则把梯度"原样分发给所有贡献过该位置的窗口"。

这里有个容易被忽视的性质：当同一个像素被多个 patch 覆盖时，它的梯度也是累加的。所以如果在 Unfold 和 Fold 之间插入了逐窗口的处理（比如 softmax 权重、自注意力），反向传播时梯度会自动按覆盖关系分配到多个路径上，这种结构在局部注意力模型里其实是被特意利用的——多个窗口共享同一区域时，梯度信息会更丰富。

5. 边界与替代：什么时候别用 Unfold

Unfold 很好用，但它不是万能的。有些场景用它是自找麻烦，我列几个实际匹配过的坑。

5.1 显存开销是最大短板

Unfold 的输出大小和"窗口重叠程度"强相关。窗口重叠越多，同一份像素会被复制到多个窗口里，显存占用成倍上涨。

拿一张(1, 3, 512, 512)的图像举例，用patch_size=16, stride=8做有重叠提取：

窗口数 L = ((512 - 16) / 8 + 1)^2 = 63^2 = 3969 每个窗口展平长度 = 3 * 16 * 16 = 768 输出形状 = (1, 768, 3969)

float32 下这个输出张量约 12MB，而原图本身只有约 3MB。也就是说 Unfold 把显存占用量放大了差不多 4 倍。如果输入是 batch=16 的 1024×1024 医学图像，显存直接爆掉不是开玩笑。碰到这种场景，我优先考虑用卷积或局部 attention 的官方实现来替代显式展开，或者走分块处理路线。

5.2 Tensor.unfold（一维）和 nn.Unfold（二维）的区别

PyTorch 里其实还有一个同名但不同层级的 API：Tensor.unfold(dim, size, step)。它作用在单个维度上，返回的是一个 view，不 copy 数据，非常适合处理一维序列、时间序列的滑窗。

两者的核心区别在于：

如果你只是做时间序列的滚动窗口特征，千万别用nn.Unfold，先得把序列强行 reshape 成 4D 是很别扭的事情。反过来，如果你要做二维图像 patch，用一维unfold得自己在 H 和 W 两个方向各调一次再组合，麻烦多了。这两个函数各管一段，不是互相替代的关系。

5.3 无重叠场景下的轻量替代

当 stride 等于 kernel_size、窗口完全不重叠时，Unfold 做的事情其实可以用原生 reshape 加 permute 一行搞定，而且完全不用额外展开数据：

def extract_patches_reshape(x, p): b, c, h, w = x.shape # 要求 h 和 w 都能被 p 整除 return x.reshape(b, c, h // p, p, w // p, p).permute(0, 2, 4, 3, 5, 1).reshape(b, -1, p * p * c)

这个方案在无重叠、且hw恰好整除时，和 Unfold 的结果完全一致，但不会产生额外的显存拷贝。patch 数量巨大的时候优势很明显。我之前在做高分辨率图像的 patch 级分类任务时，就因为这个 reshape 写法省下了接近 30% 的显存。

当然它也有局限：不支持重叠、不支持 padding、要求整除，而且一旦 stride 小于 kernel_size，这个方案立刻失效。所以结论很清晰：

• 有重叠需求、需要灵活控制 padding 或 dilation：用 Unfold；
• 无重叠、整除且追求极致显存效率：用 reshape；
• 只想沿一个维度滑窗：用Tensor.unfold；
• 目标是(B, L, D)这种 embedding：优先用Conv2d的等价写法，省显存、速度快。

5.4 一个小技巧：用 Unfold 调试卷积

最后补充一个排查问题的小技巧。有时候你写了自定义卷积逻辑，不确定卷积核到底看了哪些区域，直接把输入过一个nn.Unfold(kernel_size=k, stride=s, padding=p)，再把输出的第 i 列 reshape 成(C, k, k)，就能看到这个窗口覆盖的原始像素内容。这个调试方式比在卷积中间打印梯度清楚得多，因为 Unfold 本身没有任何权重干扰，输出就是纯纯的输入采样。

我在排查一个自定义空洞卷积的参数对齐问题时，正是靠这个办法确认了 dilation 后的采样点坐标，一眼看到底是哪个位置算错了。这类"拿 Unfold 当透视镜用"的经验，比单纯背 API 文档要实用得多。

说到底，torch.nn.Unfold不是那种天天出现在模型定义里的明星算子，但凡是涉及局部窗口提取的任务，它几乎都是最优解之一。我建议你第一次用之前先花十分钟跑一遍上面的小实验，把输出维度和排列顺序彻底搞清楚，后续再遇到和 patch 相关的处理都会顺很多。实际调参时如果发现显存紧张，优先考虑重叠度和数据拷贝，别一上来就怪 batch size 太大。先理解它的语义，再用它作为 API，你才不会在某个深夜被一个莫名其妙的 0 维度索引难住。