1. 从一次 OOM 说起unfold 返回的到底是什么前阵子帮朋友看一段滑窗推理的代码模型本身不大显存却直接顶到 40G 然后崩了。他把图像切成 512×512 的窗口重叠 128总共也就几百个窗口按理说不至于。扒了半天代码问题出在一行patches img.unfold(2, 512, 384).unfold(3, 512, 384)之后的.contiguous()上——单张 4096×4096 的图展开后中间张量是几千万个元素乘上窗口面积一拷贝就炸。而unfold本身其实一个字节都没多占。这个细节很关键torch.Tensor.unfold返回的是一个视图view不是一份新数据。它在原张量的维度上切出一个滑动窗口的观察角度只改变 size 和 stride底层 storage 完全复用。这一点和 NumPy 的sliding_window_view是同一个思路。理解了这个前提后面的所有用法和坑才有落脚点。1.1 视图不是数据一次拷贝都没发生先看最基本的签名。Tensor.unfold(dimension, size, step)三个参数在dimension这一维上以step为步长取出长度为size的连续片段。返回张量的形状规则很固定——原来这一维的长度 L 变成(L - size) // step 1然后在末尾追加一维长度为size。拿一个 4×4 的二维张量试一下import torch x torch.arange(16).reshape(4, 4) w x.unfold(0, 2, 1) # 在行方向滑窗口高 2步长 1 print(w.shape) # torch.Size([3, 4, 2]) print(w.stride()) # (4, 1, 4) print(w.data_ptr() x.data_ptr()) # True同一块内存w的形状是(3, 4, 2)含义拆开看3是窗口的起始行位置有 3 个0、1、24还是原来的列数最后的2是窗口内部的行偏移。stride 是(4, 1, 4)——注意最后一维的步长是 4正好等于原张量行方向的步长这就是顺着原内存往下走一行的意思。因为是视图对窗口的写入会直接改到原张量w[0, 0, 0] 999 print(x[0, 0]) # tensor(999)这个特性一半是福利后面做重叠相加重建时非常省事一半是陷阱不小心改了窗口源数据就脏了。判断要不要独立性就看你需不需要.clone()。经验上只要这个窗口张量会往下游传、会被别人写就先 clone别赌。1.2 Tensor.unfold 与 F.unfold 长得像作用完全不同这两个东西名字撞脸是新手最容易混淆的地方我见过不止一次有人把F.unfold的输出当成窗口列表去遍历结果维度对不上。torch.nn.functional.unfold也就是常说的im2col是给卷积做矩阵化用的输入必须是 4D 的(N, C, H, W)输出是 3D 的(N, C*kh*kw, L)其中L是所有窗口位置的总数。它会真实拷贝内存而且会把窗口内所有通道、所有像素按固定顺序拉平成一列。它天生支持dilation、padding、stride这几个卷积参数。Tensor.unfold则是通用的滑窗视图对维度没有要求2D、3D、5D 都行但不支持 dilation也不管 padding。对比项Tensor.unfoldF.unfoldim2col是否拷贝内存否纯视图是真实复制输入维度要求任意必须是 4D(N,C,H,W)输出形状原形状 追加窗口维度(N, C*kh*kw, L)支持 dilation否是支持 padding否需自己 pad是典型用途窗口遍历、重叠相加、局部统计卷积底层实现、矩阵化加速选型其实很清楚你要看窗口、要在窗口上做就地操作、要控制输出形状用Tensor.unfold你要把窗口摊成一个大矩阵丢给 GEMM用F.unfold更省事。1.3 视图的三条副作用用视图换来的性能提升是有代价的这三条我基本每次都要在注释里写一遍。第一非连续。unfold出来的张量 stride 是重叠的很多算子虽然能跑但底层会走通用路径而不是优化路径速度可能差好几倍。做matmul、reshape、送进cudnn之前先.contiguous()。第二view会失败。视图张量的 stride 不满足连续条件.view()直接报RuntimeError: view size is not compatible with input tensors size and stride。用.reshape()代替它会在必要时自动拷贝。这个报错非常高频看到基本就能确定是 stride 问题。第三保存会放大。如果你把一个从大张量 unfold 出来的小切片torch.save出去PyTorch 会连整个底层 storage 一起序列化。一个 8×8 的窗口可能拖着几个 G 的原图。要么.clone()后再存要么.contiguous()别偷懒。2. 两次 unfold 的维度落点谁在前、谁在后二维滑窗必须做两次unfold一次管行、一次管列。这里有个非常隐蔽的顺序问题——两次调用谁先谁后结果的窗口内维度顺序是不一样的而且不会报错只会在你后面做矩阵乘法时悄悄算出错误结果。2.1 从 (H, W) 到 (N, C, H, W)窗口维度一直在末尾追加规则只有一条每次unfold都在末尾追加一维并且不改变前面已有维度的索引位置。所以对于(H, W)x torch.randn(8, 8) p x.unfold(0, 3, 1) # (6, 8, 3) - (Ho, W, kh) p p.unfold(1, 3, 1) # (6, 6, 3, 3) - (Ho, Wo, kh, kw)第一次之后维度 1 还是W所以第二次unfold(1, ...)作用在W上完全正确最终形状(Ho, Wo, kh, kw)。扩展到批量和通道索引整体后移两位x torch.randn(2, 3, 32, 32) # (N, C, H, W) p x.unfold(2, 5, 1).unfold(3, 5, 1) # (N, C, Ho, Wo, kh, kw)结果(2, 3, 28, 28, 5, 5)。因为第一次 unfold 只在末尾加了维度W仍然稳稳待在索引 3 上第二次直接对着 3 号位开刀就行。记住这个顺序规律前面的维度表示窗口在哪里后面的维度表示窗口里是什么。做permute重排的时候这个区分决定了你要不要动后两维。2.2 先横后竖 vs 先竖后横一个静默出错的反例现在换一种写法先滑列再滑行x torch.arange(16).reshape(4, 4).float() a x.unfold(0, 2, 1).unfold(1, 2, 1) # (3,3,2,2) b x.unfold(1, 2, 1).unfold(0, 2, 1) # (3,3,2,2) print(a.shape, b.shape) # 都是 (3,3,2,2)形状一模一样但内容不同。a[0,0]是[[0,1],[4,5]]而b[0,0]是[[0,4],[1,5]]——窗口内部的行列被转置了。因为b的后两维含义是(kw, kh)而不是(kh, kw)。这个错误的杀伤力在于如果你后面只是做统计量求和、求均值、求最大值结果完全正确因为转置不改变这些量的值只要一涉及需要保持空间顺序的操作比如卷积权重、位置编码、可视化的 patch 重建结果就全乱。我的做法是永远固定成行在前也就是unfold(2, ...)打头、unfold(3, ...)跟上4D 情形不给第二种写法留机会。2.3 小张量自检五分钟排掉八成的维度错误写滑窗代码之前先花五分钟用一个已知的小张量验证维度顺序比调试两个小时强。probe torch.arange(36).reshape(6, 6) win probe.unfold(0, 3, 2).unfold(1, 3, 2) print(win.shape) # (2, 2, 3, 3)2 个起始行 × 2 个起始列 print(win[0, 0]) # 应该等于 probe[0:3, 0:3] print(win[1, 1]) # 应该等于 probe[2:5, 2:5] assert torch.equal(win[0, 0], probe[0:3, 0:3]) assert torch.equal(win[1, 1], probe[2:5, 2:5])win[0,1]应该等于probe[0:3, 2:5]win[1,0]应该等于probe[2:5, 0:3]。四个角都断言一遍索引和窗口内顺序就锁死了。这套断言我基本会直接留在代码里当单元测试改窗口参数时跑一遍省心。还有一点值得提醒窗口位置索引的 stride 是原张量 stride 乘以 step不是 1。行方向的窗口位置步长是4*28以 6×6 的 int64 为例原 stride 是(6,1)这意味着你在 GPU 上按窗口遍历时访存是跳跃的。如果窗口很小、图很大这种跳跃会明显拖慢速度后面第 5 节会详细说。3. 边缘、整除与 dilationunfold 交不了的那部分活unfold的设计很克制它只负责按给定的尺寸和步长取出合法窗口其余一概不管不补边、不做空洞、不管你能不能覆盖整张图。这三件事都得自己补上。3.1 Ho 的取整公式与丢边窗口数公式Ho (H - kh) // sh 1里的整除是向下去整这意味着末尾不足一个完整窗口的部分会被直接丢掉。举个直观的例子H10, kh3, sh3Ho (10-3)//3 1 21 3三个窗口覆盖0~8第 9 个像素被扔掉最后两个位置7、8、9本来可以放一个窗口但因为只有 2 个像素的余量放不下。丢掉边缘在图像分类里可能无所谓但在滑动窗口推理、拼接重建、语义分割这类任务里就是实打实的 bug——拼回去的图会缺一条边。常见的处理有两种思路补边法先 pad 到(H - kh) % sh 0让最后一个窗口刚好贴边。这是最干净的做法。尾部补窗口算出总窗口数后把最后一个窗口的起点设为H - kh允许前一个窗口与它重叠。这在检测类任务里更常用因为保证每个像素都被完整覆盖。两种方法各有取舍补边法实现简单、窗口大小一致代价是引入的填充像素会参与计算尾部补窗口不引入虚假像素但最后一个窗口的重叠会不均匀做加权融合时要额外处理。3.2 写一个能算 pad 量的通式补边的量其实一推就出来。目标是让(H_pad - kh)能被sh整除也就是H_pad ≡ kh (mod sh)。所以def pad_amount(L, k, s): 返回到能整除所需的补边量 rem (L - k) % s return 0 if rem 0 else s - rem注意这里只在不带 padding 的情况下讨论。如果已经有p的填充先加上2p再算。完整用法import torch.nn.functional as F def window_view_2d(x, kh, kw, sh, sw): H, W x.shape[-2:] pad_h pad_amount(H, kh, sh) pad_w pad_amount(W, kw, sw) if pad_h or pad_w: # 统一右/下补避免位置偏移 x F.pad(x, (0, pad_w, 0, pad_h)) return x.unfold(-2, kh, sh).unfold(-1, kw, sw)这里用-2、-1做负索引好处是 3D、4D 输入都不用改代码unfold支持负维度索引换算会自动处理。补边方向也有讲究默认往右、往下补这样窗口和原图左上角是对齐的重建时裁剪回原尺寸就完事。如果两侧都补坐标映射会多一层偏移除非有特殊理由比如做中心对齐的卷积否则别给自己找麻烦。3.3 dilation 没有参数用 as_strided 补Tensor.unfold没有dilation参数这是它和卷积最明显的差距。如果你要做空洞窗口比如膨胀卷积的感受野分析、稀疏采样的局部注意力只能自己动手。原理不复杂空洞采样本质就是把原张量在行、列方向的步长乘上dilation然后照常取窗口。as_strided可以直接指定每个维度的 stridedef window_view_dilated(x, kh, kw, sh1, sw1, dh1, dw1): H, W x.shape[-2:] Ho (H - dh * (kh - 1) - 1) // sh 1 Wo (W - dw * (kw - 1) - 1) // sw 1 s_h, s_w x.stride(-2), x.stride(-1) new_stride x.stride()[:-2] (s_h * sh, s_w * sw, s_h * dh, s_w * dw) new_shape x.shape[:-2] (Ho, Wo, kh, kw) return x.as_strided(new_shape, new_stride)窗口内部的两个 stride 分别乘上dh、dw就是空洞效果。窗口位置的两个 stride 乘上sh、sw控制滑动。其他前置维度N、C直接沿用。注意as_strided不做任何边界检查尺寸算错不会报错会安静地读到非法内存结果是随机值甚至崩溃。用它之前务必先用小张量跟unfold或者朴素切片对拍一遍确认无误后再上大张量并且优先在 CPU 上验证。我用这套东西做过一次空洞窗口的统计特征提取跟F.unfold配dilation的结果逐元素对比是完全一致的性能上因为是零拷贝比F.unfold快不少——当然前提是你后续操作能接受非连续内存。4. 把窗口用起来卷积验证、重叠推理与 patch 切分前面三节讲的都是怎么正确地把窗口取出来这一节讲取出来之后干什么。我挑三个实际项目里用得最多的场景把代码和踩过的坑都摆出来。4.1 用 unfold 手搓一遍卷积顺便验证 im2col手搓卷积最大的价值不是替代conv2d而是验证你的窗口顺序完全正确。只要手搓的结果和F.conv2d逐元素对得上后面所有滑窗代码都可以放心写。import torch import torch.nn.functional as F def conv2d_via_unfold(x, weight, biasNone, stride1, padding0): N, C, H, W x.shape OC, IC, KH, KW weight.shape assert IC C, 输入通道与卷积核不匹配 if padding 0: x F.pad(x, (padding,) * 4) Hp, Wp x.shape[-2:] OH (Hp - KH) // stride 1 OW (Wp - KW) // stride 1 # (N, C, OH, OW, KH, KW) cols x.unfold(2, KH, stride).unfold(3, KW, stride) # 整理成 (N, OH*OW, C*KH*KW)通道维度必须在窗口维度之前 cols cols.permute(0, 2, 3, 1, 4, 5).reshape(N, OH * OW, C * KH * KW) w weight.reshape(OC, -1) # (OC, C*KH*KW) out cols w.t() # (N, OH*OW, OC) if bias is not None: out out bias return out.transpose(1, 2).reshape(N, OC, OH, OW)这里最容易翻车的是permute(0, 2, 3, 1, 4, 5)这一步。为什么必须把C挪到窗口前面因为weight.reshape(OC, -1)拉平后的顺序是C → KH → KW你的cols最后一维也必须严格按这个顺序排否则乘出来的权重就对错了位置。很多人只对KH、KW敏感忽略通道顺序小通道数时还能碰巧对一旦C 1立刻出问题。验证方式很简单随机参数对拍torch.manual_seed(0) x torch.randn(2, 3, 11, 11) w torch.randn(8, 3, 3, 3) b torch.randn(8) ref F.conv2d(x, w, b, stride2, padding1) mine conv2d_via_unfold(x, w, b, stride2, padding1) print((ref - mine).abs().max().item()) # 应该在 1e-6 量级误差在1e-6以内就算对上了。如果对不上先检查permute顺序再检查padding是不是(left, right, top, bottom)的四元组顺序最后检查输出 reshape 时OH/OW有没有搞反。必须说清楚手搓版在生产环境基本没有优势。F.conv2d背后是 cudnn 的隐式 GEMM 加 Winograd比你先展开再乘要快而且省下那一大块中间矩阵的内存。手搓版的意义在于教学、调试和特殊场景比如需要自定义窗口布局、需要在窗口上做非线性的预处理。4.2 滑动窗口推理重叠相加与权重归一这是unfold最值钱的用法。大图直接送进网络会爆显存切成窗口分别推理再拼回去靠的就是视图的写回能力。思路是这样的win canvas.unfold(...)是视图对win做add_等价于往canvas上的对应位置累加。开两个画布一个装推理结果一个装覆盖次数最后相除就得到重叠区域的加权平均。import torch torch.no_grad() def sliding_infer(img, th, tw, sh, sw, fn): img: (N, C, H, W) fn: 接收 (M, C, th, tw) 的窗口批返回同形状结果 patches img.unfold(2, th, sh).unfold(3, tw, sw) # N,C,nh,nw,th,tw nh, nw patches.shape[2], patches.shape[3] flat patches.permute(0, 2, 3, 1, 4, 5).reshape(-1, *img.shape[1:3], th, tw) pred fn(flat).reshape(img.shape[0], nh, nw, *img.shape[1:3], th, tw) pred pred.permute(0, 3, 1, 2, 4, 5) # 回到 N,C,nh,nw,th,tw canvas torch.zeros_like(img) count torch.zeros_like(img) canvas.unfold(2, th, sh).unfold(3, tw, sw).add_(pred) count.unfold(2, th, sh).unfold(3, tw, sw).add_(1.0) return canvas / count.clamp_min(1.0)几个关键点。装饰器必须是torch.no_grad()。对视图做原地写入会破坏 autograd 的计算图记录一旦canvas参与了梯度追踪add_会直接报 a view of a leaf Variable that requires grad is being used in an in-place operation或者更隐蔽地在反向时报错。推理场景直接关掉梯度最省心。count的类型要和canvas一致。如果img是float16count也应该是float16不然除法时类型提升会凭空多出一份显存。注意窗口覆盖不到边缘的情况。上面的写法里count在没被覆盖的像素上仍然是 0clamp_min(1.0)会把它们变成 0 结果。如果这不可接受就在调用前把图 pad 到刚好被覆盖推理完再裁回来。如果需要更平滑的融合把count.add_(1.0)换成加一个二维 Hann 窗或者高斯核重叠区域中心权重高、边缘权重低能明显减少拼缝。代价是要多存一份权重图做的时候用同样的 unfold 视图累加就行思路完全一样。实测下来这套写法比for 循环切片 逐个写回快 3~5 倍主要省在 Python 层的循环开销和避免了反复的narrow索引。显存上因为是视图累加只多出canvas和count两张图可以接受。4.3 patch 切分从 ViT 视角看 unfold视觉 Transformer 里的 patch embedding本质就是一次不重叠的滑窗切分。常规写法是reshape permutedef patchify(x, p): B, C, H, W x.shape assert H % p 0 and W % p 0 g x.reshape(B, C, H // p, p, W // p, p) g g.permute(0, 2, 4, 1, 3, 5) # B, nh, nw, C, p, p return g.reshape(B, (H // p) * (W // p), C * p * p)这是最快的做法零拷贝reshape在连续输入下是 viewpermute是 view最后一步reshape会拷贝一次。但它要求H、W能被p整除。用unfold的版本对不整除的情况更宽容一些因为你可以先 paddef patchify_unfold(x, p): B, C, H, W x.shape ph (p - H % p) % p pw (p - W % p) % p if ph or pw: x F.pad(x, (0, pw, 0, ph)) # B, C, nh, nw, p, p g x.unfold(2, p, p).unfold(3, p, p) g g.permute(0, 2, 3, 1, 4, 5) return g.reshape(B, -1, C * p * p)注意这里的step等于size等于p所以窗口之间不重叠nh (H - p) // p 1整除时正好等于H // p和reshape版本等价。如果你想做重叠的 patch比如 Swin 里的 shifted window或者 patch 之间带 overlap 的特征提取把step改成小于p的值就行unfold会自动处理窗口数不用手算。还有一个细节值得留意permute之后立刻reshape会触发一次完整拷贝。数据量大的时候比如 batch 64、256 个 patch、每个 patch 196 个元素这次拷贝很可观。如果能接受窗口维度的顺序是(C, p, p)而不是(p, p, C)就可以跳过permute直接用unfold的结果做后续 reshape省一次大拷贝。这属于先想清楚下游要什么形状再决定要不要重排的优化没有通用答案。5. 内存、性能与报错几个必须知道的边界把功能跑通只是第一步真正让人头疼的是那些跑得通但慢得离谱和偶尔崩一次的场景。这一节把我在实际项目里撞过的墙整理一下。5.1 什么时候不该用 unfoldunfold不是万能的。有几种情况绕开它反而更好。做均值池化、最大池化这类规整的局部聚合直接用F.avg_pool2d、F.max_pool2d。它们有专门的 kernel 优化不用展开成窗口再mean。而且池化对边缘的处理ceil_mode、count_include_pad已经有明确语义自己用 unfold 实现要重新处理一遍容易错。需要标准卷积行为用conv2d。手搓版省不下内存——中间矩阵cols的大小是N × OH × OW × C × KH × KW比输入大KH*KW倍小卷积核还行7×7 的核直接把显存吃满。只是想要一次切片没有滑动用narrow或直接索引。unfold的价值在多窗口单次切片它并不比索引快。真正适合unfold的是这几种窗口之间有重叠的场景、需要对窗口做就地写回的场景、窗口布局不是标准矩形网格的场景、以及调试阶段想直观看到每个窗口内容的场景。5.2 内存放大倍数怎么估估算公式其实很简单展开后的元素数是原元素数 × (KH × KW) / (sh × sw)如果sh KH或sw KW有重叠这个倍数就大于KH*KW。举几个实测数据KHKW3, shsw1放大 9 倍KHKW7, shsw2放大49/4 ≈ 12.25倍KHKW512, shsw384前面的场景放大262144/147456 ≈ 1.78倍——看起来不大但基数是一张 4096×4096 的图乘上去就是十几个 G。关键的判断点是你需不需要 contiguous。视图本身不占内存一直保持视图的话无论放大多少倍都免费。一旦你.contiguous()、.reshape()在非连续情况下、或者送进某个要求连续输入的算子内存立刻按上面的倍数膨胀。我的做法是在代码里加一个护栏def check_window_mem(x, kh, kw, sh, sw, budget_gb8.0): elems x.numel() * (kh * kw) / (sh * sw) gb elems * x.element_size() / (1024 ** 3) if gb budget_gb: raise MemoryError(f展开后约 {gb:.2f} GB超过预算 {budget_gb} GB) return gb在写.contiguous()之前调一次比等 OOM 报错直观得多。因为 OOM 报错经常指向一个完全无关的算子因为显存碎片很难定位到真正的元凶。5.3 三个高频报错与定位先说最常见的几个我把报错原文和对应原因做成表方便直接对号入座。报错信息片段根因处理方式view size is not compatible with input tensors size and stride视图非连续.view()不适用改用.reshape()或先.contiguous()maximum size for tensor at dimension X is Y but size is Zunfold的size超过该维长度先F.pad补边或缩小窗口a view of a leaf Variable that requires grad is being used in an in-place operation对需要梯度的视图做原地写加torch.no_grad()或改写成非原地运算step is expected to be greater than 0step传了 0 或负数检查参数步长必须为正输出结果全 0 或明显错位两次unfold顺序反了或permute顺序错用小张量断言对拍maximum size ... but size is这个尤其容易踩。它的触发时机往往是大部分样本都能跑换个尺寸的输入就崩——因为批量里的图片尺寸变了。稳妥的做法是在滑窗函数入口统一做尺寸归一或者对每个样本单独算窗口参数。别指望unfold帮你兜底。另外一个不报错但是很坑的行为当step size时窗口之间会跳空。比如sh5, kh3窗口覆盖0~2、5~7、10~12中间有像素完全不被任何窗口覆盖。如果后面要做重叠相加重建这些位置会保持初始值 0看起来像图像被打了马赛克。这种错不会报任何异常只能靠可视化发现。我的习惯是每次改步长参数后顺手把count图imshow一下看有没有黑条。最后说一个性能上的实测结论。CPU 上用unfold做逐窗口的mean和用F.avg_pool2d相比慢 10 倍以上差距主要在访存模式——unfold的窗口是跨行跳跃的缓存命中率低。GPU 上差距会缩小到 2~3 倍但仍然存在。所以我现在的基本策略是验证、调试、需要自定义布局的时候用unfold功能确认后就换成对应的专用算子。两者算出来的结果对得上切换起来心里有底也顺便给手搓版本做了一次回归测试。有意思的是unfold在torch.compile下的表现比预期好因为编译后窗口遍历会被融合成一个 kernel避免了来回的 Python 开销和中间张量。如果项目已经在用torch.compile那unfold的适用边界可以放宽一些不用急着换成专用算子。