NAFNetLocal_池化窗口一致性

NAFNetLocal 中的池化窗口一致性：为什么训练与推理的聚合尺度必须对齐

主题：在 NAFNetLocal（SCA + TLC / Local_Base）中，让训练期的池化（聚合）窗口 与推理期的池化窗口 保持相近，可以显著提升最终模型效果。

视角：从原理层面解释"为什么会掉点"，并给出可直接落地的调参准则与验证方法。

0. 一句话结论

在 NAFNetLocal 里，通道注意力（SCA）的输出质量几乎完全由"统计量是在多大的空间区域上聚合出来的"决定。

训练时网络学到的是 以训练 crop 边长为聚合窗口 的统计规律；
推理时如果聚合窗口与它相差过大（典型情况：训练 256 crop、推理 1080p/4K 全图），SCA 输入的特征分布发生漂移，注意力权重失真；
表现为 PSNR/SSIM 掉点，且掉点集中在统计偏离全图均值的区域（暗部、强纹理、大块平坦区）。

把训练和推理的池化窗口调到同一量级 ------ 经验上 推理窗口 ≈ 1 ~ 1.5 × 训练 crop 边长 ------ 是这一类结构里投入产出比最高的一档优化：改一个参数，不动训练，不改结构。

1. 结构回顾：SCA 里那一次池化到底在干什么

NAFBlock 的前向（省略 LayerNorm 的 plain 变体）：

复制代码

X₁ = Conv1×1(X)                  # 通道扩张
X₂ = DWConv3×3(X₁)               # 深度可分离空间卷积
X₃ = SimpleGate(X₂)              # 通道对折相乘，替代非线性激活
X₄ = X₃ ⊙ SCA(X₃)                # 通道门控
Y  = X + β · Conv1×1(X₄)         # 残差

SCA（Simplified Channel Attention）本身只有两步：

python 复制代码

SCA(X) = W_{1×1}( P(X) )
# 训练期：P = nn.AdaptiveAvgPool2d(1)   → 输出 C×1×1
# 推理期（NAFNetLocal/TLC）：P = 局部平均池化 → 输出 C×(H/k)×(W/k)，再上采样回 H×W

两个关键观察：

P 是 SCA 唯一的"空间信息汇总"入口。X₄ = X₃ ⊙ SCA(X₃) 是纯乘性门控，门控信号除了 X₃ 自身外，全部信息来自 P(X) 这一次聚合。
W 是 1×1 卷积，线性、无激活 。它的输入分布被 P(X) 的取值分布完全决定。

所以：P 的窗口大小，就是通道注意力唯一的一个"超参数级"自由度，也是唯一一个 train/test 会不一致的操作。

NAFNetLocal 的做法（TLC: Test-time Local Converter）是：推理期把 P 从 GAP 换成 local average pooling，窗口 k 由训练尺寸推导：

复制代码

k = ⌈1.5 × train_size⌉          # NAFNet 官方实现 / NAFSSR 中采用
# 训练 crop 256  → 推理窗口 ≈ 384
# 训练 crop 512  → 推理窗口 ≈ 768

系数 1.5 的物理含义后面第 3 节展开。

2. 不一致从哪里来：patch 训练 vs 全图推理

	训练	推理（全图）
输入尺寸	256×256 / 512×512 crop	1920×1080 / 3840×2160
GAP 聚合区域	256 px 见方的邻域	整幅图
输出描述子	C×1×1（每个 crop 一个）	C×1×1（全图共用一个）
广播倍率	1 → 256（×256）	1 → 1080（×1080，随分辨率变）

差异不是"更全局"这么简单，至少有三层机制在同时起作用。

2.1 统计量分布漂移：动态范围随窗口面积收缩

设某通道特征为空间平稳场，单点方差为 σ²，窗口 A 内等效独立样本数为 N_eff(A)（要扣掉空间相关性，不是像素数）。窗口均值：

复制代码

μ_A = (1/|A|) Σ_{p∈A} X_p
Var[μ_A] ≈ σ² / N_eff(A)
Std[μ_A] ∝ 1 / √N_eff(A)

结论很直接：窗口越大 → N_eff 越大 → μ 的方差越小 → P(X) 的取值被压缩到通道均值附近，动态范围单调收缩。

