ARTICLE DETAIL

资讯详情

深耕编程入门与网站建设的一线实战洞察。

Retinanet训练前必读:模型准备与预训练权重选择指南

Retinanet训练前必读:模型准备与预训练权重选择指南 1. 前置认知模型准备到底在准备什么说到Retinanet训练自己的数据集很多教程上来就甩命令行然后就是“等训练完看结果”至于中间那个模型文件是怎么来的、为什么要用这个预训练权重、换一个行不行几乎没人讲清楚。我最初跑通整个流程也踩了不少坑回头看最关键的其实是第一步模型准备。这一步没理顺后面跑起来全是玄学。先给还没入门的读者把概念对齐一下Retinanet是一个单阶段目标检测网络2017年由FAIR提出核心卖点是引入了Focal Loss来解决正负样本极度不平衡的问题配上FPN特征金字塔做多尺度检测在速度和精度之间取得了非常理想的平衡。和两阶段的Faster R-CNN相比它不需要RPN区域提议网络这一步而是直接在特征图的每个位置上做密集预测因此推理速度更快在当时COCO榜单上的精度也完全不输两阶段方法。那“模型准备”在这个流程里具体指哪些事情以我自己的实践经验来看至少包括四个方面确定你的任务类型和骨干网络版本准备好代码库和运行环境下载并校验预训练权重理解配置文件里和模型相关的参数、修改迁移学习策略。这四件事看着不起眼但每件都有坑。比如预训练权重选择错误会导致Loss直接不收敛骨干网络输入尺寸设置不对会导致训练到一半显存爆掉。这篇就把我踩过的坑和验证过可行的做法一次讲清楚。需要说明的是这篇是整个系列的第二篇上一篇讲了数据准备和标注格式这一篇专注在“从零把模型跑起来之前的所有准备工作”。按照这个顺序来后面训练那篇你会顺很多。2. 模型结构认知与版本选型2.1 Retinanet为什么值得继续用在写模型准备之前我默认你已经大概了解Retinanet的原理但考虑到可能有跳着读的朋友还是快速过一遍核心设计因为你只有理解了模型的结构逻辑才知道准备模型文件的时候要关注什么。Retinanet的检测器主要由两部分组成主干网络Backbone和检测头Head。主干网络负责提取图片特征通常用ResNet系列配合FPN结构把不同层级的语义信息融合起来。重点在检测头部分Retinanet在FPN的每一层输出上都会接两个并行的子网络一个负责分类一个负责回归框位置。分类子网络输出的是每个anchor属于各个类别的概率回归子网络输出的是anchor和真实框之间的偏移量。和Faster R-CNN不同Retinanet会在特征图的每个位置铺上多个尺寸和比例的anchor算下来一张图会有上万个候选框但其中绝大多数都是背景。这种密集预测方式直接导致正负样本比例严重失衡于是Focal Loss登场了它通过调制因子降低易分类样本的损失权重让模型集中精力去学习那些难分类的少数样本。这个设计放到今天来看依然能打。虽然YOLO系列迭代了很多版本但Retinanet在逻辑上非常简单清晰改动空间大非常适合做一些定制化实验也方便理解单阶段检测器的工作原理。所以即便你已经用过YOLO回来接触Retinanet也不会浪费时间。2.2 选择骨干网络ResNet50还是ResNet101模型准备的第一步其实不是下载文件而是想清楚自己要用哪个Backbone。Retinanet最常用的两个主干是ResNet50和ResNet101其他还有MobileNet、VGG之类的选择但实际项目里ResNet系列占绝对主流。我的建议很直接新手无脑选ResNet50显存不紧张且有精度要求的任务再上ResNet101。原因有两个。第一ResNet50在精度和显存消耗之间更均衡作为默认配置出现在几乎所有开源项目中参考资料最多问题排查最简单。第二ResNet101并非在所有数据集上都有明显优势尤其是你的自定义数据集只有几千张图片时更深的网络反而更容易过拟合训练时间还翻倍。从显存角度来看输入尺寸512x512、Batch Size为2时ResNet50的显存占用大约在5-6GB同样的配置换成ResNet101会到8GB左右。如果你手头只有一张8GB显存的消费级显卡ResNet50几乎是唯一稳妥的选择。还有一个小点很多人会忽略Retinanet实际是可以在推理时改变输入尺寸的因为FPN和检测头都是全卷积结构不依赖固定的输入大小。但你训练时用的尺寸会直接影响anchor的尺度设计所以一旦定下来训练过程中尽量别乱改。2.3 代码库选择官方实现还是第三方复现Retinanet的代码实现有很多版本先把这个选型问题讲了因为不同代码库的模型文件结构、配置方式、存储格式都有差异选错了后续会很难受。目前主流的选择有这么几个第一是Facebook官方Detectron2里的实现这个版本最规范模型文件是.pth格式配置文件用yaml适合正经做研究和项目落地但对环境要求比较高PyTorch版本、CUDA版本都得匹配好新手容易在装环境环节卡住。第二是fizyr/keras-retinanet这个算是Retinanet普及度最高的第三方实现基于Keras框架模型文件是.h5格式训练脚本非常直白对新手极其友好。缺点是Keras的老版本和一些新环境有兼容问题TensorFlow 2.x时代需要踩一些坑。第三是各种基于PyTorch的轻量复现比如Pytorch-Retinanet代码量少读起来不费劲适合想深入理解原理的人。说实话我不太愿意在这篇里指定“唯一推荐”因为这取决于你后续要做什么。但如果你第一次跑通流程我建议选fizyr版本或者一个简洁的PyTorch版本因为它们的配置文件简单日志直观出错容易定位。我自己用的是PyTorch系列后面讲解也以PyTorch版本为主但思路和坑点对其他框架同样适用。有一点必须提醒无论你选择哪个代码库都要先读一遍README和模型加载部分的源码别直接开跑。我就遇到过有人拿着Detectron2的权重文件往fizyr代码里塞两边resnet的state_dict命名方式完全不同直接报错。3. 模型文件准备与代码环境搭建3.1 预训练权重下载哪个、从哪下这是整个模型准备环节的重头戏。Retinanet作为一个检测模型Backbone部分通常使用在ImageNet上预训练好的ResNet权重做初始化这样模型一开始就具备基础的特征提取能力而不是从零开始瞎猜。实践表明使用预训练权重后同样数据量下训练收敛速度会快很多最终精度也更高。那具体下载哪个文件取决于你的代码库如果你用Detectron2官方会直接提供在COCO上训练好的完整Retinanet权重包含backbone和检测头后缀名为.pkl如果你用fizyr版本它的发布页面会提供ResNet50和ResNet101的ImageNet预训练权重格式是.h5还会给出在COCO上训练好的完整模型如果你用某些PyTorch复现一般会让你单独去下载ResNet的state_dict或者直接用torchvision.models.resnet50(pretrainedTrue)来自动加载。这里有一个很重要但常被忽视的问题backbone预训练和完整模型预训练是两个概念。Backbone预训练的意思是只有ResNet那部分用了ImageNet权重检测头是随机初始化的。你在自己的数据集上训练时检测头要从头学。完整模型预训练则是在COCO上训练好的整套网络如果你用这个做初始化检测头也带了COCO的先验知识这在你的数据和COCO差异较大时反而可能起反作用因为anchor的分布和类别语义都不同。所以我的建议是常规情况下只使用Backbone预训练权重就足够了。只有在你的数据集和COCO非常相似比如通用目标检测常见类别时才考虑用COCO完整模型继续微调。下载地址方面各代码库的GitHub Release页面基本都有不要图省事从奇怪的网盘下载很容易拿到被改动过的权重文件尤其是.h5格式出了问题排查起来极其痛苦。我之前就吃过亏下载了一个自称是ResNet50的权重文件训练Loss一直不正常后来逐层对比才发现是有人改了网络结构重新导出的。3.2 权重文件的校验与体积认知知道下载地址之后有必要先把“模型文件到底长什么样”这件事讲清楚因为你拿到一个.h5或者.pth文件如果连里面有多少层、什么结构都不知道后面排查问题就等于瞎猜。以ResNet50为例它的网络结构是一个7x7卷积接BatchNorm和ReLU然后是4个Stage每个Stage分别有3、4、6、3个Bottleneck残差块。最终在ImageNet上训练好的state_dict大概包含100多个key总文件体积一般在90-100MB左右FP32精度。如果你下载的ResNet50权重只有几十MB那几乎可以断定是半精度或者被裁剪过的版本用起来隐患很大。完整版RetinanetResNet50 FPN 两个检测头在COCO上训练好的权重体积大约在130-150MB稍大一些也正常。校验方式也很简单在你的代码环境里加载一次权重把层级结构打印出来快速确认几个关键点第一个卷积层的权重shape是不是(64, 3, 7, 7)最后一个Stage的输出通道是不是2048检测头部分的层是否存在如果只加载Backbone权重这一步没有是正常的。这一步30秒就能完成却能在正式训练前拦截掉80%的“加载模型失败”类报错。我每换一次代码库或权重来源都会做一次这个验证已经成了肌肉记忆。3.3 环境搭建版本匹配是第一优先级模型文件准备好之后接下来要确保代码跑得起来。环境问题在Retinanet训练中出现的概率非常高而且报错五花八门但追根溯源基本都是版本不匹配。PyTorch版本的Retinanet对环境和依赖的要求相对清爽Python 3.8以上PyTorch 1.10到2.x都行配合对应的torchvision再加一个OpenCV用于图像读取和可视化就够了。如果你是TensorFlow/Keras路线就麻烦一点老版本代码往往依赖Keras 2.2.4、TensorFlow 1.14之类的上古版本在新机器上装起来相当折腾。我的建议优先走PyTorch路线。关于CUDA这里分享一个实操经验不要盲目追求最新版CUDA先查你的显卡驱动支持的最高CUDA版本再选择对应的PyTorch安装命令。如果你用的是RTX 30系显卡CUDA 11.x系列就够了不需要非得用12.x。更稳妥的办法是直接用官方提供的Docker镜像省掉一堆依赖冲突。不过如果你是第一次跑我还是建议在本地环境装这样出问题你能看到完整报错也能更深入理解框架的工作方式。用Docker虽然省事但出了问题你连环境变量都可能找不到。4. 配置文件中的模型参数解读与调整4.1 从配置文件看懂模型怎么搭起来的拿到代码库之后不要急着跑训练建议先把配置文件完整读一遍。这个习惯能帮你少走很多弯路。以PyTorch-Retinanet的配置文件为例核心参数通常包括这些backbone选择ResNet50还是ResNet101直接决定了模型体积和显存占用num_classes你的检测目标类别数注意这里是否包含背景类。有些代码库的num_classes要填“类别数1”有些则不需要这个必须看源码确认image_size训练时输入图片的尺寸常见的是512x512或800x800影响显存和检测精度anchor_sizes和anchor_ratiosanchor的尺寸和长宽比设置官方默认值通常是针对COCO数据集调好的自定义数据集建议先沿用默认值batch_size每次迭代送入模型的图片数量受限于显存大小max_iter或epochs训练迭代次数。其中最容易踩坑的就是num_classes。我见过不止一个人在这里填错导致训练能跑通但推理时类别数对不上输出结果全是错的。另外anchor_sizes和anchor_ratios这两个参数在自定义数据集上值得花时间研究。Retinanet在FPN的每一层会设置不同尺度的anchor默认配置适合通用目标检测场景。如果你的检测目标有特殊的长宽比比如检测长条形物体、文本行这种默认参数可能不太合适。但需要提醒的是修改anchor参数是一个牵一发动全身的操作需要你重新统计自己数据集中目标框的尺寸分布再做针对性调整。新手阶段建议先用默认值跑通后期再优化。4.2 迁移学习策略哪些层需要冻结模型文件加载进来后还有一个必须想清楚的问题训练时哪些层要参与更新哪些层要冻结。很多人会走入一个误区认为使用预训练模型就必须把所有层都冻住只训练最后几层。这其实是把迁移学习想简单了。对于Retinanet来说检测头部分分类子网络和回归子网络是必须要训练的因为它们的参数是随机初始化的不训练就没有任何检测能力。问题在于Backbone。如果数据量非常少比如只有几百张且你的任务和ImageNet分类任务在特征层面高度相似那么冻结Backbone前几个Stage只训练最后几个Stage和检测头能有效防止过拟合训练速度也更快。如果数据量中等几千张我的经验是直接全量微调但适当降低学习率。因为Retinanet本身的模型容量不算特别夸张全量微调在数据量尚可的情况下效果最好。如果数据量很大且和预训练分布差异明显那就更简单了全量训练甚至可以考虑把学习率调高一些。实际操作层面冻结层的方式很简单以PyTorch为例把对应层参数的requires_grad设为False。4.3 学习率与优化器的模型侧考量配置模型的时候学习率是和模型结构强相关的。Retinanet的检测头是随机初始化的它的梯度尺度和Backbone的预训练参数梯度尺度有差异这会导致一个常见问题如果学习率设置不当检测头的Loss在训练初期会剧烈波动。这在单阶段检测器里尤其明显因为分类分支要面对成千上万个anchor每个anchor都有梯度贡献随机初始化的检测头一开始根本分不清前景背景产生的梯度信号又大又乱。应对方法有两种。一是设置更低的初始学习率比如1e-4级别让检测头慢慢适应。二是在代码里对Backbone和检测头使用不同的学习率这个做起来也不复杂把参数分成两组Backbone用较小学习率检测头用较大学习率。很多开源代码都支持这种配置你可以直接通过配置项开启。优化器方面我见过用SGD和Adam都有成功的案例。SGD配上动量0.9是经典配置收敛稳定Adam收敛快但最终精度有时不如SGD调到位。如果你追求省心Adam1e-4起步可以快速看到Loss下降的趋势等模型能正常收敛之后再换SGD精调也行。5. 实操中的常见问题与排查技巧5.1 显存不足应该怎么调训练前最让人头疼的报错之一就是CUDA out of memory。这背后其实是一个动态平衡问题显存占用由模型参数、中间特征图、梯度和优化器状态共同决定任何一个环节超了都会爆。给几个排查思路按优先级排第一步检查batch_size是不是太大。这是最直接的变量从2开始尝试逐步减半。如果2都爆再看输入尺寸。第二步检查输入尺寸image_size。Retinanet的显存占用和输入尺寸大致是平方关系——不是线性是平方。从800改成512显存占用可以下降接近五分之二。第三步检查是否开了混合精度。PyTorch的torch.cuda.amp是现在训练检测模型的标配不仅显存占用下降训练速度还更快。大部分人只要用了AMP8GB显存跑512x512的ResNet50完全没问题。还有一个冷门但有效的技巧在配置文件里调整num_workers。训练时数据加载的进程数量开得过大也会额外占用一部分显存。如果num_workers4爆显存而num_workers0正常那就是数据加载的问题。5.2 加载权重时报错Key mismatch这个问题出现的概率极高尤其是你从不同来源下载权重文件时。所谓Key mismatch本质上是权重文件里的参数名和你代码里模型结构生成的参数名对不上。最常见的原因有两个一是Backbone版本不一致。同样是ResNet50有的代码库会用torchvision.models.resnet50自带的结构有的会自己重写ResNet类尽管结构相同但每一层的变量命名方式完全不同。比如一个叫layer1.0.conv1.weight另一个可能叫backbone.body.layer1.0.conv1.weight前缀都不一样。二是类别数不一致。如果你用的权重是在COCO上训练的完整RetinanetCOCO有80个类别输出层的权重shape是(80, ...)而你的数据集只有10个类别检测头的shape就不匹配。这种错位特别容易造成误判让人以为是权重文件坏了。排查方法是打印出权重的所有key再打印出模型的所有state_dict key做一次diff。这个过程虽然繁琐但绝对是值得的。5.3 Loss不收敛或NaN的处理如果模型文件准备环节出了问题训练阶段的症状往往是Loss不收敛、剧烈震荡或者直接变成NaN。出现NaN九成原因是学习率过大梯度爆炸。先把学习率降到1e-5试一下如果还有NaN再检查是不是数据本身的问题比如标注框出现了负坐标、宽高为0的异常框等。这些脏数据会在计算IoU或者Smooth L1 Loss时产生非数值结果。如果Loss能降但在0.5-1附近徘徊不下那可能是类别不平衡问题过于严重Focal Loss也难以处理。这时候要检查你的anchor设置和训练数据看是不是正样本数量太少。一个很基础但有用的统计算一下训练集中每张图片平均有几个目标框。如果这个值小于0.5问题的根源在于数据而不是模型。另外一个容易被忽略的坑数据增强。如果训练代码默认开启了一些激进的数据增强比如随机裁剪、旋转到很大角度在小数据集上也可能导致Loss震荡不收敛。排查时可以先把数据增强全部关掉看Loss是否稳定下降如果是再逐个开启增强策略定位问题。我把这些问题整理成一张速查表方便你对照排查症状可能原因排查方向CUDA out of memory输入尺寸过大、Batch过大、未开AMP减小尺寸、Batch减半、开启混合精度Key mismatch权重来源与代码库不匹配、类别数不同打印权重key逐个diffLoss为NaN学习率过大、数据包含异常标注降低学习率、清洗标注数据Loss不下降学习率过低、正样本过少调高学习率检查anchor配置训练速度极慢未开AMP、数据加载瓶颈开启AMP、调整num_workers推理结果全错num_classes配置不对核对类别数和背景类处理方式5.4 关于预训练权重的一个补充心得最后再分享一个我最近常用的技巧可以在加载预训练权重后冻结Backbone的前几个Stage在自定义数据集上跑约10到20个迭代让检测头先“找到感觉”。之后再解冻全部层用小学习率继续训练。这个策略在数据量不足的情况下特别有效能明显缓解随机初始化检测头在训练初期输出剧烈波动的问题。操作起来就是在代码里截断优化器的参数列表或者手动改requires_grad。10到20个迭代损失的时间很少换来的是训练稳定性的大幅提升这笔账非常划算。回到标题的核心模型准备并不是“下载一个权重文件”这么简单。它是对模型结构、参数初始化、训练策略的一次整体规划。前面这些工作做到位了真正开始训练的时候你就能完全专注于调参和数据分析省去大量为环境、文件、配置问题来回折腾的时间。我自己第一次跑通Retinanet模型准备这步就花了一个多星期但那之后训练过程几乎没出过岔子。把这个过程记录下来希望对正在走这条路的读者有帮助。
返回列表