做图像patch块提取的时候很多人的第一反应是写双层循环挨个切片再torch.stack。我第一次做滑动窗口裁剪时也是这么干的直到后来在项目里发现torch.nn.Unfold可以一步把整张图的patch全部取出来速度快、代码短而且能直接接在神经网络里做可微操作。简单理解Unfold就是PyTorch对im2col的标准封装输入一张N,C,H,W的图像输出变成N, C*kH*kW, L这种“把每个滑动窗口摊平成列”的张量全程不需要自己管理索引。这篇文章会从原理、参数公式、实操代码、踩坑记录、扩展应用几个角度把Unfold讲透适合正在写自定义卷积、滑动窗口网络、图像复原和patch分类相关代码的人。1. Unfold的工作逻辑从patch循环到im2col1.1 写循环切patch的痛点先看一段最直觉的patch提取代码假设输入是x形状是(N, C, H, W)想提取大小为k、步长为s的所有块def extract_patches_by_loop(x, k3, s1): 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:ik, j:jk] 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)不填充、步长为sUnfold的输出是patches unfold(x) patches.shape (N, C * kH * kW, L)其中L表示总共有多少个窗口等于输出特征图空间位置的个数。举个例子输入(1, 3, 4, 4)窗口大小3x3步长1无padding这时能在4x4图上放(4-31)^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_sizek)(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完全一致参数作用默认值注意事项kernel_size窗口大小可传int或(tuple)无和卷积的kernel一致stride窗口每次滑动的步长1越大窗口数量越少padding输入四周补0的宽度0补0会进入patchdilation窗口内元素之间的间隔1大于1时窗口实际覆盖范围变大这四个参数组合起来决定窗口覆盖范围。dilation很多同学容易忽略但它非常有用。比如想把一个3x3的窗口变成“间隔一点的3x3”实际感受野达到5x5可以设置dilation2这样不需要真正放大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坐标矩阵的行列数一旦理解错了后面重建原图时索引很容易出错。以大图HW5、窗口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。这正好对应卷积中padding1保持空间尺寸的情况直观而且好记。2.3 用一个具体的数值例子验证为了彻底把公式和形状对应起来我建议在本地跑一下这个最简单的验证代码import torch x torch.arange(1 * 1 * 4 * 4, dtypetorch.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_size3, dilation1, padding0, stride1) patches unfold(x) print(patches.shape) # torch.Size([1, 9, 4])此时kH3, kW3单通道所以第二维是9。四个窗口分别对应原图中的左上、右上、左下、右下四个3x3块。第一个窗口的展开结果是0,1,2,4,5,6,8,9,10正好是左上角那个3x3方块按行读取。遇到这类问题用固定的连续数值构造输入一眼就能看出行列读取顺序是否正确。2.4 实际项目中的参数选择经验在实际项目中参数选择更多取决于后续任务需求。如果目的是做滑动窗口推理通常希望有少量重叠来避免窗口边界的信息割裂这时可以在stride小于kernel_size比如窗口8x8、步长4。如果想要在通道维做池化前的局部特征提取一般取padding(kernel-1)//2、stride1这样输出窗口数量和原图尺寸一致方便后续按像素位置对齐。另外dilation在需要扩大感受野而不增加参数量的场景下很实用。我做过一个分割任务希望每个窗口能看到更广的上下文但显存又有限于是把kernel_size3, dilation2实际感受野相当于5x5patch数量不变内存开销没有明显增长。不过使用dilation后要注意窗口内的元素不再连续做可视化时如果直接reshape成二维图像会看到间隔排列的像素不要以为是自己代码写错了。3. 实操用Unfold提取patch、查看patch、拼回原图3.1 最简提取与维度打印先写一个完整可运行的流程直接从张量构造开始import torch import torch.nn.functional as F # 构造一个单通道4x4输入 x torch.arange(1 * 1 * 4 * 4, dtypetorch.float32).view(1, 1, 4, 4) # 实例化Unfold unfold torch.nn.Unfold(kernel_size3, dilation1, padding0, stride1) # 提取patch patches unfold(x) print(patches.shape) # torch.Size([1, 9, 4])这里patches的第二维是C*kH*kW9第三维是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, stride1, padding0): 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, dtypepatches.dtype) count torch.zeros(C, H, W, dtypepatches.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(min1) return img这个函数里最关键的是top和left的计算。窗口编号idx换算成行号和列号之后要减去padding才能得到窗口在原图中的起始位置。如果加了padding窗口的某些区域实际上落在原图外面对应部分要舍弃否则会越界报错。最后用计数做平均能让边界处因为补0产生的暗边被淡掉。实测下来只要参数设置正确重建出的图和原图基本一致误差只有浮点精度级别。3.4 封装成可复用的工具函数实际项目中我会把Unfold相关的操作封装成一个类方便在多个地方复用class PatchExtractor: def __init__(self, kernel_size, stride1, padding0, dilation1): self.unfold torch.nn.Unfold( kernel_sizekernel_size, dilationdilation, paddingpadding, stridestride ) 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, strideself.stride, paddingself.padding )这样一个类既能提取patch又能重建原图用起来很顺手。我在做图像修复和超分辨率实验时经常需要把图切成patch输入模型再把模型输出的patch拼回原图这种封装能节省很多重复代码。需要注意的是reconstruct函数里传入的x_shape是原始输入的形状如果中间经过了通道数变化要对应调整。4. 常见问题与踩坑记录维度、顺序、padding和显存4.1 问题速查表先列一个我在调试Unfold过程中遇到的高频问题速查表方便你直接对照问题现象原因解决办法报错说输入必须是4D输入缺少batch维度用unsqueeze(0)加一维patch内容顺序不对输入可能是NHWC先permute(0,3,1,2)成NCHW重建图有黑边padding补0混入了patch重建时按有效区域计数平均输出patch数量不是预期值stride或padding算错用公式floor((H2p-k)/s1)核对显存突然不够patch矩阵太大减小窗口、增大stride或分块处理梯度不回传使用了detach或单独处理确保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_sizewindow_size, stridewindow_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, biasNone, stride1, padding0): n, c, h, w x.shape c_out, c_in, kh, kw weight.shape unfold torch.nn.Unfold(kernel_size(kh, kw), paddingpadding, stridestride) 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、步长8Unfold输出的中间张量直接撑爆了显存最后改成把图像切成四块分别处理才解决。另外如果只是想在数据预处理阶段提取patch并不需要梯度可以加上with torch.no_grad():同时把结果转换到CPU或者磁盘存储避免GPU显存和训练模型叠加。Unfold是一个看起来简单、用起来很爽、但细节容易出错的操作我的建议是在自己的项目里先用小张量把索引关系验证一遍再上全尺寸数据这样可以省去大量排查时间。我在实际使用中养成了一个习惯每次构建新的Unfold流程时都会先打印patches.shape再随机取其中一列和原图对应区域做切片对比确认stride和padding都算对了才继续往下写。有一次我把padding方向理解反了导致训练loss怎么都降不下去浪费了一整天最后是用这份对照结果定位到问题的。希望你看完这篇之后能少走我走过的这些弯路。
