在 8GB 显存的 RTX 3070 上训练一个有效的水下显著性检测网络,看起来像是一个"戴着镣铐跳舞"的任务。但经过精心设计,SG-JEDNet-Tiny 不仅在资源约束下完成了训练,还在多个水下数据集上取得了与全尺寸模型可比的结果。
256 分辨率输入的理由
水下显著性检测的标准输入通常是 352×352 或 384×384。我们选择 256×256 并非妥协,而是基于对水下图像退化特征的观察:悬浮颗粒和散射效应导致高频细节本身就不可靠——放大分辨率带来的"伪细节"反而可能误导网络。256×256 恰好保留了显著性判断所需的形状和颜色信息,同时将特征图的内存占用降低了约 45%。
AdamW + Cosine Annealing 调参经验
在多个消融实验中,AdamW 配合 cosine annealing 的组合表现出最稳定的收敛曲线。几个关键参数:
- 学习率:初始 lr=1e-3 配合 5 个 epoch 的 linear warmup,峰值后 cosine 衰减到 1e-6。
- Weight Decay:设为 1e-4——这个值来自 grid search 在 {1e-5, 1e-4, 1e-3, 1e-2} 上的对比实验。1e-3 及以上会过度抑制特征调制模块的参数,1e-5 则几乎无效。
- β 参数:保持默认 (0.9, 0.999),但发现 β2 降至 0.99 在水下数据这种高噪声场景下略微提升了稳定性。
Mixed Precision 节省 40% 显存
开启 PyTorch 的自动混合精度(AMP)后,前向传播中的卷积和矩阵乘法自动切换到 FP16。实际显存从 6.8GB 降至 4.1GB——节省了约 40%。更关键的是,AMP 的梯度缩放机制(gradient scaling)在 GFM 模块的梯度调制中工作良好:微小的纹理梯度在 FP16 下不会被 underflow 截断,因为 scaler 会先放大再缩小。
Early Stop Patience=20 的设定依据
这个值不是拍脑袋定的。我们绘制了 200 个 epoch 的验证损失曲线,发现验证损失在 80-120 epoch 之间出现平坦区,但之后还会出现约 1.5% 的 F-measure 提升(130-160 epoch 之间)。Patience=15 有时会过早停止,错过第二波提升;Patience=20 在 95% 的实验中都等到了最优 checkpoint。200 epoch 后验证损失趋于平稳,F-measure 波动在 ±0.3% 以内。
Batch Size=8 在 8GB 显存下的极限
256×256 输入 + batch size=8 的组合在 8GB 显存上已是极限。尝试过 gradient accumulation 来模拟更大 batch(8×4=32),但发现 BN 层的统计量在小 batch 下已经足够稳定——因为 256 分辨率下特征图的 spatial 维度较大,每个样本贡献的统计量已经足够多。增大虚拟 batch 带来的收益几乎为零(F-measure +0.1%),但训练时间增加了 3.8 倍。
SG-JEDNet-Tiny 的经验总结为一句话:不要在硬件约束面前硬刚,要利用约束倒逼设计。256 分辨率的选择、AMP 的应用、patience 的调优,都是"戴着镣铐跳舞"的产物——而最终结果证明,镣铐有时候反而是最好的编舞。