Faster R-CNN代码讲解之 train_mobilenet.py

资料：
github代码链接：https://github.com/WZMIAOMIAO/deep-learning-for-image-processing
b站一个不错的up主讲解视频：https://www.bilibili.com/video/BV1of4y1m7nj?t=99&p=2
数据集
数据集使用Pascal VOC2012 （共20个分类）
Pascal VOC2012 train/val数据集下载地址：http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar
文件结构
├── 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标签文件

1.检查是否有可用的gpu设备，有的话默认使用第一块gpu没有的话使用cpu

device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    print("Using {} device training.".format(device.type))

2.检查保存权重的文件夹是否存在，不存在则创建

if not os.path.exists("save_weights"):
    os.makedirs("save_weights")

3.图像预处理模块
这里的图像预处理模块儿跟分类网络里的图像预处理模块儿使用的方式不一样。在目标检测中如果你将你的图片进行随机翻转，那么你的真实框的坐标也要进行一个翻转，这点需要注意一下。
1）transforms.Compose（[函数1，函数2]）：图像处理操作的集合
2）transforms.ToTensor()：把图片转换成张量形式
3）transforms.RandomHorizontalFlip（）：随机水平翻转函数

data_transform = {
        "train": transforms.Compose([transforms.ToTensor(),
                                     transforms.RandomHorizontalFlip(0.5)]),
        "val": transforms.Compose([transforms.ToTensor()])
    }

3.检查VOC数据集文件夹是否存在
1）VOC_root = “./”：把VOC数据集的根目录保存为VOC_root，路径是在当前项目下。
2）检查路径是否存在，不存在将输出VOCdevkit dose not in path。

 VOC_root = "./"
    if os.path.exists(os.path.join(VOC_root, "VOCdevkit")) is False:
        raise FileNotFoundError("VOCdevkit dose not in path:'{}'.".format(VOC_root))

4.数据集的定义和载入
调用VOC2012DataSet来定义我们的数据集，VOC2012DataSet（）函数在my_dataset.py中定义。

#load train data set
# VOCdevkit -> VOC2012 -> ImageSets -> Main -> train.txt
train_data_set = VOC2012DataSet(VOC_root, data_transform["train"], "train.txt")

定义好数据集之后再通过torch包下的DataLoader（）函数对数据进行一个载入。注意这里的collate_fn是自定义的，因为读取的数据包括image和targets，不能直接使用默认的方法合成batch。
1）torch.utils.data.DataLoader（）：载入数据。
2）batch_size要根据GPU的显存来设置。
3）num_workers在liunx系统下要设置为4或者8，让它采用一个多性能的方式来进行图像读取和预处理，能够加快模型训练。

batch_size = 8
nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers
print('Using %g dataloader workers' % nw)
train_data_loader = torch.utils.data.DataLoader(train_data_set,
                                                batch_size=batch_size,
                                                shuffle=True,
                                                num_workers=nw,
                                                collate_fn=train_data_set.collate_fn)

同样的方法来定义验证集的数据集以及载入我们的datasets。

#load validation data set
#VOCdevkit -> VOC2012 -> ImageSets -> Main -> val.txt
    val_data_set = VOC2012DataSet(VOC_root, data_transform["val"], "val.txt")
    val_data_set_loader = torch.utils.data.DataLoader(val_data_set,
                                                      batch_size=batch_size,
                                                      shuffle=False,
                                                      num_workers=nw,
                                                      collate_fn=train_data_set.collate_fn)

5.设置模型的类别
VOC2012数据集一共20个类别，但在这里把背景作为一个类别0包括进去了。

#create model num_classes equal background + 20 classes
    model = create_model(num_classes=21)
    print(model)

然后把模型指派到设备当中，如果有GPU的话它会指派到GPU上。

model.to(device)

6.训练过程
训练过程在这里分为了两步。
**第一步，**首先冻结前置特征提取网络权重（backbone），使用优化器优化（训练）rpn以及fasterrcnn后半部分的网络 ，即，first frozen backbone and train 5 epochs。
（因为现在下载的权重是backbone的权重，而rpn的权重部分以及Fasterrcnn后半部分的权重还是一个初始化的状态，所以需要先去训练rpn网络的权重和fasterrcnn后半部分的权重。）

#冻结backbone网络的权重
 for param in model.backbone.parameters():
        param.requires_grad = False
# 定义一个优化器
 params = [p for p in model.parameters() if p.requires_grad]
 optimizer = torch.optim.SGD(params, lr=0.005,
                              momentum=0.9,weight_decay=0.0005)
#开始训练，训练一个epoch，每10次迭代打印一次
 num_epochs = 5
    for epoch in range(num_epochs):
            utils.train_one_epoch(model,optimizer,train_data_loader,device,e poch,print_freq=50,train_loss=train_loss,train_lr=learning_rate)
