一、源码链接

Faster R-CNN源码链接https://pan.baidu.com/s/1SQjyLXD47H11ke05OXYSsQ?pwd=gt4f

二、环境配置

* Python3.6/3.7/3.8
* Pytorch1.7.1(注意：必须是1.6.0或以上，因为使用官方提供的混合精度训练1.6.0后才支持)
* pycocotools(Linux:`pip install pycocotools`; Windows:`pip install pycocotools-windows`(不需要额外安装vs))
* Ubuntu或Centos(不建议Windows)
* 最好使用GPU训练
* 详细环境配置见`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: 简易的预测脚本，使用训练好的权重进行预测测试
├── validation.py: 利用训练好的权重验证/测试数据的COCO指标，并生成record_mAP.txt文件
└── pascal_voc_classes.json: pascal_voc标签文件

四、预训练权重下载地址

## 预训练权重下载地址（下载后放入backbone文件夹中）：
* MobileNetV2 weights(下载后重命名为`mobilenet_v2.pth`，然后放到`bakcbone`文件夹下): https://download.pytorch.org/models/mobilenet_v2-b0353104.pth
* Resnet50 weights(下载后重命名为`resnet50.pth`，然后放到`bakcbone`文件夹下): https://download.pytorch.org/models/resnet50-0676ba61.pth
* ResNet50+FPN weights: https://download.pytorch.org/models/fasterrcnn_resnet50_fpn_coco-258fb6c6.pth
* 注意，下载的预训练权重记得要重命名，比如在train_resnet50_fpn.py中读取的是`fasterrcnn_resnet50_fpn_coco.pth`文件，
不是`fasterrcnn_resnet50_fpn_coco-258fb6c6.pth`，然后放到当前项目根目录下即可。

五、训练集

## 数据集，本例程使用的是PASCAL VOC2012数据集
* Pascal VOC2012 train/val数据集下载地址：http://host.robots.ox.ac.uk/pascal/VOC/voc2012/VOCtrainval_11-May-2012.tar
* 如果不了解数据集或者想使用自己的数据集进行训练，请参考我的bilibili：https://b23.tv/F1kSCK
* 使用ResNet50+FPN以及迁移学习在VOC2012数据集上得到的权重: 链接:https://pan.baidu.com/s/1ifilndFRtAV5RDZINSHj5w 提取码:dsz8

六、训练方法及注意事项

## 训练方法
* 确保提前准备好数据集
* 确保提前下载好对应预训练模型权重
* 若要训练mobilenetv2+fasterrcnn，直接使用train_mobilenet.py训练脚本
* 若要训练resnet50+fpn+fasterrcnn，直接使用train_resnet50_fpn.py训练脚本
* 若要使用多GPU训练，使用`python -m torch.distributed.launch --nproc_per_node=8 --use_env train_multi_GPU.py`指令,`nproc_per_node`参数为使用GPU数量
* 如果想指定使用哪些GPU设备可在指令前加上`CUDA_VISIBLE_DEVICES=0,3`(例如我只要使用设备中的第1块和第4块GPU设备)
* `CUDA_VISIBLE_DEVICES=0,3 python -m torch.distributed.launch --nproc_per_node=2 --use_env train_multi_GPU.py`

## 注意事项
* 在使用训练脚本时，注意要将`--data-path`(VOC_root)设置为自己存放`VOCdevkit`文件夹所在的**根目录**
* 由于带有FPN结构的Faster RCNN很吃显存，如果GPU的显存不够(如果batch_size小于8的话)建议在create_model函数中使用默认的norm_layer，
即不传递norm_layer变量，默认去使用FrozenBatchNorm2d(即不会去更新参数的bn层),使用中发现效果也很好。
* 训练过程中保存的`results.txt`是每个epoch在验证集上的COCO指标，前12个值是COCO指标，后面两个值是训练平均损失以及学习率
* 在使用预测脚本时，要将`train_weights`设置为你自己生成的权重路径。
* 使用validation文件时，注意确保你的验证集或者测试集中必须包含每个类别的目标，并且使用时只需要修改`--num-classes`、`--data-path`和`--weights-path`即可，其他代码尽量不要改动

七、大概看一下训练过程（train_mobilenetv2.py）


   
    
     
      
     
     
      
       def 
       main():
      
     
    
     
      
     
     
      
           device = torch.device(
       "cuda:0" 
       if torch.cuda.is_available() 
       else 
       "cpu")
      
     
    
     
      
     
     
          
       print(
       "Using {} device training.".
       format(device.
       type))
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
          
       # 用来保存coco_info的文件
      
     
    
     
      
     
     
      
           results_file = 
       "results{}.txt".
       format(datetime.datetime.now().strftime(
       "%Y%m%d-%H%M%S"))
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
          
       # 检查保存权重文件夹是否存在，不存在则创建
      
     
    
     
      
     
     
          
       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 = 
       "./"  
       # VOCdevkit
      
     
    
     
      
     
     
      
           aspect_ratio_group_factor = 
       3
      
     
    
     
      
     
     
      
           batch_size = 
       8
      
     
    
     
      
     
     
      
           amp = 
       False  
       # 是否使用混合精度训练，需要GPU支持
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
          
       # 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_dataset = VOCDataSet(VOC_root, 
       "2012", data_transform[
       "train"], 
       "train.txt")
      
     
    
     
      
     
     
      
           train_sampler = 
       None
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
          
       # 是否按图片相似高宽比采样图片组成batch
      
     
    
     
      
     
     
          
       # 使用的话能够减小训练时所需GPU显存，默认使用
      
     
    
     
      
     
     
          
       if aspect_ratio_group_factor >= 
       0:
      
     
    
     
      
     
     
      
               train_sampler = torch.utils.data.RandomSampler(train_dataset)
      
     
    
     
      
     
     
              
       # 统计所有图像高宽比例在bins区间中的位置索引
      
     
    
     
      
     
     
      
               group_ids = create_aspect_ratio_groups(train_dataset, k=aspect_ratio_group_factor)
      
     
    
     
      
     
     
              
       # 每个batch图片从同一高宽比例区间中取
      
     
    
     
      
     
     
      
               train_batch_sampler = GroupedBatchSampler(train_sampler, group_ids, batch_size)
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
      
           nw = 
       min([os.cpu_count(), batch_size 
       if batch_size > 
       1 
       else 
       0, 
       8])  
       # number of workers
      
     
    
     
      
     
     
          
       print(
       'Using %g dataloader workers' % nw)
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
          
       # 注意这里的collate_fn是自定义的，因为读取的数据包括image和targets，不能直接使用默认的方法合成batch
      
     
    
     
      
     
     
          
       if train_sampler:
      
     
    
     
      
     
     
              
       # 如果按照图片高宽比采样图片，dataloader中需要使用batch_sampler
      
     
    
     
      
     
     
      
               train_data_loader = torch.utils.data.DataLoader(train_dataset,
      
     
    
     
      
     
     
      
                                                               batch_sampler=train_batch_sampler,
      
     
    
     
      
     
     
      
                                                               pin_memory=
       True,
      
     
    
     
      
     
     
      
                                                               num_workers=nw,
      
     
    
     
      
     
     
      
                                                               collate_fn=train_dataset.collate_fn)
      
     
    
     
      
     
     
          
       else:
      
     
    
     
      
     
     
      
               train_data_loader = torch.utils.data.DataLoader(train_dataset,
      
     
    
     
      
     
     
      
                                                               batch_size=batch_size,
      
     
    
     
      
     
     
      
                                                               shuffle=
       True,
      
     
    
     
      
     
     
      
                                                               pin_memory=
       True,
      
     
    
     
      
     
     
      
                                                               num_workers=nw,
      
     
    
     
      
     
     
      
                                                               collate_fn=train_dataset.collate_fn)
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
          
       # load validation data set
      
     
    
     
      
     
     
          
       # VOCdevkit -> VOC2012 -> ImageSets -> Main -> val.txt
      
     
    
     
      
     
     
      
           val_dataset = VOCDataSet(VOC_root, 
       "2012", data_transform[
       "val"], 
       "val.txt")
      
     
    
     
      
     
     
      
           val_data_loader = torch.utils.data.DataLoader(val_dataset,
      
     
    
     
      
     
     
      
                                                         batch_size=
       1,
      
     
    
     
      
     
     
      
                                                         shuffle=
       False,
      
     
    
     
      
     
     
      
                                                         pin_memory=
       True,
      
     
    
     
      
     
     
      
                                                         num_workers=nw,
      
     
    
     
      
     
     
      
                                                         collate_fn=val_dataset.collate_fn)
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
          
       # create model num_classes equal background + 20 classes
      
     
    
     
      
     
     
      
           model = create_model(num_classes=
       21)
      
     
    
     
      
     
     
          
       # print(model)
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
      
           model.to(device)
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
      
           scaler = torch.cuda.amp.GradScaler() 
       if amp 
       else 
       None
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
      
           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)
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
      
           init_epochs = 
       5
      
     
    
     
      
     
     
          
       for epoch 
       in 
       range(init_epochs):
      
     
    
     
      
     
     
              
       # train for one epoch, printing every 10 iterations
      
     
    
     
      
     
     
      
               mean_loss, lr = utils.train_one_epoch(model, optimizer, train_data_loader,
      
     
    
     
      
     
     
      
                                                     device, epoch, print_freq=
       50,
      
     
    
     
      
     
     
      
                                                     warmup=
       True, scaler=scaler)
      
     
    
     
      
     
     
      
               train_loss.append(mean_loss.item())
      
     
    
     
      
     
     
      
               learning_rate.append(lr)
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
              
       # evaluate on the test dataset
      
     
    
     
      
     
     
      
               coco_info = utils.evaluate(model, val_data_loader, device=device)
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
              
       # write into txt
      
     
    
     
      
     
     
              
       with 
       open(results_file, 
       "a") 
       as f:
      
     
    
     
      
     
     
                  
       # 写入的数据包括coco指标还有loss和learning rate
      
     
    
     
      
     
     
      
                   result_info = [
       f"{i:.4f}" 
       for i 
       in coco_info + [mean_loss.item()]] + [
       f"{lr:.6f}"]
      
     
    
     
      
     
     
      
                   txt = 
       "epoch:{} {}".
       format(epoch, 
       ' '.join(result_info))
      
     
    
     
      
     
     
      
                   f.write(txt + 
       "\n")
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
      
               val_map.append(coco_info[
       1])  
       # pascal 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=
       3,
      
     
    
     
      
     
     
      
                                                          gamma=
       0.33)
      
     
    
     
      
     
     
      
           num_epochs = 
       20
      
     
    
     
      
     
     
          
       for epoch 
       in 
       range(init_epochs, num_epochs+init_epochs, 
       1):
      
     
    
     
      
     
     
              
       # train for one epoch, printing every 50 iterations
      
     
    
     
      
     
     
      
               mean_loss, lr = utils.train_one_epoch(model, optimizer, train_data_loader,
      
     
    
     
      
     
     
      
                                                     device, epoch, print_freq=
       50,
      
     
    
     
      
     
     
      
                                                     warmup=
       True, scaler=scaler)
      
     
    
     
      
     
     
      
               train_loss.append(mean_loss.item())
      
     
    
     
      
     
     
      
               learning_rate.append(lr)
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
              
       # update the learning rate
      
     
    
     
      
     
     
      
               lr_scheduler.step()
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
              
       # evaluate on the test dataset
      
     
    
     
      
     
     
      
               coco_info = utils.evaluate(model, val_data_loader, device=device)
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
              
       # write into txt
      
     
    
     
      
     
     
              
       with 
       open(results_file, 
       "a") 
       as f:
      
     
    
     
      
     
     
                  
       # 写入的数据包括coco指标还有loss和learning rate
      
     
    
     
      
     
     
      
                   result_info = [
       f"{i:.4f}" 
       for i 
       in coco_info + [mean_loss.item()]] + [
       f"{lr:.6f}"]
      
     
    
     
      
     
     
      
                   txt = 
       "epoch:{} {}".
       format(epoch, 
       ' '.join(result_info))
      
     
    
     
      
     
     
      
                   f.write(txt + 
       "\n")
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
      
               val_map.append(coco_info[
       1])  
       # pascal mAP
      
     
    
     
      
     
     
       
      
     
    
     
      
     
     
              
       # save weights
      
     
    
     
      
     
     
              
       # 仅保存最后5个epoch的权重
      
     
    
     
      
     
     
              
       if epoch 
       in 
       range(num_epochs+init_epochs)[-
       5:]:
      
     
    
     
      
     
     
      
                   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)

先判断是否有可用的GPU，若没有则使用CPU进行训练。

data_transform定义了图像预处理的函数。

VOC_root指定了VOC数据集的根目录。

VOCDataSet定义我们的数据集（通过mydataset.py文件定义我们自己的数据集）

DataLoader进行数据载入。

训练过程首先冻结前置特征提取网络权重（backbone），训练rpn以及最终预测网络部分。随后解冻前置特征提取网络权重（backbone），接着训练整个网络权重。

转载：https://blog.csdn.net/qq_41694024/article/details/128490710

查看评论

飞道的博客

飞道的博客

个人资料

文章分类

文章存档

阅读排行

评论排行

推荐文章

Faster RCNN网络源码解读（Ⅱ） --- Faster RCNN源码使用