FaterRCNN代码学习总结1

1.项目结构

backbone为FaterRCNN主干网络部分的代码，包括其权重，此开源项目提供了VGG和resnet50两种，并提供了对应的权重文件，代码的作者认为backbone部分已经在coco数据集上训练完善了，因此在新的训练中可以直接调用coco数据集的权重文件作为权重初始化。 data为自己创建的COCO数据格式的月球陨石坑数据集 network_files为除backbone以外的结构部分，如roi_head、rpn_function和boxes、transform等对预测框和图像进行处理转换的脚本。 save_weights为训练过程每一轮权重的存放目录 train_utils是一个自定义的Python工具模块，主要包含训练过程中常用的辅助函数和工具类

2.train文件

2.1模块导入

[lang:Python] import
1
2
3
4
5
6
7
8
9
10
11
import os
import datetime
import torch
from tqdm import tqdm  # 进度条显示库
import transforms  # 数据预处理变换
from network_files import FasterRCNN, FastRCNNPredictor  # 模型结构
from backbone import resnet50_fpn_backbone  # 骨干网络
from my_dataset import VOCDataSet  # VOC数据集读取(未使用)
from datasets import CocoDataset  # COCO格式数据集读取
from train_utils import GroupedBatchSampler, create_aspect_ratio_groups  # 数据采样工具
from train_utils import train_eval_utils as utils  # 训练评估工具

2.2模型创建

主要函数为def create_model(num_classes, load_pretrain_weights=True): 在该函数中调用backbone文件夹中的resnet50_fpn作为backbone ，之后使用FasterRCNN创建完整的model，导入预训练权重，更改roi_head中的分类头以适应自定义类别数，最后返回model（固定套路）。
[lang:Python] create_model
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
def create_model(num_classes, load_pretrain_weights=True):
 
    backbone = resnet50_fpn_backbone(pretrain_path="/disk527/Commondisk/a804_qkf/vscodeproject/code/faster_rcnn/backbone/resnet50.pth",
                                     norm_layer=torch.nn.BatchNorm2d,
                                     trainable_layers=3)
    # 训练自己数据集时不要修改这里的91，修改的是传入的num_classes参数
    model = FasterRCNN(backbone=backbone, num_classes=91)
    if load_pretrain_weights:
        # 载入预训练模型权重
        # https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth
        weights_dict = torch.load("/disk527/Commondisk/a804_qkf/vscodeproject/code/faster_rcnn/backbone/fasterrcnn_resnet50_fpn_coco.pth", map_location='cpu')
        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.3主函数流程

1.device 设备设置
2.data_transform = 数据增强设置
3.文件初始化
4.train_datasets、val_datasets加载，通过dataset进行数据集加载
5.train_data_loader、val_data_loader进行批次划分加载，使用函数torch.utils.data.DataLoader
6.创建模型 model = create_model(num_classes=args.num_classes + 1)
7.model.to(device)
8.优化器设置，用到函数：torch.optim.选择需要的迭代器
9.学习率设置，用到函数：torch.optim.lr_scheduler.StepLR
10.恢复训练设置（接着上次训练结果继续训练，可有可无）

[lang:Python] train
1
2
3
4
5
6
7
8
9
if args.resume != "":
        checkpoint = torch.load(args.resume, map_location='cpu')
        model.load_state_dict(checkpoint['model'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        lr_scheduler.load_state_dict(checkpoint['lr_scheduler'])
        args.start_epoch = checkpoint['epoch'] + 1
        if args.amp and "scaler" in checkpoint:
            scaler.load_state_dict(checkpoint["scaler"])
        print("从epoch{}继续训练...".format(args.start_epoch))

11.训练循环，主要针对train_loss、learning_rate、val_map进行更新，在循环前提前创建空列表，在循环中通过append进行记录，使用optimizer.step()进行参数更新

for epoch in tqdm(range(args.start_epoch, args.epochs), desc=”总进度”):
训练一个epoch： 使用函数utils.train_one_epoch（）；
记录指标： .appeend() 主要针对train_loss、learning_rate、val_map进行记录；
参数更新： lr_scheduler.step()；
在验证集上评估：utils.evaluate（）；
写入结果： with open；
保存模型权重： 权重文件要用到torch.save函数，在调用权重文件时，对应的用torch.load函数，这两个函数相关联；
循环结束
循环结束后，对训练过程的map loss 进行绘制

[lang:Python] plot_loss
1
2
3
4
5
6
7
8
9
# plot loss and lr curve
if len(train_loss) != 0 and len(learning_rate) != 0:
    from plot_curve import plot_loss_and_lr
    plot_loss_and_lr(train_loss, learning_rate)
 
# plot mAP curve
if len(val_map) != 0:
    from plot_curve import plot_map
    plot_map(val_map)

2.4 train_one_epoch函数

rain_one_epoch函数并不位于trian文件中，该作者将其放至在了下面的文件下，但在train文件中进行了调用。目录：faster_rcnn/train_utils/train_eval_utils.py。
[lang:Python] def
1
2
def train_one_epoch(model, optimizer, data_loader, device, epoch,
                    print_freq=50, warmup=False, scaler=None):

执行一个完整的训练周期
参数:
model: 待训练的PyTorch模型
optimizer: 优化器(如Adam/SGD)
data_loader: 训练数据加载器
device: 计算设备(‘cuda’/‘cpu’)
epoch: 当前周期序号
print_freq: 日志打印间隔(默认50个batch)
warmup: 是否启用学习率预热(默认False)
scaler: 混合精度训练的梯度缩放器(默认None)
1.初始化设置：设置为训练模式、指标记录器、添加学习率记录、日志头信息
2.warmup预热策略，避免训练初期因随机初始化导致梯度不稳定，逐步提高学习率

[lang:Python] warm
1
2
3
4
if epoch == 0 and warmup is True:
    warmup_factor = 1.0 / 1000  # 初始学习率缩放因子(从0.001倍开始)
    warmup_iters = min(1000, len(data_loader) - 1)  # 预热步数(最多1000步)
    lr_scheduler = utils.warmup_lr_scheduler(optimizer, warmup_iters, warmup_factor)

3.循环训练

[lang:Python] for
1
for i, [images, targets] in enumerate(metric_logger.log_every(data_loader, print_freq, header)):
数据转移到指定设备，使用.to(device)
前向传播：model(images,targets)
损失求和:：losses_reduced = sum(loss for loss in loss_dict_reduced.values())
更新滑动平均损失
检查损失是否有效
反向传播与参数更新：清空梯度optimizer.zero_grad()、反向传播losses.backward() 、参数更新optimizer.step()、学习率更新lr_scheduler.step()
记录指标