FasterRCNN源码解析(一)-——跑通代码
这个系列是对哔哩哔哩up主霹雳吧啦Wz所出的FasterRCNN源码解析的视频进行一个记录以及加上自己理解(可能没有多少,更多的是对数据类型怎么变换的进行一个记录),首先学习源码的第一步就是先跑通目标代码
这里附上霹雳吧啦Wz的github链接:https://github.com/WZMIAOMIAO/deep-learning-for-image-processing
课程中的代码都在git中,大家可以自行下载
环境配置
- Python3.6或者3.7
- Pytorch1.6(注意:必须是1.6.0或以上,因为使用官方提供的混合精度训练1.6.0后才支持)
- pycocotools(Linux: pip install pycocotools;Windows:pip install
pycocotools-windows(不需要额外安装vs)) - Ubuntu或Centos(不建议Windows)
- 最好使用GPU训练
- 详细环境配置见原作者github中的requirements.txt文件
文件结构
- ├── backbone: 特征提取网络,可以根据自己的要求选择
- ├── network_files: Faster R-CNN网络(包括Fast R-CNN以及RPN等模块)
- ├── train_utils: 训练验证相关模块(包括cocotools)
- ├── my_dataset.py: 自定义dataset用于读取VOC数据集
- ├── train_mobilenet.py: 以MobileNetV2做为backbone进行训练
- ├── train_resnet50_fpn.py: 以resnet50+FPN做为backbone进行训练
- ├── train_multi_GPU.py: 针对使用多GPU的用户使用
- ├── predict.py: 简易的预测脚本,使用训练好的权重进行预测测试
- ├── valisation.py: 利用训练好的权重验证/测试数据的COCO指标,并生成record_mAP.txt文件
- └── pascal_voc_classes.json: pascal_voc标签文件
跑通代码
作者在视频中跑的是mobilnet模型,这里我们尝试跑一下res50+fpn的模型
1. create_model
这个类是定义模型的部分。
这里需要注意的是
backbone = resnet50_fpn_backbone()
会自动的冻结部分底层权重
代码如下(示例):
def create_model(num_classes):
backbone = resnet50_fpn_backbone()
# 训练自己数据集时不要修改这里的91,修改的是传入的num_classes参数
model = FasterRCNN(backbone=backbone, num_classes=91)
# 载入预训练模型权重
# https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth
weights_dict = torch.load("./backbone/fasterrcnn_resnet50_fpn_coco.pth")
missing_keys, unexpected_keys = model.load_state_dict(weights_dict, strict=False)
if len(missing_keys) != 0 or len(unexpected_keys) != 0:
print("missing_keys: ", missing_keys)
print("unexpected_keys: ", unexpected_keys)
# get number of input features for the classifier
in_features = model.roi_heads.box_predictor.cls_score.in_features
# replace the pre-trained head with a new one
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)
return model
2. main
这就是训练的主干部分了,
主要步骤有:
- 声明本次训练是在GPU上还有在CPU上进行
- 定义
data_transform
图像预处理函数 - 获取voc数据集路径
- 获取数据集
- 对数据集(验证集和训练集)进行载入操作
- 然后实例化训练模型
- 定义优化函数
- 定义学习率及其衰减策略
- 对每个epoch开始进行训练(并对训练权重进行保存)
- 绘制 损失值 和 学习率 曲线
- 绘制mAP曲线
def main(parser_data):
device = torch.device(parser_data.device if torch.cuda.is_available() else "cpu"