而 W_{1×1} 是在训练期的取值分布 上拟合出来的线性映射。推理时 P(X) 的动态范围变小了，W·P 的输出系统性偏离训练分布：

注意力权重的对比度被拉低，通道门控趋于"平均化"；
网络好不容易学到的通道选择能力被稀释 → 等价于"注意力没开"。

用一个可测量的量刻画它：

复制代码

r = Std_test[ P(X) ] / Std_train[ P(X) ]

r 越接近 1，CA 越工作在设计点上。全图推理 + GAP 时 r ≪ 1；换成 crop 级别的 local pooling 后 r → 1。

2.2 空间自适应性：1 个描述子 vs H·W/k² 个描述子

GAP 输出 C×1×1 → 整幅图共享一组通道权重。但图像复原任务里不同区域的统计差异极大：天空 vs 建筑纹理、暗部 vs 高光、欠曝区 vs 正常区。

训练时，GAP 对 crop 得到的其实是"这块 256 区域的通道统计"------ 网络隐式地学到了"按局部内容自适应"的行为；
推理时，GAP 得到"全图统计 "，强迫全图共用一套权重 ------ 这是训练期从未出现过的行为模式。

local pooling 输出 C×(H/k)×(W/k)，上采样回 H×W 后，CA 变成 逐位置 的通道注意力。这让推理期行为回到训练期见过的模式：每个位置的统计量，都来自一个 crop 级别的邻域。

这也是为什么掉点往往集中在"统计偏离全图均值"的区域 ------ 那些位置恰恰是最需要偏离平均权重的。

2.3 广播/上采样倍率随分辨率漂移

从 pooled map 恢复到 H×W：

GAP：1 → H，倍率 = H，随输入分辨率变化；
local pooling：H/k → H，倍率 = k，固定，与分辨率无关。

固定倍率意味着：对任意分辨率输入，CA 的空间平滑先验是恒定的。而 GAP 的广播倍率随分辨率走，1080p 和 4K 下网络表现不一致 ------ 对于"一套权重跑多分辨率"的落地场景，这一点是硬伤。

3. 一致性准则（可直接当规范用）

定义三个量：

k_train：训练期 SCA 实际感受的聚合窗口边长（用 GAP 时 = crop 边长 S）
k_test：推理期 SCA 的聚合窗口边长（local pooling kernel，或 GAP 时 = 输入边长 H）
一致性比 ρ = k_test / k_train

ρ 区间	行为	建议
ρ ≪ 1（如 0.25）	窗口比训练 crop 还小，N_eff 不足，统计量噪声大；网络训练期从未见过这么局部的聚合	✗ 掉点
ρ ∈ $1, 1.5$	统计量分布与训练期基本一致，同时保留逐位置自适应	✓ 推荐工作区间
ρ ∈ $1.5, 2$	统计量更稳、多带一点上下文，通常仍优于 GAP	✓ 可接受，需扫描确认
ρ ≫ 1（如 4~8）	退化为近似 GAP，train-test 不一致重新出现	✗ 掉点，且越接近全图越差

关于 1.5 这个系数： 取略大于 1，是因为 (a) 稍大的窗口让每个位置的均值估计更稳（降低估计方差），(b) 给门控多带一点周边上下文，避免注意力在相邻位置抖动。它是一个可靠的起点，不是理论最优 ------ 用第 5 节的扫描实验在你的数据上确认峰值。

核心原则一句话：不是"越全局越好"，也不是"越精细越好"，而是"和训练时看到的一样大"。

4. 落地检查清单

4.1 最常见的坑：`train_size` 与真实 crop 不一致

NAFNetLocal 的推理窗口是从 train_size 推出来的。很多掉点根本不是算法问题，是配置错配：

复制代码

配置里写 train_size = 256  →  k_test ≈ 384
实际 dataloader crop = 512 →  k_train = 512
→ ρ ≈ 0.75，落在欠聚合区

必须三者对齐：

复制代码

dataset 的 gt_size / patch_size
    == network_g 中声明的 train_size
    == 推理时 convert(train_size=(1, 3, S, S)) 传入的 S

排查建议：训练脚本启动时把实际 crop 尺寸打一行日志，和模型 config 里声明的值做一次 assert。

