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(pretrained=True)来自动加载。
这里有一个很重要但常被忽视的问题: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,那几乎可以断定是半精度或者被裁剪过的版本,用起来隐患很大。
完整版Retinanet(ResNet50 + 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_ratios:anchor的尺寸和长宽比设置,官方默认值通常是针对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调到位。如果你追求省心,Adam+1e-4起步可以快速看到Loss下降的趋势,等模型能正常收敛之后再换SGD精调也行。
5. 实操中的常见问题与排查技巧
5.1 显存不足应该怎么调
训练前最让人头疼的报错之一就是CUDA out of memory。这背后其实是一个动态平衡问题:显存占用由模型参数、中间特征图、梯度和优化器状态共同决定,任何一个环节超了都会爆。
给几个排查思路,按优先级排:
第一步,检查batch_size是不是太大。这是最直接的变量,从2开始尝试,逐步减半。如果2都爆,再看输入尺寸。
第二步,检查输入尺寸image_size。Retinanet的显存占用和输入尺寸大致是平方关系——不是线性,是平方。从800改成512,显存占用可以下降接近五分之二。
第三步,检查是否开了混合精度。PyTorch的torch.cuda.amp是现在训练检测模型的标配,不仅显存占用下降,训练速度还更快。大部分人只要用了AMP,8GB显存跑512x512的ResNet50完全没问题。
还有一个冷门但有效的技巧:在配置文件里调整num_workers。训练时数据加载的进程数量开得过大,也会额外占用一部分显存。如果num_workers=4爆显存而num_workers=0正常,那就是数据加载的问题。
5.2 加载权重时报错Key mismatch
这个问题出现的概率极高,尤其是你从不同来源下载权重文件时。所谓Key mismatch,本质上是权重文件里的参数名和你代码里模型结构生成的参数名对不上。
最常见的原因有两个:
一是Backbone版本不一致。同样是ResNet50,有的代码库会用torchvision.models.resnet50自带的结构,有的会自己重写ResNet类,尽管结构相同,但每一层的变量命名方式完全不同。比如一个叫layer1.0.conv1.weight,另一个可能叫backbone.body.layer1.0.conv1.weight,前缀都不一样。
二是类别数不一致。如果你用的权重是在COCO上训练的完整Retinanet,COCO有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逐个diff |
| Loss为NaN | 学习率过大、数据包含异常标注 | 降低学习率、清洗标注数据 |
| Loss不下降 | 学习率过低、正样本过少 | 调高学习率,检查anchor配置 |
| 训练速度极慢 | 未开AMP、数据加载瓶颈 | 开启AMP、调整num_workers |
| 推理结果全错 | num_classes配置不对 | 核对类别数和背景类处理方式 |
5.4 关于预训练权重的一个补充心得
最后再分享一个我最近常用的技巧:可以在加载预训练权重后,冻结Backbone的前几个Stage,在自定义数据集上跑约10到20个迭代,让检测头先“找到感觉”。之后再解冻全部层,用小学习率继续训练。
这个策略在数据量不足的情况下特别有效,能明显缓解随机初始化检测头在训练初期输出剧烈波动的问题。操作起来就是在代码里截断优化器的参数列表,或者手动改requires_grad。10到20个迭代损失的时间很少,换来的是训练稳定性的大幅提升,这笔账非常划算。
回到标题的核心:模型准备并不是“下载一个权重文件”这么简单。它是对模型结构、参数初始化、训练策略的一次整体规划。前面这些工作做到位了,真正开始训练的时候你就能完全专注于调参和数据分析,省去大量为环境、文件、配置问题来回折腾的时间。我自己第一次跑通Retinanet,模型准备这步就花了一个多星期,但那之后训练过程几乎没出过岔子。把这个过程记录下来,希望对正在走这条路的读者有帮助。