PyTorch二维滑动窗口实现：unfold、as_strided与性能避坑指南

前阵子做图像分块推理我又一次被 PyTorch 的滑动窗口坑到了后半夜。需求说起来特别简单把一张 512×512 的特征图按 7×7 的窗口、步长 1 切开对每个窗口做一次局部归一化再拼回去。我一开始写了个双重 for 循环跑一次要十几秒改完用unfold之后掉到几十毫秒中间还踩了视图不连续、梯度对不上、内存翻十几倍这几个坑。所以这篇就把 torch 实现 tensor 二维数据的滑动窗口 这件事掰开揉碎讲一遍从底层 stride 到底是什么到三条实现路线怎么选再到完整的可复用代码和踩坑速查表。滑动窗口这个操作本质上回答的是一个问题——我想看数据的局部。卷积是它、池化是它、图像分块推理是它、时序特征构造是它、局部注意力窗口还是它。会用for循环切片的人不少但知道unfold返回的是视图、知道as_strided的 stride 该怎么算、知道窗口重叠时反向传播的梯度会累加的人就不多了。这篇适合两类人一类是写过卷积但没手撸过滑窗的想搞清楚 PyTorch 在底下干了什么另一类是已经在用F.unfold但遇到性能或者内存问题想找替代方案的。代码我全部跑过参数推导过程也一并写出来你可以直接抄。1. 先把需求说清楚二维滑窗到底在切什么很多人觉得滑窗是个实现细节其实它的语义比想象中丰富。同样是切窗口不同场景对窗口的要求完全不一样选错实现路线后面的性能和正确性都会跟着塌。所以在动手写代码之前我习惯先把需求归类确定自己到底要的是哪一种窗口。1.1 三种典型场景三种不同的窗口语义第一种是稠密重叠窗口最典型的就是卷积。窗口在 H 和 W 两个方向上逐像素滑动步长通常为 1相邻窗口重叠率极高。这种场景的诉求是我要一次拿到所有窗口然后做矩阵运算。特点是展开后的元素数量是原图的 k×k 倍内存压力最大。第二种是分块切分典型场景是大图分块推理、ViT 的 patch embedding。窗口之间通常不重叠步长等于窗口大小。这种场景的诉求是切完之后还要能原样拼回去所以除了切你还得考虑怎么对齐、怎么处理边界余数。第三种是局部聚合比如滑动平均、滑动方差、局部最大值。这种场景其实你根本不需要真的把窗口展开用avg_pool2d、max_pool2d就能一步到位速度快一个数量级。很多人上来就unfold然后mean结果是能算对但慢得离谱。把需求归到这三类里路线基本就定了。我见过最常见的错误就是把第三种硬生生做成第一种代码写得很通用跑起来却比专用算子慢十倍。1.2 为什么 for 循环切片不能用先看一个反面教材也是我最初的写法# 千万别在正式代码里这么写 def naive_windows(x, k, s): H, W x.shape out [] for i in range(0, H - k 1, s): row [] for j in range(0, W - k 1, s): row.append(x[i:ik, j:jk]) out.append(torch.stack(row)) return torch.stack(out)这段代码逻辑上没错但问题有三个。第一Python 层面循环了几万次每次切片都要创建新的 Tensor 对象开销全在解释器和内存分配上GPU 根本喂不饱。第二torch.stack会真正复制数据窗口不重叠还好一旦重叠内存直接按重叠倍数翻。第三它拿不到任何向量化收益你在 GPU 上跑和在 CPU 上跑的耗时差不多等于白买了卡。正确的做法是让 PyTorch 在 C 层面一次性把窗口索引算出来。而这个索引的学名就叫 stride。2. 核心原理stride 才是滑窗的真正底层PyTorch 的 Tensor 是一块连续内存加上一组如何解释这块内存的元数据。这组元数据就是 shape、stride、storage_offset 三件套。理解了这一点滑窗就不再是切片而是换一种方式解释同一块内存。2.1 从一维 unfold 讲起Tensor.unfold(dimension, size, step)这个 API 的作用是在指定维度上按size大小、step步长滑动把每个窗口里的元素作为新维度追加到末尾。它返回的是视图不复制任何数据。import torch x torch.arange(10) print(x.unfold(0, 4, 2).shape) # torch.Size([4, 4]) print(x.unfold(0, 4, 2)) # tensor([[0, 1, 2, 3], # [2, 3, 4, 5], # [4, 5, 6, 7], # [6, 7, 8, 9]])一维长度 10窗口 4步长 2窗口数量是(10 - 4) // 2 1 4最后一个窗口起点是 6覆盖到 9刚好贴着边界。如果长度是 11那么窗口数还是 4最后一个窗口只覆盖到 9元素 10 就被丢掉了——这个细节后面会专门讲。用生活化的说法原来的张量像是一排连续的座位unfold相当于发给你一把取号尺从第 0 号开始连续取 4 个座位算一组然后往前挪 2 个再取。座位本身没变只是取法变了。2.2 二维窗口就是两次 unfold 的链式调用二维滑窗不需要新 API把一维的 unfold 用两次就行。这里有一个非常容易搞错的点第二次调用时维度索引要跟着变。因为unfold把新维度追加到末尾第一次 unfold 之后原来的最后一维W位置就从-1变成了-2。所以正确写法是x torch.arange(36).reshape(6, 6) # 第一次在 H 维dim-2上滑窗 y x.unfold(-2, 3, 1) # (4, 6, 3) - 原 H 变成 4W 还是 6末尾多一个 3 # 第二次此时 W 在 dim-2 的位置上 z y.unfold(-2, 3, 1) # (4, 4, 3, 3)最终形状是(nH, nW, kh, kw)含义是第几个行窗口、第几个列窗口、窗口内第几行、窗口内第几列。这个维度顺序我很喜欢因为它把空间索引和窗口内索引分得很清楚做后续的mean、softmax都很直观。如果写成y.unfold(-1, 3, 1)你就是在窗口内部再滑一次结果形状变成(4, 6, 3, 1)完全不是你要的东西而且不会报错。这类错误最难查因为 shape 看着差不多对。2.3 as_strided一步到位的 stride 推导两次unfold虽然好用但有个小毛病它只支持在已有的维度上操作遇到 batch、channel 都得手动加循环。而as_strided可以直接把整套 stride 摊开写一步到位。推导思路很简单。假设输入x形状是(N, C, H, W)连续内存下 stride 是(C*H*W, H*W, W, 1)。我们希望输出形状是(N, C, nH, nW, kh, kw)那么每一维的 stride 应该是多少输出维度含义需要的 stride推导理由NbatchC*H*W跟原来一样跳一整张图C通道H*W跟原来一样跳一整个通道nH行窗口序号sh * W窗口往下挪sh行nW列窗口序号sw * 1窗口往右挪sw列kh窗口内行W窗口内逐行走跳一整行kw窗口内列1窗口内逐列走跳一个元素写成代码def sw_as_strided(x, k, s): if isinstance(k, int): k (k, k) if isinstance(s, int): s (s, s) kh, kw k sh, sw s H, W x.shape[-2], x.shape[-1] if H kh or W kw: raise ValueError(f窗口 {k} 大于输入尺寸 {(H, W)}) nH (H - kh) // sh 1 nW (W - kw) // sw 1 sH, sW x.stride(-2), x.stride(-1) shape x.shape[:-2] (nH, nW, kh, kw) strides x.stride()[:-2] (sh * sH, sw * sW, sH, sW) return x.as_strided(shape, strides, storage_offsetx.storage_offset())这里有两个细节值得单独拎出来。一是 stride 用的是x.stride(-2)而不是硬编码W这样即使输入不是标准连续布局比如是从某个更大的张量里切出来的算出来的索引依然正确。二是storage_offset显式传进去虽然默认行为就是继承输入的 offset但写出来更保险也方便后来人阅读。注意as_strided是 PyTorch 里少数几个不做越界检查的 API。stride 算错不会报错会静默读到别的数据甚至读到未分配的显存导致进程崩溃。所以这个函数写完一定要用unfold的结果做逐元素比对别凭感觉相信它。3. 三条实现路线对比与选型到这一步手里的选项其实有三个链式unfold、F.unfold、as_strided。它们都能实现二维滑窗但适用场景和坑点完全不同。下面逐个说清楚最后给一张选型表。3.1 路线A链式 Tensor.unfold这是我最常用的默认方案因为它宽容度最高。任意维度的张量都能用不用先 reshape 成 4D也不用管是不是连续内存。def sw_unfold(x, k, s): if isinstance(k, int): k (k, k) if isinstance(s, int): s (s, s) kh, kw k sh, sw s return x.unfold(-2, kh, sh).unfold(-2, kw, sw)对(H, W)输入返回(nH, nW, kh, kw)对(N, C, H, W)输入返回(N, C, nH, nW, kh, kw)对(B, T, H, W)这种奇形怪状的输入也能直接用前缀维度原样保留。这一点在写通用工具函数时特别省心。它的缺点是不支持 dilation空洞窗口。如果你要做空洞卷积那种稀疏采样Tensor.unfold做不到得自己用索引或者走as_strided。3.2 路线BF.unfoldtorch.nn.functional.unfold是给卷积做 im2col 用的输入必须是严格的 4D(N, C, H, W)输出是(N, C*kh*kw, L)其中L nH * nW。它最大的好处是原生支持padding和dilation而且输出的内存布局就是卷积核展开后的标准布局可以直接喂给matmul。import torch.nn.functional as F def sw_funfold(x, k, s1, p0, d1): # x: (N, C, H, W) N, C, H, W x.shape kh, kw (k, k) if isinstance(k, int) else k sh, sw (s, s) if isinstance(s, int) else s ph, pw (p, p) if isinstance(p, int) else p dh, dw (d, d) if isinstance(d, int) else d u F.unfold(x, kernel_size(kh, kw), stride(sh, sw), padding(ph, pw), dilation(dh, dw)) # (N, C*kh*kw, L) nH (H 2*ph - dh*(kh-1) - 1) // sh 1 nW (W 2*pw - dw*(kw-1) - 1) // sw 1 u u.reshape(N, C, kh, kw, nH, nW) return u.permute(0, 1, 4, 5, 2, 3) # (N, C, nH, nW, kh, kw)注意permute之后张量就不再连续了后面如果要做view会报错得先contiguous()或者改用reshape。这一点跟路线 A 的坑是一样的。窗口数的计算公式跟无 padding 版本不一样多了2p和d*(k-1)。这个公式背后是有效窗口大小 d*(k-1)1空洞就是在窗口内插入空隙所以有效覆盖长度变长了。3.3 路线Cas_strided 手工构造as_strided是最底层的方案前面已经写过实现。它的优势是能一次性把 batch、channel、窗口全部摊开中间没有reshape和permute的开销输出的视图直接就是目标形状不需要二次重排。它的劣势也很明显不安全、不可读、不支持 dilation除非你自己推导、对内存布局有隐含假设。我在生产代码里只在两种情况下用它一是需要极致性能的热路径二是我需要在窗口上做读写而不是只读。提示如果只是想读unfold永远优先。只有在unfold无法表达你的需求比如需要跨维度组合滑窗时才考虑as_strided。3.4 三条路线的选型对照表维度链式 unfoldF.unfoldas_strided输入维度要求任意≥2 维严格 4D任意≥2 维输出形状(..., nH, nW, kh, kw)(N, C*kh*kw, L)自定义是否视图是否会分配新内存是支持 padding否需外部F.pad是否需外部F.pad支持 dilation否是需自己推导安全性高高低不检查越界适用场景通用工具函数卷积/im2col极致性能、特殊布局关于是否视图这一行很多人低估了它的影响。F.unfold会真的分配一块新内存元素数量是N*C*kh*kw*nH*nW窗口重叠时这个数字会非常吓人。而unfold和as_strided只是改元数据创建成本接近于零——但代价是只要你对视图做任何非原地计算PyTorch 还是得按需物化内存一样会炸。所以视图省内存这个说法只在你不做后续运算时成立。4. 完整实操从零写一个可复用的滑窗工具前面都是铺垫这一节把它落地成一个能直接放进项目里的模块。我会把接口设计、padding 对齐、反向传播验证这三块讲透因为这三块是最容易出问题的地方。4.1 窗口数量与边界对齐的计算先说窗口数量公式。给定长度L、窗口k、步长s窗口数是n (L - k) // s 1这个公式里藏着两个容易被忽略的事实。第一个事实(L-k)不能整除s不代表数据没覆盖完。举个真实例子L224, k7, s2(224-7)217217 % 2 1看着除不尽但n 217//2 1 109最后一个窗口起点是108*2 216覆盖到2167-1 222可原图索引到 223也就是说最后一行 223 没被覆盖。等等我算错了重新算216 7 223覆盖到索引 222确实差一个。所以正确的判据是最后一个窗口的右端点(n-1)*s k是否等于L。等于说明刚好贴边小于说明尾巴被丢了大于是不可能的公式保证了不会越界。第二个事实如果你想让最后一个窗口刚好贴边需要满足(L - k) % s 0。不满足的话有三个选择——减小步长可能改变语义、增大窗口改变语义、或者补 padding 后再切。补多少 padding 是个需要算的问题。设左右各补p要求(L 2p - k) % s 0即2p ≡ (k - L) (mod s)。举个具体例子L224, k8, s3(224-8)216216 % 3 0不用补。再看L225, k8, s3217 % 3 1需要解2p ≡ 2 (mod 3)p1可以补完L227(227-8)219219 % 3 0对齐。注意这里的p是单边补的量实际补完长度是L 2p。如果s是偶数而2p永远是偶数那么奇偶性无法靠 padding 调整这时候就要换思路比如把p拆成不对称的左右补法。def calc_pad_to_fit(L, k, s): 返回需要单边补的像素数保证 (L 2p - k) % s 0 r (L - k) % s if r 0: return 0 # 需要 2p ≡ -r (mod s) for p in range(s): if (2 * p) % s (-r) % s: return p return 0 # s 为偶数且 r 为奇数时可能无解需另想办法这个函数对s为偶数、r为奇数的情况返回 0 是认输这时候你得改用非对称 padding或者接受尾部裁剪。我在实际项目里更倾向后者因为尾部裁剪的语义更清楚不容易引入人为的边界信息。4.2 核心实现一个带 padding 和对齐检查的封装把前面的东西合起来写一个我实际在用的版本。设计上做了三个取舍输入支持任意前缀维度、padding 用外部F.pad而不是塞进unfold、返回前不做contiguous把决定权交给调用方。import torch import torch.nn.functional as F def sliding_window_2d(x, kernel_size, stride1, padding0, align_tailFalse, modeconstant, value0.0): x: (..., H, W) kernel_size: int 或 (kh, kw) stride: int 或 (sh, sw) padding: int 或 (ph, pw) align_tail: True 时自动补 padding使最后一个窗口贴边 返回: (..., nH, nW, kh, kw) 的视图或张量 kh, kw (kernel_size, kernel_size) if isinstance(kernel_size, int) else kernel_size sh, sw (stride, stride) if isinstance(stride, int) else stride ph, pw (padding, padding) if isinstance(padding, int) else padding H, W x.shape[-2], x.shape[-1] if align_tail: ph max(ph, calc_pad_to_fit(H 2 * ph, kh, sh)) pw max(pw, calc_pad_to_fit(W 2 * pw, kw, sw)) if ph or pw: # F.pad 的参数顺序是 (left, right, top, bottom)从最后一维开始 x F.pad(x, (pw, pw, ph, ph), modemode, valuevalue) H2, W2 x.shape[-2], x.shape[-1] if H2 kh or W2 kw: raise ValueError(f窗口 {(kh, kw)} 大于输入 {(H2, W2)}) return x.unfold(-2, kh, sh).unfold(-2, kw, sw)F.pad的参数顺序是个经典坑。(pw, pw, ph, ph)的含义是最后一维左边补 pw、右边补 pw倒数第二维上边补 ph、下边补 ph。我第一次写的时候按(ph, ph, pw, pw)传结果宽高补反了shape 没报错但窗口全乱查了两小时。mode参数值得多说一句。constant补 0语义最干净但会在边界引入人为的暗边reflect镜像反射适合图像任务能减轻边界效应replicate复制边缘像素适合特征图不会引入不存在的值。这三个怎么选取决于你后续处理对边界有多敏感。做归一化的时候补 0 会把边界窗口的均值拉低我一般用replicate。4.3 三种实现的等价性验证代码写完了不能只看 shape得验证数值。我习惯用同一份输入跑三条路线逐元素比对这样能一次性把 stride 算错、维度顺序搞反的问题都抓出来。torch.manual_seed(0) x torch.randn(2, 3, 10, 12) a sw_unfold(x, 3, 2) b sw_as_strided(x, 3, 2) c sw_funfold(x, 3, 2) print(a.shape, b.shape, c.shape) # 都应该是 (2, 3, 4, 5, 3, 3) assert torch.allclose(a, b), as_strided 与 unfold 不一致 assert torch.allclose(a, c), F.unfold 与 unfold 不一致这里我特意选了一个除不尽的尺寸H10, k3, s2(10-3)//2 1 4最后窗口起点是 6覆盖到 8索引 9 被丢了。W12, k3, s2(12-3)//2 1 5最后起点 8覆盖到 10索引 11 被丢了。如果三条路线对尾部的处理不一致这个断言会立刻炸出来。再补一个非连续输入的场景这是最容易翻车的地方x_big torch.randn(2, 3, 20, 20) x_slice x_big[:, :, ::2, ::2] # 非连续stride 是 (1200, 400, 40, 2) a sw_unfold(x_slice, 3, 1) b sw_as_strided(x_slice, 3, 1) assert torch.allclose(a, b), 非连续输入下 as_strided 出错只要as_strided里 stride 是从x.stride()动态取的这种情况就能过。如果硬编码成(C*H*W, H*W, W, 1)这里必挂。这也是我坚持用x.stride(-2)的原因。4.4 反向传播重叠窗口的梯度会累加很多人担心视图操作会不会切断梯度。答案是不会。unfold和as_strided都是可微的梯度会正确地回传到原张量而且——这是重点——重叠区域的梯度会按覆盖次数累加。验证方法很直接。造一个 5×5 的全 1 输入窗口 3×3、步长 1对窗口求和后反传得到的梯度矩阵每个位置的值应该正好等于该位置被多少个窗口覆盖。x torch.ones(5, 5, requires_gradTrue) w sw_as_strided(x, 3, 1) # (3, 3, 3, 3) w.sum().backward() print(x.grad) # tensor([[1., 2., 3., 2., 1.], # [2., 4., 6., 4., 2.], # [3., 6., 9., 6., 3.], # [2., 4., 6., 4., 2.], # [1., 2., 3., 2., 1.]])这个矩阵是个很好的心算校验器中心被 9 个窗口覆盖所以是 9四个角各被 1 个窗口覆盖所以是 1边上的位置按距离递减。如果你拿到的梯度不长得像这样那一定是 stride 或者窗口数算错了。在真实训练里这个性质意味着什么意味着如果你用滑窗做局部归一化然后求和或者平均重叠区域会获得更多梯度模型会更关注重叠多的地方。这未必是你想要的。解决办法是在聚合的时候除以覆盖次数做一次归一化这在 overlap-add 的重建场景里几乎是必须的。x torch.randn(8, 8, requires_gradTrue) w sw_as_strided(x, 3, 1) cover sw_as_strided(torch.ones(8, 8), 3, 1).sum(dim(-2, -1)) # 每个窗口内的有效元素数 # 直接求和会让重叠区权重偏高 loss_naive w.mean(dim(-2, -1)).sum() # 按覆盖次数归一化各位置权重才均衡 loss_norm (w.sum(dim(-2, -1)) / cover).sum()不过要注意cover的张量形状是(nH, nW)除以它的时候要做广播对齐而且像cover这种通过sum算出来的东西本身是可微的最好用torch.no_grad()包一下或者直接构造常量避免污染计算图。5. 踩坑记录与排查速查表这一节全是血泪。我把自己和同事遇到过的、网上问得最多的几类问题整理出来配上报错原文和成因方便你直接对上号。5.1 常见报错与成因对照报错信息根本原因解决方式view size is not compatible with input tensors size and strideunfold结果是视图不连续不能直接view改用reshape或先contiguous()maximum size for tensor at dimension 0 is 5 but size is 7窗口大于输入长度先F.pad或者检查参数step must be greater than zero步长传了 0检查参数步长必须 ≥1more than one element of the written-to tensor refers to a single memory location对unfold视图做原地写先clone()再写CUDA error: an illegal memory accessas_strided的 stride 越界重新推导 stride加断言输出 shape 少了两个维度第二次unfold用错了维度索引两次都用-2view和reshape那个坑我要多说两句。unfold返回的张量 stride 是花式的PyTorch 无法保证它按行优先连续存放所以view会直接拒绝。reshape的处理策略是能不动就不动动不了就复制所以它一定成功但你可能在不知情的情况下触发了一次大复制。如果你的代码在reshape之后显存暴涨八成就是这里。原地写那个报错也很有意思。unfold产生的是重叠视图同一块内存被多个逻辑元素引用PyTorch 没法判断你想改哪一个所以干脆禁止写入。想改就先clone()代价是要复制一份完整的展开数据。如果只是想在某几个窗口上做修改用索引赋值会比全量克隆省得多。5.2 隐性问题视图共享、内存暴涨、静默错误第一类是改了视图原张量也变了。因为unfold是视图理论上改视图会影响原张量但实际上 PyTorch 又禁止你直接改视图所以这个特性表现为半可用状态。我的建议是干脆把视图当只读需要修改就复制。第二类是内存暴涨而且不报 OOM 而是变慢。举个例子1×3×512×512的输入窗口 7×7步长 1nH nW 506展开后的元素数是3 × 506 × 506 × 49 ≈ 3763 万float32 下约 143 MB。看着还行但如果你同时保留了好几层中间结果或者 batch 是 16那就是 2.3 GB 起步。更隐蔽的是很多操作会在内部再复制一份比如permute之后做contiguous、softmax之前做类型提升同一份数据在显存里可能同时存在三四份。提示写滑窗代码时养成一个习惯在关键节点打印torch.cuda.memory_allocated() / 1024**2看看每一步涨了多少。这比事后查 OOM 高效得多。第三类是静默错误最危险。典型症状是 shape 完全正确数值全错而且模型还能勉强训练只是指标上不去。原因通常是维度顺序搞反了(nH, nW, kh, kw)写成了(nH, nW, kw, kh)。当kh kw时shape 检查完全没有区分能力。我现在的做法是在实现里对所有kh ! kw的用例都加一组测试比如用(3, 5)的窗口跑一遍再和手动循环的结果比对。5.3 性能优化能不用滑窗就别用滑窗这句话听起来像废话但真的是我踩了无数坑之后的结论。滑窗展开的本质是把数据复制 k×k 倍只要你是走物化路线性能上限就被卡死了。所以在动手写unfold之前先问自己三个问题。问题一我是不是只想做聚合如果是求和、平均、最大值、方差这类操作avg_pool2d/max_pool2d直接一步到位不需要展开速度通常快一个数量级。滑动方差可以拆成E[x²] - E[x]²两项各用一次avg_pool2d就能算出来注意用 float64 或者 two-pass 方式保证数值稳定不然大数值下会掉精度。问题二我是不是在做卷积如果是而且权重固定直接用F.conv2d。手写 im2col matmul 只有在权重动态生成、或者需要自定义聚合函数时才有意义。标准卷积的 cuDNN 实现已经优化到极致你手写的版本很难打赢。问题三我是不是只需要部分窗口如果后续只用到少数几个窗口那就别全量展开用索引切片按需取。展开是 O(n) 的操作按需取可能只需要 O(1)。我在做局部注意力的时候就用过这招只在注意力的 query 周围取一个小窗口而不是把整个序列展开。# 反例为了取几个窗口把整张图展开了 all_win sw_unfold(x, 7, 1) # (506, 506, 7, 7)白白分配 sel all_win[100:110, 200:210] # 正例只取需要的区域 sub x[100:100106, 200:200106] # 多取 6 是为了覆盖窗口 sel sw_unfold(sub, 7, 1) # (10, 10, 7, 7)这个优化在分块推理里特别有用。大图分块的时候块与块之间要留出窗口的余量这样每块边缘的窗口也能算对而不是每块都从整图重新展开。还有一个技巧是交错物化。如果你确实需要全量展开但后续运算是逐窗口的、可分离的比如先对窗口内做行方向的卷积再做列方向的可以先只展开一个方向算完再展开另一个方向。这样中间结果的大小从k²降到k内存能省一个数量级。这个思路在实现大核卷积的时候很常用。最后分享一个我自己的收尾习惯任何滑窗代码合入之前我都会加三条断言——shape 断言、和unfold基准的数值断言、以及一次小尺寸的梯度断言。这三条过了基本就不会在半夜被训练崩掉的消息叫醒了。