4.2 训练期就直接上 local pooling（根治方案）

TLC 是"训练 GAP、推理 local"的补丁方案，两者仍有残余差距（GAP 是全局的、广播方式也不同）。更彻底的做法是训练期就用同一个 local pooling 算子：

python 复制代码

class LocalSCA(nn.Module):
    def __init__(self, c, window=384):
        super().__init__()
        self.window = window          # 训练 / 推理 共用同一个值
        self.conv = nn.Conv2d(c, c, kernel_size=1, bias=True)

    def forward(self, x):
        h, w = x.shape[-2:]
        k = self.window
        s = F.avg_pool2d(x, k, k, 0)                                   # 非重叠平均池化
        s = F.interpolate(s, size=(h, w), mode='bilinear', align_corners=False)
        return x * self.conv(s)

好处：训练与推理执行完全相同的算子，只有输入分辨率不同，gap 从根上消掉；
代价：训练期多一点显存/计算；需要保证 crop 边长 ≥ window（否则要写退化分支）。

折中方案 ：训练期按概率混合 GAP / local（如各 50%），让 W 对两种聚合尺度都鲁棒，推理固定用 local。这对"推理分辨率不确定"的场景尤其划算。

4.3 多分辨率训练 / 渐进 crop

若推理分辨率跨度大（720p ~ 4K），单一窗口很难全覆盖：

训练期用多尺度 crop（256 / 384 / 512 随机），让 W 见过不同 N_eff 下的统计分布，扩大 ρ 的容忍区间；
或按分辨率分档部署：每一档各自取 k_test，保证 ρ ≈ 1.2。

4.4 推理 tiling 时，窗口要跟着 tile 走

大图切 tile（如 512×512、overlap 64）后，每个 tile 的输入尺寸回到训练 crop 量级，此时要重新算 ρ：

若继续用 1.5 × train_size 的大窗口，当 window > tile_size 时，池化退化成"每个 tile 一个值"，等价于 GAP-on-tile：可接受，但空间自适应性丧失；
建议设 k_test = min(1.5 × train_size, tile_size)，或直接按 tile 尺寸重算 ρ。

4.5 `fast_imp` 近似要最后再加

NAFNet 的 NAFAvgPool2d 提供 fast_imp（可分离/近似实现，加速大窗口池化）。开启后统计量与精确均值有偏差，等价于轻微改变了聚合核形状。

顺序很重要 ：先关掉 fast_imp 把窗口调对、测准 baseline，最后再单独评估 fast_imp 的精度损失。否则你分不清掉点来自窗口错配还是近似误差。

4.6 base_size 模式 vs train_size 模式

NAFAvgPool2d 支持两种推导方式：

train_size 模式：窗口由训练 crop 推出 → 目标是"匹配训练分布"，即本文主张的做法；
base_size 模式：窗口由推理基准尺寸推出 → 适合"推理分辨率固定且已知"的部署。

默认用 train_size 模式。 只有在推理分辨率被硬性固定、且和训练 crop 差距极大时，才考虑 base_size 显式指定。

5. 怎么验证"我调对了"

5.1 最小可复现扫描实验

固定一份已训好的权重（训练 crop = S）；
验证集上把 k_test 从 S/2 扫到 4S，例如 {128, 192, 256, 384, 512, 768, 1024}；
画 PSNR/SSIM vs ρ 曲线 → 预期是单峰曲线，峰值落在 ρ ∈ $1, 1.5$ ；
同时记录 CA 输入的分布统计（Std[P(X)]、均值、W·P 输出范围），验证峰值点与 r ≈ 1 的位置重合。

第 4 步是关键：如果指标峰值与 r ≈ 1 重合，就证明了因果链是"窗口错配 → 分布漂移 → 掉点"，而不是碰巧。

5.2 两个补充 diagnostic

分区指标 ：把验证图切成平坦区 / 纹理区 / 暗区分别算 PSNR。窗口不匹配时，掉点集中在统计偏离全图均值的区域，全图平均指标容易被掩盖，只盯总 PSNR 会误判；
分辨率一致性：同一张图下采样到 1/2、1/4 再测，指标曲线越平，说明 CA 的尺度一致性越好。这条对"一套权重跑多分辨率"的落地场景是必测项。