# evaluate on the test dataset在测试数据集上评估
            utils.evaluate(model, val_data_set_loader, device=device, mAP_list=val_mAP)
#保存训练到的权重
    torch.save(model.state_dict(), "./save_weights/pretrain.pth")

在这里插入一个小知识点，关于冻结训练的理解和使用。
冻结训练的作用：当我们已有部分预训练权重，这部分预训练权重所应用的那部分网络是通用的，如骨干网络，那么我们可以先冻结这部分权重的训练，将更多的资源放在训练后面部分的网络参数，这样使得时间和资源利用都能得到很大改善。然后后面的网络参数训练一段时间之后再解冻这些被冻结的部分，这时再全部一起训练。（转自：https://blog.csdn.net/yanghao201607030101/article/details/109898944）
冻结操作：

for param in model.backbone.parameters():
	param.requires_grad = False

解冻：

for param in model.backbone.parameters():
    param.requires_grad = True

更多冻结操作：https://www.zhihu.com/question/311095447/answer/589307812
优化器
torch.optim是实现各种优化算法的包。
要使用torch.optim，必须构造一个optimizer对象。这个对象能保存当前的参数状态并且基于计算梯度更新参数。要构造一个Optimizer，你必须给它一个包含参数（必须都是Variable对象）进行优化。然后，您可以指定optimizer的参数选项，比如学习率，权重衰减等。

torch.optim.SGD(params, lr=<required parameter>, momentum=0, dampening=0, weight_decay=0, nesterov=False)

params (iterable) – 可迭代的参数以优化或命令定义参数组
lr – 学习率
momentum (float, optional) – 动量因子 (default: 0)
weight_decay (float, optional) – 衰减(L2惩罚) (default: 0)
第二步，训练整个fasterrcnn网络的权重
解冻前置特征提取网络权重（backbone），接着训练整个网络权重 。

# 冻结backbone部分底层权重
for name, parameter in model.backbone.named_parameters():
    split_name = name.split(".")[0]
    if split_name in ["0", "1", "2", "3"]:
        parameter.requires_grad = False
    else:
        parameter.requires_grad = True
# 定义优化器
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=0.005,
                            momentum=0.9, weight_decay=0.0005)
# learning rate scheduler，每隔5步学习率乘以0.33
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
                                               step_size=5,
                                               gamma=0.33)
#训练，训练一个epoch，每50次迭代打印一次
num_epochs = 20
for epoch in range(num_epochs):
    utils.train_one_epoch(model, optimizer, train_data_loader,
                          device, epoch, print_freq=50,
                          train_loss=train_loss, train_lr=learning_rate)
    # 更新学习率
    lr_scheduler.step()
    # 在测试数据集上评估
    utils.evaluate(model, val_data_set_loader, device=device, mAP_list=val_mAP)
    # 保存权重，从第10个epoch开始保存，在这里10个epoch之后才开始收敛，所以从第10个以后开始保存。
    if epoch > 10:
        save_files = {
            'model': model.state_dict(),
            'optimizer': optimizer.state_dict(),
            'lr_scheduler': lr_scheduler.state_dict(),
            'epoch': epoch}
        torch.save(save_files, "./save_weights/mobile-model-{}.pth".format(epoch))

整体的train_mobilenet代码

import os

import torch
import torchvision
from torchvision.ops import misc

import transforms
from network_files.faster_rcnn_framework import FasterRCNN
from network_files.rpn_function import AnchorsGenerator
from backbone.mobilenetv2_model import MobileNetV2
from my_dataset import VOC2012DataSet
from train_utils import train_eval_utils as utils


def create_model(num_classes):
    # https://download.pytorch.org/models/mobilenet_v2-b0353104.pth
    backbone = MobileNetV2(weights_path="./backbone/mobilenet_v2.pth").features
    backbone.out_channels = 1280

    anchor_generator = AnchorsGenerator(sizes=((32, 64, 128, 256, 512),),
                                        aspect_ratios=((0.5, 1.0, 2.0),))

    roi_pooler = torchvision.ops.MultiScaleRoIAlign(featmap_names=['0'],  # 在哪些特征层上进行roi pooling
                                                    output_size=[7, 7],   # roi_pooling输出特征矩阵尺寸
                                                    sampling_ratio=2)  # 采样率

    model = FasterRCNN(backbone=backbone,
                       num_classes=num_classes,
                       rpn_anchor_generator=anchor_generator,
                       box_roi_pool=roi_pooler)

    return model


