~johan's blog
返回专辑
DL
·Johan·5 分钟阅读

初始化：训不动时

换激活函数却不改初始化，或最后一层 init 过大，会导致 loss 不降或梯度爆炸——Kaiming/Xavier 和 zero init head 是基本配对方案。

#深度学习#初始化#训练
初始化：训不动时

1. loss 第一步就 NaN，或长期不降

换激活函数（ReLU→GELU）却不改 init，或 classification head init 过大——第一步 loss NaN，或 loss plateau 在 random 水平。PyTorch 默认 nn.Linear/nn.Conv2d 在 1.x+ 已 Kaiming init，但 自定义层、Transformer、新加 head 要显式 init；load pretrained 后不要 re-init backbone，否则预训练权重被覆盖，微调等于 from scratch。

训不动时排查顺序：init 与激活是否配对 → LR 是否过大 → label/data 是否有 bug → 架构是否过深需 warmup。别一上来改架构，init+LR 能解决大半「不收敛」假象。第一步 loss 若在 random guess 附近震荡但不降，常见是 head logits 尺度不对 或 学习率对 head/backbone 未分层；第一步即 NaN 则优先查 init 与 AMP scale。自定义 Module 若忘记 init 子层，会继承 PyTorch 默认，但与父模块激活不匹配时仍 plateau。Conv-BN-ReLU 块中若手动 insert 新 Conv 却未 init，会出现 局部层爆炸而其余正常 的假象，hook 中间层 std 可快速定位。

2. Kaiming 与 Xavier 配对

ReLU/LeakyReLU 用 Kaiming（He），fan_out + nonlinearity='relu'：

python
def init_weights(m):
    if isinstance(m, (nn.Linear, nn.Conv2d)):
        nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
        if m.bias is not None:
            nn.init.zeros_(m.bias)
model.apply(init_weights)

sigmoid/tanh 用 Xavier（Glorot）。GELU/SiLU 网络实践常仍用 Kaiming 或 Xavier，但 Transformer 有专门 recipe：nn.init.normal_(p, std=0.02) 对 linear；embedding 有时更小 std；residual 路径 scale 1/sqrt(2*num_layers) 防 Post-LN 深网爆炸。Bias 通常 zero。BatchNorm weight=1, bias=0 默认 OK，别 random init BN scale。

Kaiming 方差依据 （ReLU）；Xavier 用 。fan_in/fan_out 对 Conv 要乘 kernel 面积。Transformer attention 的 in_proj 与 FFN 两层常统一 std=0.02，是经验 recipe 而非严格推导；换 Pre-LN 深 24 层时仍建议配 residual scale。分类 head 最后一层可用更小 std（如 0.01）或 zero init weight，使 initial logits 接近 0，softmax 不饱和。

3. 微调时的 init 边界

微调：pretrained backbone 小 LR，新 head 随机 init + 大 LR（often 10×）。model.apply(init_fn) 递归所有 submodule——load pretrained 后只对 新 head（和 optionally 新 neck）call init，别对整网 apply。检测：backbone ImageNet 权重，RPN/ROI head Kaiming init。分割：decoder head random init，encoder 可冻结若干 epoch 再解冻。

全零 init 全层 → 对称 dead，梯度无法打破。分类 head bias 设负值（focal loss 场景使 initial loss 接近 uniform）是常见 trick，与 weight init 配合。最后一层 weight std 过大 → 初始 logits 极端，softmax 饱和，gradient vanish。load checkpoint 顺序契约：build model → load pretrained → init 新增模块 only。若用 strict=False load，须 log missing/unexpected keys，避免 head 仍随机而误以为已 load。冻结 backbone 若干 epoch 时，head init 质量决定早期 loss 能否下降；解冻时 LR 应再降一档，防止破坏 pretrained 特征。两阶段微调（先训 head 再 unfreeze top-k blocks）时，第二阶段勿 re-init 已训 head，除非更换 loss 或类别数变化。

4. 第一层 activation 方差检查

init 后跑 forward（无 backward），看第一层 output activation 的 std——应在 量级，非 0 非 。vanish/explode 时查 init 是否与激活匹配、load checkpoint 后是否 duplicate init。Transformer 可看 embedding 后 first block output norm。工具：register_forward_hook 打 log 或单次 batch 统计。

多层网络应扫 若干 depth 的 activation std，不只第一层——深 12 层若中间层 std 逐层衰减到 1e-4，Post-LN 下常见，需 init scale 或 Pre-LN。AMP 下 first step 若 GradScaler scale 骤降，可能是 init 导致 overflow，而非 LR alone。记录 {layer_name, std, max} 到 wandb，换 init recipe 时可对比。微调节点：backbone 第一层 std 应与 pretrained 一致，若 diff 非零说明 re-init 泄漏。

5. 案例：微调误 apply init 抹掉预训练

某 ViT 微调，load 21k 预训练后 model.apply(init_weights) 整网 re-init，100 epoch 仍不如 random ResNet。抽样 backbone 权重与 checkpoint diff 非零。改为只 init 新 classification head，10 epoch 超 pretrain-from-scratch。说明 load 与 init 顺序 是微调契约。

团队 fix：封装 load_pretrained_then_init_head(model, ckpt, init_fn, head_names)，单测断言 backbone 权重与 ckpt max_abs_diff==0。文档写清禁止对 model 全局 apply 的 lint 规则。类似 bug 也出现在 resume 训练误 reset optimizer 与 误 re-init BN——同一类「顺序契约」问题，应进 code review checklist。上线前对 backbone 抽样 ten 层权重 hash，与预训练 manifest 比对，可一键抓 re-init 回归。

6. 验收

  • 第一步 forward 后第一层 activation std 合理；第一步 loss 非 NaN。
  • 微调：backbone 权重与 pretrained checkpoint 一致（抽样 max diff = 0）。
  • 换激活 ReLU↔GELU，固定 seed 对比 100 step loss。
  • head-only init vs 整网 re-init ablation。
  • AMP first step：GradScaler scale 不应 instant 归零。

微调验收增加 10 step loss 斜率：head-only init 应明显低于 random backbone；若平行，查 LR group 与 init 是否同时错误。Detection 场景对 FPN 新增层也应单独 init，而非假设 load_state_dict strict=False 后随机层会自动被 optimizer 覆盖。

初始化是训练可训性的前置条件；Kaiming/Xavier 与激活配对，微调 backbone 不 re-init 是基本契约。训不动先查 init 与 LR，再查 data。新层加入已训网络时，除 init head，还应检查 optimizer param group 是否含新层——漏加则 head 不更新，表现像 init 错误。分布式训练各 rank init 须一致（固定 seed），否则 BN/SyncBN 与 EMA 行为会漂移。ResNet 系 zero_init_residual 对 last BN gamma 置零使初始 block 近似 identity，是 init 与结构耦合 的范例——换 backbone 时连同这类 flag 一并核对。

← 全部文章

johan's blog