5.3 主观检查

客观指标之外，重点看两类 artifacts：

大块平坦区（天空、墙面）是否出现色偏 / 整体亮度漂移；
暗部是否出现噪声残留或过平滑 ------ 这些都是通道门控被"平均化"的典型症状。

6. 什么时候不需要折腾这个

收益不明显的情况：

推理输入分辨率与训练 crop 本来就接近（如固定 256 输入的片上 / ISP 直出链路）：GAP 的 ρ 天然 ≈ 1；
SCA 数量很少（0 ~ 4 个）：CA 对整体贡献有限，收益被淹没；
退化在空间上高度均匀（纯高斯噪声、均匀模糊）：全局统计与局部统计差别不大，ρ 的容忍区间更宽。

收益最明显的场景（优先做）：

训练 crop 小（128 / 192）但推理分辨率高（1080p / 4K）------ ρ 差一个数量级；
SCA 堆叠多（NAFNet 的 encoder 深层 + middle block 大量堆叠）；
图像内容统计空间差异大（大光比夜景、天空+建筑混合、ISP 局部 tone mapping 后的图）；
一套权重要覆盖多分辨率。

7. 结论

SCA 里的池化不是一个"顺手写的下采样"。它是通道注意力唯一的统计入口 ，决定了 W 看到的输入分布，也决定了门控是"自适应"还是"平均化"。

训练与推理的聚合窗口一旦错配，会同时引入三个问题：

统计量动态范围漂移 → W 的输入偏离训练分布，注意力对比度被拉低；
空间自适应性丧失 → 全图共用一组通道权重，训练期从未见过这种行为；
广播倍率随分辨率变化 → 多分辨率下表现不一致。

工程上最省事、最有效的动作只有三条：

对齐 train_size 与真实训练 crop ------ 消除配置型错配（回报最高、成本最低）；
推理窗口按 ρ ≈ 1 ~ 1.5 取 ------ 默认 1.5 × crop 是可靠起点，再用一次扫描实验确认峰值；
追求极致时把 local pooling 写进训练 ------ 让训练与推理执行同一个算子，从根上消掉 gap。

一句话记住它：通道注意力好不好，不看它多复杂，看它聚合的那一块地，是不是和训练时看到的一样大。

参考

Chen et al., Simple Baselines for Image Restoration（NAFNet, ECCV 2022）--- NAFBlock / SCA / NAFNetLocal
Chu et al., Improving Image Restoration by Revisiting Global Information Aggregation（TLC, ECCV 2022）--- train-test inconsistency 与 Test-time Local Converter
Chu et al., NAFSSR: Stereo Image Super-Resolution Using NAFNet（CVPR 2022 NTIRE）--- TLSC 应用于 CA，local pooling size 取训练 patch 的 1.5×
官方实现：megvii-research/NAFNet 中 NAFNetLocal / Local_Base / naf_avgpool2d.py

NAFNetLocal_池化窗口一致性

NAFNetLocal 中的池化窗口一致性：为什么训练与推理的聚合尺度必须对齐

0. 一句话结论

1. 结构回顾：SCA 里那一次池化到底在干什么

2. 不一致从哪里来：patch 训练 vs 全图推理

2.1 统计量分布漂移：动态范围随窗口面积收缩

2.2 空间自适应性：1 个描述子 vs H·W/k² 个描述子

2.3 广播/上采样倍率随分辨率漂移

3. 一致性准则（可直接当规范用）

4. 落地检查清单

4.1 最常见的坑：train_size 与真实 crop 不一致

4.2 训练期就直接上 local pooling（根治方案）

4.3 多分辨率训练 / 渐进 crop

4.4 推理 tiling 时，窗口要跟着 tile 走

4.5 fast_imp 近似要最后再加

4.6 base_size 模式 vs train_size 模式

5. 怎么验证"我调对了"

5.1 最小可复现扫描实验

5.2 两个补充 diagnostic

5.3 主观检查

6. 什么时候不需要折腾这个

7. 结论

参考

4.1 最常见的坑：`train_size` 与真实 crop 不一致

4.5 `fast_imp` 近似要最后再加