1. 这个全卷积掩码自编码器到底想解决什么问题第一次翻到 ConvNeXtV2 的论文时我盯了半天没反应过来一个纯卷积的骨干网络为什么要去做掩码自编码掩码建模这事儿不是 Transformer 的主场吗——把一块块 patch 挖掉喂进编码器让解码器把挖掉的部分补回来这套流程天生适合注意力机制因为注意力可以很自然地把这里有个洞当作一个 token 处理。卷积怎么搞卷积核是滑窗你把某块像素填成 0它和真实的黑色像素、和补零 padding 三者混在一起网络根本分不清。ConvNeXtV2 给出的答案是把这件事从头设计一遍。它由三个咬合在一起的部分组成一套针对全卷积主干重新设计的掩码预训练框架全卷积掩码自编码器、一个叫 GRN 的归一化层全局响应归一化以及配合稀疏卷积的稀疏化前向实现。三者不是并列关系而是因果链——因为要用卷积做掩码建模掩码后的零值会污染卷积核所以需要稀疏卷积因为稀疏输入会让通道响应趋同、特征退化所以需要 GRN因为有了 GRN掩码预训练才真正训得起来微调后才能在 1000 类图像分类这类基准上超过同等规模的监督预训练模型。这篇内容适合三类人看。第一类是正在选型的算法工程师手上有个检测或者分割的下游任务想找一个比 ViT 更省显存、部署更友好的骨干ConvNeXtV2 是很值得纳入候选的。第二类是想复现掩码预训练的研究型同学这套流程的每一步都有坑从掩码比例到损失计算范围错一处 loss 就是平的。第三类是只想抄一个 GRN 模块插到自己网络里的人——说实话这是整篇论文里性价比最高的部分一个二十行的模块几乎零额外开销能实打实缓解通道响应塌缩。我个人认为这篇工作最值得琢磨的地方不在于刷了多少点而在于它示范了一种思路当你想把一个为 A 架构设计的方法搬到 B 架构上时不要只搬流程要搬完之后检查 B 架构的生理结构会不会让流程失效失效了就针对性打补丁。ConvNeXtV2 就是一次完整的移植—诊断—打补丁过程。2. 整体设计思路三个组件为什么必须一起出现2.1 从监督预训练到掩码预训练的路线选择卷积网络的预训练范式在过去几年基本是被监督学习统治的拿一个大的有标注数据集训一遍再迁移到下游。这条路的天花板很明显——标注成本。自监督在 Transformer 阵营已经被掩码建模验证过一轮效果甚至能超过监督预训练那卷积为什么不行早先的尝试主要有两条。一条是把卷积网络和对比学习结合问题是需要精心构造正负样本对对数据增强策略极其敏感训起来调参成本高。另一条是把卷积网络塞进自编码器框架里做像素重建早期的做法是直接把整张图编码再解码不做掩码结果网络学到的多是低频颜色统计学不到有判别力的中高频结构迁移效果一般。ConvNeXtV2 选的是掩码重建这条最难但天花板最高的路。难点在哪难在卷积的归纳偏置和掩码这个操作是冲突的。卷积假设输入在空间上是连续、稠密、邻域相关的掩码人为制造了空洞破坏了这种连续性。Transformer 没有这个假设它靠注意力全局建模位置可以被挖掉而不影响其它位置的计算。所以要让卷积吃下掩码输入必须解决卷积核看不见被挖掉的区域这件事。我踩过一个很典型的坑早期我自己试着在卷积主干上套标准掩码自编码流程就是把被 mask 的 patch 全部置零然后正常前向。训了三十个 epochloss 降到某个值之后完全不动重建出来的图是一片模糊的色块。当时以为是学习率的问题调了半天没用。后来把中间层特征拿出来看才明白——通道之间的激活值几乎一模一样相当于网络把所有通道学成了同一个东西特征维度白白浪费了。这就是后来论文里明确讨论的特征坍塌也是 GRN 要解决的问题。2.2 三件套的分工与因果链把三个组件拆开看各自解决的是不同层级的问题缺一个都跑不通我用一张表把它们的职责和失效后果对应起来。组件解决的痛点缺失后的直接后果额外开销全卷积掩码自编码框架卷积主干无法直接消费掩码输入掩码语义被卷积核污染重建目标失真训练期一次性推理期零成本稀疏卷积前向零值填充与真实像素、边缘 padding 混淆掩码边界处特征被稀释重建发糊依赖稀疏算子实现复杂度上升GRN 全局响应归一化稀疏输入导致通道响应趋同特征坍塌通道利用率骤降微调上不去每个 block 两个 C 维参数可忽略因果链是这样的为了做掩码建模必须处理稀疏输入于是引入稀疏卷积稀疏卷积让特征图在某些位置不存在反向传播时这些位置拿不到梯度通道统计分布被压缩进而出现通道间响应高度相似GRN 在通道维度做一次对比式归一化把主导通道的响应放大、把萎靡通道压下去强行把通道间差异拉开从而恢复特征的表达能力。这里有个设计细节值得说GRN 被放在 block 的升维之后、降维之前也就是在 4 倍宽的高维空间里做归一化。为什么不在原始维度做因为高维空间里通道数多统计量更稳而且归一化的结果马上要被 1×1 卷积压回低维等于让这次归一化直接影响降维卷积的输入分布作用更直接。代价几乎没有因为高维空间的张量只是中间产物不落地到显存里长期驻留。2.3 与 Transformer 版掩码自编码的关键差异很多人第一反应是这不就是卷积版 MAE 吗方向对但差异比想象中大这些差异恰好决定了实现时的选型。第一个差异是没有位置编码。Transformer 版必须给每个 token 加位置编码否则注意力是置换不变的模型不知道谁在哪。卷积天然带位置信息滑窗的局部性本身就是一种强位置先验所以 ConvNeXtV2 里完全不需要位置编码也就省掉了掩码后位置编码怎么处理、变分辨率时位置编码怎么插值这一堆麻烦事。第二个差异是掩码粒度更粗。Transformer 版常用 16×16 的 patch掩码比例 75%ConvNeXtV2 用 32×32 的 patch掩码比例 60%。为什么32×32 的 patch 意味着卷积核的感受野能覆盖到的邻域更大块内信息更完整掩码比例降到 60% 是因为卷积的感受野有限遮得太狠会剩下太多孤岛稀疏卷积处理起来效率低且边界效应严重。第三个差异是编码器和解码器的非对称程度。Transformer 版解码器很轻几个 block 就够ConvNeXtV2 的解码器也是轻量的但要处理的是从稀疏特征图恢复到稠密像素这个更别扭的操作因为它需要在被遮掉的位置上也生成预测而主干在这些位置根本没有输出。实现上通常是在解码器里先把稀疏特征图补零变稠密再用轻量卷积上采样。3. FCMAE把掩码自编码器塞进全卷积主干的具体做法3.1 掩码策略与 patchify 的设计取舍预训练阶段的输入处理和我们平时做分类完全不一样。分类时 stem 通常是 4×4、步长 4 的卷积加一个 LayerNorm把 224×224 的图降到 56×56。预训练时这个 stem 要换掉改成 32×32、步长 32 的卷积把图直接切成 7×7 个 patch每个 patch 对应 32×32×3 的原始像素。我一开始觉得这个设计很怪——为什么不用大 kernel 的卷积做 patch embedding后来自己想明白了掩码自编码的重建目标必须是原始像素级的信息如果 stem 用了 4×4 步长 4那后面所有的下采样都已经把细节丢掉了解码器根本没法恢复出 32×32 的细粒度纹理。所以预训练 stem 必须一次性切到 32×32让每个 token 保留完整的块内像素重建才有可能。掩码的具体操作是在 patch 维度上做的一共 49 个 patch224/32 的平方随机选 60% 遮掉剩下约 20 个 patch 送进网络。注意这里遮掉的含义和 Transformer 版不同——Transformer 是把 token 直接删掉序列变短卷积版因为要保持空间结构不能真的删只能把这些位置标记为无效让稀疏卷积跳过它们。注意掩码比例 60% 是配合 32×32 的块大小定的如果你把块大小改小到 16×16比例要往上调70%~75%否则剩下的可见区域太多任务太简单学不到东西反过来块大小改大到 64比例要降到 50% 左右。这是一个联动参数不要单独调。3.2 稀疏卷积到底解决了什么这是整个流程里最容易被含糊带过的部分但恰恰是能否复现成功的关键。问题分两层。第一层假设你把被 mask 的位置填 0然后正常做 7×7 深度卷积。卷积核在滑过可见区域和遮罩区域的交界处时会把那些 0 当成真实的像素值参与加权求和。结果是可见区域边缘的特征被稀释了——本来一个通道的响应是 1.0靠近边界的输出可能变成 0.5。这会让网络学到一个错误的先验靠近边界的特征就是弱的。而推理时根本没有掩码这个先验就变成了纯粹的噪声。第二层更微妙卷积的零填充本身就是用 0 补边界。mask 填的 0 和 padding 补的 0 在数值上完全一样网络没法区分这里是图像外部和这里是被我遮掉的两者的处理逻辑应该完全不同。稀疏卷积的做法是改变计算的定义域只对有效位置做卷积输出也只落在有效位置上无效位置的输出直接不存在而不是 0。这样卷积核在计算时邻域里只统计真实存在的点被遮掉的位置不参与加权边界处也不会引入假的零值贡献。数学上等价于在稀疏张量上定义一个局部算子实现上需要维护坐标索引。具体实现有两条路我在不同项目里都用过路线一用子流形稀疏卷积算子。需要额外装一个稀疏卷积库构造时把特征和坐标一起打包卷积层会在坐标集合上运算。优点是语义完全正确边界干净缺点是依赖重版本兼容性差我自己就遇到过算子库和框架版本对不上、编译失败半天的经历而且调试起来很痛苦中间张量没法像普通张量那样随手打印。路线二归零近似。正常用稠密卷积但在每次卷积之后把被 mask 的位置重新置零同时保证所有卷积层的 bias 设为 False。这样做的效果是无效位置永远保持严格的零卷积核读到的邻域里只有真实像素和被强制归零的位置边界处依然会有稀释但至少不会把错误信号往远处传播。归零近似的代价要说清楚7×7 深度卷积的感受野半径是 3也就是说在有效区域的边缘 3 个像素内输出特征仍然是欠统计的。对于 32×32 的块来说边缘 3 个像素占比大约 34%不算小。实测下来归零近似能跑到稀疏卷积大概 95% 的效果对于个人复现、小规模验证完全够用如果要追论文数字还是得上真正的稀疏算子。另外有个容易忽略的点LayerNorm 和 GRN 都是沿通道维计算的对每个空间位置独立地跨通道做归一化不会把信息从被遮位置扩散到可见位置所以在归零近似里这两个层是安全的不需要特殊处理。GELU 也一样GELU(0)0零进零出。真正会跨空间混合的只有卷积层只需要盯住卷积就够了。这个观察省了我不少事——一开始我甚至想把 LayerNorm 也改成掩码感知的后来发现完全没必要。3.3 重建目标与损失范围重建目标这一块细节决定成败。核心做法是在计算损失之前对每个 patch 内部的像素做一次标准化减掉这个 patch 的均值、除以标准差。这一步叫逐块归一化。为什么必须做因为如果不做网络会倾向于先学这块区域的平均颜色把大量容量浪费在低频信息上。逐块归一化把每个 patch 的均值方差都抹平了网络被迫去学块内的相对结构也就是纹理和边缘这才是迁移到下游任务时有用的东西。损失只在被遮掉的 patch 上计算可见 patch 不参与。这一点很多人会写错包括我自己第一版。如果全程都算损失那 40% 的可见区域网络可以直接抄输入梯度被这些简单题主导被遮区域的梯度被稀释模型学不到真正的补全能力。伪代码大致是这样# 预训练阶段每个 step 的核心逻辑 images images.cuda(non_blockingTrue) # (B, 3, 224, 224) target patchify(images, patch_size32) # (B, 49, 3*32*32) target per_patch_normalize(target) # 每个 patch 内做标准化 x, mask random_masking(target, mask_ratio0.6) # mask: (B, 49) 布尔True 表示被遮 pred model(x, mask) # 只在可见位置做稀疏前向再补全 loss ((pred - target) ** 2)[mask].mean() # 只统计被遮位置 loss.backward() optimizer.step()注意[mask]这个索引它要求 pred 和 target 都被重排成先所有可见、再所有被遮的紧凑顺序否则布尔索引对不上。这个顺序重排是掩码自编码里最容易写错的地方之一我在两个不同项目里都栽过症状是 loss 能降但重建出来的图错位。提示调试顺序重排最简单的方法是用小数据过拟合——只取 4 张图训 200 步看能不能把重建误差压到接近 0 并且重建图肉眼可辨。如果压不下去八成是顺序或者掩码索引错了而不是模型容量不够。4. GRN一个二十行的模块解决通道塌缩4.1 特征坍塌长什么样怎么量化先说清楚现象。掩码预训练跑起来之后如果你把某个 stage 的输出特征拿出来计算通道之间的余弦相似度矩阵正常情况应该是一个接近对角占优的矩阵——不同通道学到不同模式彼此差异明显。发生坍塌时这个矩阵会几乎全是接近 1 的值意味着所有通道的输出几乎线性相关。更直观的一个指标是死亡通道比例统计每个通道在所有空间位置上的标准差如果某个通道的方差接近 0说明它在任何输入下都几乎输出同一个值这个通道就废了。坍塌严重时死亡通道能占到 40% 以上。我在一次失败的预训练里测过中间 stage 的死亡通道比例到了 52%难怪怎么调学习率都上不去。坍塌的成因值得稍微展开。随机掩码意味着每次前向的输入分布差异极大——这次这张图遮的是左上角下次遮的是中间。卷积核要在这种极度不确定的输入分布上工作最安全的策略就是把所有通道都学成接近恒等映射输出一个平滑的、不依赖输入的响应这样任何掩码模式下都不会产生剧烈误差。这是典型的局部最优陷阱loss 确实在缓慢下降但学到的东西没有判别力。Transformer 那边很少出现这个问题因为注意力的 softmax 天然会做竞争性归一化——一部分 token 的权重高另一部分就被压低通道间差异是被机制本身保证的。卷积没有这个机制所以需要人为补一个。4.2 GRN 的计算过程与实现GRN 的思路非常直接在通道维度做一次谁强谁负责的对比归一化。分四步对每个通道计算它在空间维度上的 L2 范数得到一个长度为 C 的向量代表这个通道整体的响应强度用每个通道的范数除以所有通道范数的均值得到归一化系数把原特征乘上这个系数广播到空间维乘一个可学习的缩放 γ 加一个可学习的偏移 β再残差加回原输入。第三步是核心范数大的通道被放大范数小的通道被压制通道之间的相对关系被显式强化。这一步实际上是在做软性的通道竞争。代码不长但有几个细节必须写对class GRN(nn.Module): def __init__(self, dim, eps1e-6): super().__init__() self.eps eps self.gamma nn.Parameter(torch.zeros(1, 1, 1, dim)) self.beta nn.Parameter(torch.zeros(1, 1, 1, dim)) def forward(self, x): # x: (B, H, W, C)channels-last 布局 gx torch.norm(x, p2, dim(1, 2), keepdimTrue) # (B, 1, 1, C) nx gx / (gx.mean(dim-1, keepdimTrue) self.eps) return self.gamma * (x * nx) self.beta x细节一gamma和beta都初始化为 0。这意味着训练开始时 GRN 的输出恒等于输入x整个模块是一个纯恒等映射。这是非常有必要的——预训练早期特征分布极不稳定如果 GRN 一开始就大幅改写特征很容易把训练带偏。随着训练推进γ 从 0 慢慢长起来GRN 的作用是渐进注入的。这个零初始化 残差的模式在 ConvNeXt 系列里反复出现层缩放也是同一个思路核心就是让新增模块在训练初期不添乱。细节二dim(1, 2)对应的是 H 和 W。如果你用的是 channels-first 布局维度要改成dim(2, 3)gamma的形状也要改成(1, C, 1, 1)。这个我踩过直接抄代码不改布局结果归一化在错误的维度上做等于把 batch 和通道混在一起算了效果比不加还差。细节三分母用的是通道范数的均值而不是求和。论文里写的是求和形式实现里用均值两者只差一个常数因子 C而这个因子会被可学习的 γ 完全吸收所以效果等价。但如果你从论文公式直接推导代码会困惑一下。4.3 为什么放在这个位置以及它能带来什么GRN 在 block 里的位置是固定的深度卷积 → LayerNorm → 1×1 升维4 倍→ GELU →GRN→ 1×1 降维 → 层缩放 → 残差。放在 GELU 之后、降维之前。选这个位置的理由我琢磨过一阵。放在 GELU 之后是因为 GELU 会把负半轴压到接近零输出的分布已经比较干净这时候做范数统计更稳定不会被激活前的正负抵消干扰。放在降维之前是因为要在高维空间做通道竞争——4 倍宽的通道里有更多候选者参与竞争筛选出来的组合更丰富如果先降到原始维度再竞争可选空间小了一半以上。至于能不能放在 block 开头、LayerNorm 之前我试过效果差一些因为深度卷积的输出还没经过任何归一化不同通道的数值尺度差异巨大直接算范数会被极值通道主导。所以位置不是随便放的它依赖于前面的归一化把尺度拉齐了。GRN 带来的收益在不同任务上不一样。分类任务上主要是微调后的绝对精度提升因为通道利用率高了等效于模型容量被用满。检测和分割任务上感受更明显因为这些任务依赖细粒度的特征差异通道塌缩会直接抹掉这些差异。还有一个隐性收益是训练稳定性——加了 GRN 之后预训练对学习率的敏感度明显下降我试过在 1e-4 到 2e-4 的区间里变化loss 曲线都还挺稳不加 GRN 的时候超过 1.5e-4 就开始抖。注意GRN 不替代 LayerNorm两者是串联关系。LayerNorm 管的是把每个位置的通道分布拉回标准正态GRN 管的是在通道之间重新分配权重。有一次我图省事把 LayerNorm 换成了 GRN结果训练直接崩因为 GRN 对输入的绝对尺度不做约束前几层数值会爆。5. 动手实现从模块搭建到预训练跑通5.1 环境准备与工程结构先把预期说清楚这套预训练配方在原始工作里是用几十甚至上百张卡、几千的 batch 跑的个人环境不可能复现同等规模。所以我的建议是分成两个层次——模块级和流程级要严格对齐规模级要大幅缩减。先把流程跑通、验证每个组件确实起作用再谈规模。模块级指的是 GRN、block 结构、损失定义、掩码逻辑这些必须写对一处不对就全盘白费。规模级指的是 batch size、图像分辨率、epoch 数、模型宽度这些可以等比缩。我在一张 24G 显存的卡上跑过完整的预训练流程验证用的配置是图像 128×128patch 32×32每个样本只有 16 个 patch模型宽度压到 Base 的四分之一batch 192200 个 epoch。loss 能稳定下降、重建图能看出轮廓、微调到一个小规模分类数据集上比随机初始化高 8 个点以上这就说明流程是对的。目录结构建议这么分后面调参的时候好定位proj/ models/ grn.py # GRN 模块 block.py # ConvNeXtV2 block backbone.py # 四阶段主干 sparse_ops.py # 稀疏卷积的两种实现封装 data/ patchify.py # 切块 逐块归一化 随机掩码 pretrain.py # 预训练主循环 finetune.py # 微调流程 configs/ tiny_pretrain.yaml tiny_finetune.yaml5.2 Block 的完整实现一个 block 里最容易出错的是通道布局的来回切换。为了性能卷积希望用 channels-last但 LayerNorm 和 GRN 的定义都是沿最后一维所以整个 block 内部统一用 channels-last 最省事进出 block 的时候不用转。class ConvNeXtV2Block(nn.Module): def __init__(self, dim, layer_scale_init1e-6, drop_path0.0): super().__init__() self.dw nn.Conv2d(dim, dim, 7, padding3, groupsdim, biasFalse) self.norm LayerNorm(dim, eps1e-6, data_formatchannels_last) self.pw1 nn.Conv2d(dim, 4 * dim, 1, biasTrue) self.act nn.GELU() self.grn GRN(4 * dim) self.pw2 nn.Conv2d(4 * dim, dim, 1, biasTrue) self.gamma nn.Parameter(layer_scale_init * torch.ones(dim)) \ if layer_scale_init 0 else None self.drop_path DropPath(drop_path) if drop_path 0. else nn.Identity() def forward(self, x): inp x # (B, H, W, C) x self.dw(x.permute(0, 3, 1, 2)).permute(0, 2, 3, 1) x self.norm(x) x self.pw1(x.permute(0, 3, 1, 2)) x self.act(x) x self.grn(x.permute(0, 2, 3, 1)).permute(0, 3, 1, 2) x self.pw2(x) x x.permute(0, 2, 3, 1) if self.gamma is not None: x x * self.gamma return inp self.drop_path(x)层缩放的初始值是个值得单独说的参数。原始监督预训练里用 1e-6非常保守目的是让残差分支在训练初期贡献几乎为零。但掩码预训练的场景不一样——预训练本身提供了很强的正则掩码就是一种数据增广而且预训练需要模型快速建立起有效的特征表达1e-6 会让残差分支很长时间内学不动。实践中我会把预训练阶段的初始值放到 1e-4 量级微调阶段再改回 1e-6 重新初始化。这个改动在我自己的小规模实验里带来了大约 0.6 到 1.2 个点的提升不稳定但方向一致。当然这属于超参不同数据规模下最优值会漂移建议在小模型上先扫一遍。还有一个实操细节pw1和pw2的 bias 建议保留为 True而深度卷积的 bias 设为 False因为后面紧跟 LayerNormbias 会被消掉纯属浪费。在归零近似的稀疏实现里pw1、pw2的 bias 要额外注意——如果 bias 非零被 mask 位置的输出就不是严格的 0 了必须在卷积后重新置零。所以稳妥的做法是预训练阶段把所有卷积的 bias 全关掉微调阶段再按常规配置打开。5.3 预训练循环与关键超参超参这块我给出一个可参考的配方同时说明每一项背后的逻辑方便你按自己的算力缩放。参数参考值缩放建议说明输入分辨率224128~160分辨率下降后 patch 数变少掩码比例可略降patch 大小32保持 32不要动改了整套掩码逻辑都要重调掩码比例0.60.5~0.65与 patch 大小联动batch size4096192~512按卡数线性缩学习率同步线性缩学习率1.5e-41e-4~1.2e-4AdamWbetas 取 (0.9, 0.95)权重衰减0.050.05预训练阶段偏大是正常的warmup40 epoch10~20小规模可以缩短但不能取消训练轮数800200~400小数据集上更少就够权重 EMA0.99990.999~0.9999微调从 EMA 权重开始batch 和学习率的线性缩放这条经验很硬我试过 batch 从 512 缩到 128 但学习率不变训练前期直接发散loss 冲到几百然后一路平着走。反过来把学习率缩了但 batch 不缩收敛慢得像蜗牛200 epoch 才到别人 80 epoch 的位置。warmup 千万别省。掩码预训练的前几百步梯度非常大因为重建误差初始很高直接上大学习率会让模型跳进一个很差的局部解而且自己出不来——你会看到 loss 先降后猛涨然后卡死调都调不回来。我一般用 10 到 20 个 epoch 做线性 warmup小规模验证时可以压到 5 个 epoch但至少要覆盖前 2000 个 step。EMA 这块很多人会忽略。预训练的主权重在训练末期会有不小的波动EMA 权重相当于对最近若干步做平滑通常能比主权重高 0.3 到 0.8 个点的下游精度。实现上注意 EMA 的更新要在梯度同步之后做别在optimizer.step()之前否则会把这一步的梯度影响也算进去。另外 EMA 只对浮点参数做不要把整型 buffer 也卷进去。预训练循环主体框架如下重点看损失那几行for epoch in range(total_epochs): model.train() for images in loader: images images.cuda(non_blockingTrue) target patchify(images, 32) # (B, N, 3*32*32) target per_patch_normalize(target) x, mask random_masking(target, 0.6) pred model(x, mask) # 只在可见位置前向 pred unshuffle(pred, mask) # 重排回原 patch 顺序 loss ((pred - target) ** 2)[mask].mean() # 只算被遮的 patch optimizer.zero_grad(set_to_noneTrue) loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() ema.update(model) scheduler.step() if epoch % 20 0: visualize_reconstruction(model, val_sample) # 定期看重建效果clip_grad_norm_那行是我自己加的原始配方里不一定有。理由是掩码预训练早期个别 batch 的梯度范数会突然很大尤其是随机遮到了全是平坦区域的 patch 时重建误差会被放大加个 1.0 的裁剪能明显减少训练中的尖峰。这个操作对最终精度没有负面影响纯粹是稳定性保险。另外强烈建议每 20 个 epoch 做一次重建可视化。这是最直观的诊断手段正常情况应该能看出被遮区域的颜色、大致纹理、物体轮廓如果重建出来是纯色块或者一片噪点别继续跑了先排查问题。我见过太多人盯着 loss 曲线跑了几十个小时最后发现是掩码索引写反了。5.4 显存与吞吐的估算方法这部分给个估算框架因为每个人的卡不一样硬给数字没意义。显存占用主要分四块模型参数、梯度、优化器状态、激活值。以 Base 规模约 89M 参数为例混合精度训练时参数 fp16 占 178MBfp32 主权重 356MBAdamW 的一阶二阶动量各 356MB梯度 178MB加起来大约 1.4GB。这部分是固定的跟 batch 无关。大头在激活值。如果开梯度检查点激活值大概能压到原来的三分之一到四分之一不开的话分辨率 224、batch 256 的 Base 模型激活值能吃到 20GB 以上。所以个人环境上梯度检查点基本是必开的代价是训练速度慢 20% 到 30%。吞吐方面一个粗略的参考是24G 显存的卡128 分辨率、宽度的四分之一版本、batch 192、开梯度检查点大概能跑到每秒 60 到 90 个样本。这个吞吐下200 epoch 在 10 万张图上要跑六七个小时。如果你打算在百万级数据上跑得按周算这时候建议先在小数据上把流程和超参敲定再一次性上大规模别反复试。6. 微调与下游迁移权重怎么接、参数怎么调6.1 stem 替换与权重对接预训练和微调的模型结构只有一处不同stem。预训练是 32×32 步长 32 的卷积加 LayerNorm微调是 4×4 步长 4 的卷积加 LayerNorm后面四个 stage 的结构完全一致。所以权重对接的做法是加载预训练权重忽略掉 stem 的那两个键其余全部直接拷贝。ckpt torch.load(pretrain_ema.pth, map_locationcpu) state ckpt[model] msg model.load_state_dict(state, strictFalse) print(missing:, msg.missing_keys) # 应该只有 stem 相关的 2 个键 print(unexpected:, msg.unexpected_keys) # 应该为空missing_keys里如果出现了 stem 之外的任何键说明你的 stage 划分或者 block 数量跟预训练时不一致这时候别急着训先对齐结构。我遇到过一次是深度卷积的 bias 键原因是预训练时关了 bias、微调时开了这种形状不一致会被load_state_dict静默跳过strictFalse下结果那一层的 bias 一直是随机初始化的值微调精度莫名低两个点查了半天。提示加载完权重之后逐层打印每个参数的mean和std如果某一层全是 0 或者全是一个常数基本可以断定这层没加载上。这个检查比看missing_keys更可靠因为形状不匹配的键在strictFalse下不会报错。微调阶段还有两个约定俗成的调整层缩放的初始值改回 1e-6因为现在是监督微调不需要那么激进的残差以及把所有卷积的 bias 打开。6.2 分辨率迁移与多尺度特征卷积主干在分辨率迁移上比 Transformer 舒服太多这件事值得单独夸一下。因为 LayerNorm 沿通道维、GRN 沿通道维、池化和卷积都在局部空间操作整个网络对输入的空间尺寸没有任何硬约束。224 训完直接喂 288 或者 384前向不会报错精度通常会涨一点因为测试时可以用更高的分辨率提取更细的特征。这一点在部署时很实用——同一个权重端侧用小分辨率、云端用大分辨率不用改模型。但有个前提要注意训练时如果全程固定 224测试时突然上 384性能提升可能不如预期因为归一化的统计量虽然不受影响但深度卷积 7×7 的感受野在图像中的相对比例变了。我一般的做法是微调的最后 10 个 epoch 里混入一部分 288 分辨率的样本让模型适应一下尺度变化测试时再用 288 就能稳定吃到收益。多尺度特征输出也是现成的四个 stage 的输出步长分别是 4、8、16、32直接接到检测或者分割的 neck 上就能用不需要像 Transformer 那样额外设计特征金字塔的接口。这也是 ConvNeXtV2 在很多下游任务里被选中的实际原因之一——迁移成本低。6.3 学习率与分层衰减微调的学习率比预训练低一个数量级一般 5e-5 起步batch 1024 的时候可以到 1e-4。AdamW权重衰减 0.05100 个 epoch余弦退火warmup 5 个 epoch。分层学习率衰减这个技巧在卷积主干上效果比 Transformer 更明显。做法是给不同深度分配不同的学习率乘子浅层给小的乘 0.5 到 0.7深层给接近 1 的。原因是浅层学到的是通用边缘纹理这部分知识在预训练里已经很成熟微调时大改反而破坏深层学的是任务相关的语义组合需要充分适配新任务。我通常按 stage 分组给四个 stage 分别配 0.5、0.65、0.8、1.0 的乘子stem 给 0.3。这个配置在中等规模数据集几万到几十万张上比较稳。如果是极小的数据集几千张更好的策略是直接冻结前两个 stage只训后两个 stage 加分类头等 loss 稳定了再解冻全网络用更小的学习率微调。还有一个反直觉的点微调阶段的数据增强要比预训练弱。掩码预训练本身就是很强的增广模型已经见过各种信息缺失的情况微调时再上重度随机擦除、大比例裁剪容易让模型过度适应残缺输入在完整图像上的表现反而下降。我一般微调时只用随机裁剪加翻转轻微的色彩抖动可选。7. 常见问题与排查实录7.1 症状到原因的速查表踩过的坑整理成表按症状查比较快。症状最可能的原因排查与处理loss 前几十步快降之后彻底平掉掩码比例过低任务太简单提到 0.6或缩小 patch 到 16重建图是纯色块 / 一片模糊目标没做逐块归一化或损失算了全部 patch检查per_patch_normalize是否生效检查损失是否只在 mask 位置重建图能看出轮廓但有明显错位掩码重排逻辑写错预测和目标的 patch 顺序对不上用 4 张图过拟合测试200 步内应能重建清晰中间层通道响应高度相似缺 GRN或 GRN 的 γ 一直没长起来检查 GRN 是否被冻结打印 γ 的绝对值均值应稳步上升预训练 200 epoch 后微调精度和随机初始化差不多用了训练末期的原始权重或 stem 没换换 EMA 权重核对missing_keys只含 stem显存占用远超预期稀疏实现没生效实际走的是稠密前向检查稀疏分支是否被条件判断跳过加打印统计实际有效位置数分布式训练某张卡空转批归一化的跨卡同步开销全网络用 LayerNorm GRN不含任何跨样本统计层微调精度比论文低 2~3 个点layer decay 没配、分辨率没适配按 stage 配分层学习率末尾混入高分辨率样本7.2 三个我亲身踩过、文档里不会写的坑坑一γ 全零初始化导致的假死。GRN 的 γ 零初始化是对的但如果优化器把 GRN 的参数放进了 weight decay 组而 γ 初始是 0那么梯度更新之后立刻被 weight decay 拉回 0 附近永远长不起来。表面看训练正常、loss 也在降但 GRN 实际上完全没起作用等于没加。排查方法是训几个 epoch 后打印gamma.abs().mean()正常的曲线应该从 0 稳步爬升到 0.1 往上。解决也简单把 GRN 的 γ、β 和所有 LayerNorm 的参数从 weight decay 组里排除单独放一组weight_decay0。这个坑我怀疑坑过不少人因为它的症状和GRN 本来就没用一模一样。坑二稀疏卷积的坐标顺序与稀疏张量的内存布局混淆。用子流形稀疏卷积实现的时候坐标张量的顺序必须严格是 (batch, z, y, x) 这类固定格式而特征张量是 (N, C)。我一开始把特征做成了 (C, N) 的通道优先布局算子库不报错但输出的值完全不对——因为它在错误的维度上做了运算。这类问题几乎不可能从 loss 曲线上看出来只能通过对比稀疏前向和稠密前向在无掩码情况下的输出是否完全一致来验证。在引入稀疏算子后第一件事永远是做一次无掩码的一致性对比误差应该在 1e-5 以内。这个习惯我后来一直保留着。坑三EMA 的实现把 buffer 也更新了。有些模型里 LayerNorm 的running_mean之类是 buffer虽然 LayerNorm 本身没有但如果你混用了 BN 或者别的统计层就有。EMA 更新时如果不筛选requires_grad的参数会把 buffer 也做滑动平均导致推理时的归一化统计和训练时对不上精度掉得很突然。正确做法是只对model.parameters()里requires_gradTrue的部分做 EMA或者在保存时只保存参数部分、重建时再补 buffer。7.3 一个诊断特征坍塌的实用小工具最后分享一个我自己写的诊断脚本五分鐘就能判断预训练有没有走上歪路。思路是给指定 stage 的输出挂一个 hook累积若干 batch 的激活值然后算两个指标一是通道间的平均余弦相似度越接近 1 越危险二是死亡通道比例通道方差小于阈值的占比。torch.no_grad() def diagnose(model, loader, stage_idx2, n_batch8): feats [] def hook(m, i, o): feats.append(o.detach().float().mean(dim(2, 3))) # (B, C) h model.stages[stage_idx].register_forward_hook(hook) for i, (x, _) in enumerate(loader): model(x.cuda()) if i n_batch: break h.remove() f torch.cat(feats, 0) # (N, C) f F.normal
阅读完成 · 觉得有帮助?