def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")#检查是否有可用的gpu设备，有的话默认使用第一块gpu没有的话使用cpu。
    print("Using {} device training.".format(device.type))

    # 检查保存权重文件夹是否存在，不存在则创建
    if not os.path.exists("save_weights"):
        os.makedirs("save_weights")

    data_transform = {
        "train": transforms.Compose([transforms.ToTensor(),
                                     transforms.RandomHorizontalFlip(0.5)]),
        "val": transforms.Compose([transforms.ToTensor()])
    }

    VOC_root = "./"
    # check voc root
    if os.path.exists(os.path.join(VOC_root, "VOCdevkit")) is False:
        raise FileNotFoundError("VOCdevkit dose not in path:'{}'.".format(VOC_root))

    # load train data set
    # VOCdevkit -> VOC2012 -> ImageSets -> Main -> train.txt
    train_data_set = VOC2012DataSet(VOC_root, data_transform["train"], "train.txt")
    # 注意这里的collate_fn是自定义的，因为读取的数据包括image和targets，不能直接使用默认的方法合成batch
    batch_size = 8
    nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8])  # number of workers
    print('Using %g dataloader workers' % nw)
    train_data_loader = torch.utils.data.DataLoader(train_data_set,
                                                    batch_size=batch_size,
                                                    shuffle=True,
                                                    num_workers=nw,
                                                    collate_fn=train_data_set.collate_fn)

    # load validation data set
    # VOCdevkit -> VOC2012 -> ImageSets -> Main -> val.txt
    val_data_set = VOC2012DataSet(VOC_root, data_transform["val"], "val.txt")
    val_data_set_loader = torch.utils.data.DataLoader(val_data_set,
                                                      batch_size=batch_size,
                                                      shuffle=False,
                                                      num_workers=nw,
                                                      collate_fn=train_data_set.collate_fn)

    # create model num_classes equal background + 20 classes
    model = create_model(num_classes=21)
    # print(model)

    model.to(device)

    train_loss = []
    learning_rate = []
    val_mAP = []

    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    #  first frozen backbone and train 5 epochs                   #
    #  首先冻结前置特征提取网络权重（backbone），训练rpn以及最终预测网络部分 #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    for param in model.backbone.parameters():
        param.requires_grad = False

    # define optimizer
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.SGD(params, lr=0.005,
                                momentum=0.9, weight_decay=0.0005)

    num_epochs = 5
    for epoch in range(num_epochs):
        # train for one epoch, printing every 10 iterations
        utils.train_one_epoch(model, optimizer, train_data_loader,
                              device, epoch, print_freq=50,
                              train_loss=train_loss, train_lr=learning_rate)

        # evaluate on the test dataset
        utils.evaluate(model, val_data_set_loader, device=device, mAP_list=val_mAP)

    torch.save(model.state_dict(), "./save_weights/pretrain.pth")

    # # # # # # # # # # # # # # # # # # # # # # # # # # # #
    #  second unfrozen backbone and train all network     #
    #  解冻前置特征提取网络权重（backbone），接着训练整个网络权重  #
    # # # # # # # # # # # # # # # # # # # # # # # # # # # #

    # 冻结backbone部分底层权重
    for name, parameter in model.backbone.named_parameters():
        split_name = name.split(".")[0]
        if split_name in ["0", "1", "2", "3"]:
            parameter.requires_grad = False
        else:
            parameter.requires_grad = True

    # define optimizer
    params = [p for p in model.parameters() if p.requires_grad]
    optimizer = torch.optim.SGD(params, lr=0.005,
                                momentum=0.9, weight_decay=0.0005)
    # learning rate scheduler
    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
                                                   step_size=5,
                                                   gamma=0.33)
    num_epochs = 20
    for epoch in range(num_epochs):
        # train for one epoch, printing every 50 iterations
        utils.train_one_epoch(model, optimizer, train_data_loader,
                              device, epoch, print_freq=50,
                              train_loss=train_loss, train_lr=learning_rate)
        # update the learning rate
        lr_scheduler.step()

        # evaluate on the test dataset
        utils.evaluate(model, val_data_set_loader, device=device, mAP_list=val_mAP)

        # save weights
        if epoch > 10:
            save_files = {
                'model': model.state_dict(),
                'optimizer': optimizer.state_dict(),
                'lr_scheduler': lr_scheduler.state_dict(),
                'epoch': epoch}
            torch.save(save_files, "./save_weights/mobile-model-{}.pth".format(epoch))

    # 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)

    # model.eval()
    # x = [torch.rand(3, 300, 400), torch.rand(3, 400, 400)]
    # predictions = model(x)
    # print(predictions)


if __name__ == "__main__":
    version = torch.version.__version__[:5]  # example: 1.6.0
    # 因为使用的官方的混合精度训练是1.6.0后才支持的，所以必须大于等于1.6.0
    if version < "1.6.0":
        raise EnvironmentError("pytorch version must be 1.6.0 or above")

    main()