faster-rcnn 训练自己的数据集
本文详细介绍了如何使用Faster R-CNN模型对自定义数据集进行训练,并提供了完整的操作指南:
模型概述:Faster R-CNN是一种基于卷积神经网络(CNN)的数据集检测算法,结合候选框(RPN)和分类器(R-CNN),通过关注机制提升检测效果。
数据准备:
- 数据集需按VOC格式组织。
- 需划分训练与验证数据集。
- 使用工具生成2007train.txt和2007val.txt标签文件。
预处理:- 通过Python脚本将VOC格式标注文件转换为适合Faster R-CNN的数据格式。
- 使用utils.py中的函数完成目标检测任务。
配置文件调整:- 修改classes_path以指定类别名称。
- 更新trainannotationpath和valannotationpath指向生成的标签文件路径。
模型训练:- 使用提供的PyTorch实现库进行模型下载和安装。
- 设置超参数(如学习率、世代数等)并执行训练过程。
- 提供了详细的代码示例和命令行运行说明。
结果分析:- 通过验证集评估模型性能。
- 计算平均精度(mAP)以评估模型效果。
整篇文章结构清晰,内容详实,并附带了完整的代码示例和运行说明,对读者理解和实践具有很高的参考价值。
本博客专注于自建数据集基于VOC格式的Faster-RCNN模型训练实践。
基于快速RCNN算法的改进而形成。该算法由两大核心组件构成,请参考下图进行了解。第一部分是一个用于提取候选框的全卷积神经网络(RPN),第二部分则是一个基于候选框进行目标检测的快照式RCNN架构。其整体检测流程整合在一个神经网络架构中。为了实现多目标跟踪功能,Faster R-CNN采用了关注机制(attention mechanism)。具体而言,RPN组件指示模型应将注意力集中在特定区域。

这篇文档详细阐述了利用Faster R-CNN算法完成自定义数据集训练的具体方法,并且提供了常见问题及其应对策略。
源码基于开源代码实现,并可访问其完整代码库以获取详细技术文档
该库基于深度学习框架PyTorch实现了Faster R-CNN算法,在VOC数据集上进行了训练与验证。该库基于深度学习框架PyTorch实现了Faster R-CNN算法,在VOC数据集上进行了训练与验证. 为了方便用户使用与扩展功能,请考虑在GitHub上创建个人账户并参与代码贡献.

该库实现了基于PyTorch框架的Faster R-CNN算法的开发,并支持基于VOC数据集格式的数据训练。
代码结构如下图所示

其中,
① VOCdevkit为存放数据集的目录(VOC格式)
② img为测试图片目录
③ img_out 为存放测试结果的目录
记录文件用于存储训练数据的日志信息以及模型的存储位置
⑤ model_data 为相关配置文件
⑥ nets为网络代码
⑦ utils为相关工具代码
⑧ 其他的为训练,检测等相关代码
1、准备数据集
需要准备VOC格式的数据集(以VOC格式数据集为例)

数据集中,Annotations为标注文件(xml格式),JPEGImages为图片文件
2、数据集预处理
① 划分训练与验证数据集
将VOC数据集按照特定要求划分为训练集与验证集
可以调用如下代码执行,下面代码是按9:1的比例划分的
# coding:utf-8
import os
import random
import argparse
parser = argparse.ArgumentParser()
#xml文件的地址,根据自己的数据进行修改 xml一般存放在Annotations下
parser.add_argument('--xml_path', default='VOC2007/Annotations', type=str, help='input xml label path')
#数据集的划分,地址选择自己数据下的ImageSets/Main
parser.add_argument('--txt_path', default='VOC2007/ImageSets/Main', type=str, help='output txt label path')
opt = parser.parse_args()
trainval_percent = 1.0 # 训练集和验证集所占比例。 这里没有划分测试集
train_percent = 0.9 # 训练集所占比例,可自己进行调整
xmlfilepath = opt.xml_path
txtsavepath = opt.txt_path
total_xml = os.listdir(xmlfilepath)
if not os.path.exists(txtsavepath):
os.makedirs(txtsavepath)
num = len(total_xml)
list_index = range(num)
tv = int(num * trainval_percent)
tr = int(tv * train_percent)
trainval = random.sample(list_index, tv)
train = random.sample(trainval, tr)
file_trainval = open(txtsavepath + '/trainval.txt', 'w')
file_test = open(txtsavepath + '/test.txt', 'w')
file_train = open(txtsavepath + '/train.txt', 'w')
file_val = open(txtsavepath + '/val.txt', 'w')
for i in list_index:
name = total_xml[i][:-4] + '\n'
if i in trainval:
file_trainval.write(name)
if i in train:
file_train.write(name)
else:
file_val.write(name)
else:
file_test.write(name)
file_trainval.close()
file_train.close()
file_val.close()
file_test.close()
② 格式转换,voc转annotation,生成对应的txt文件
下面代码执行成功后,将生成2007_train.txt与2007_val.txt两个文件
import os
import random
import xml.etree.ElementTree as ET
import numpy as np
from utils.utils import get_classes
#--------------------------------------------------------------------------------------------------------------------------------#
# annotation_mode用于指定该文件运行时计算的内容
# annotation_mode为0代表整个标签处理过程,包括获得VOCdevkit/VOC2007/ImageSets里面的txt以及训练用的2007_train.txt、2007_val.txt
# annotation_mode为1代表获得VOCdevkit/VOC2007/ImageSets里面的txt
# annotation_mode为2代表获得训练用的2007_train.txt、2007_val.txt
#--------------------------------------------------------------------------------------------------------------------------------#
annotation_mode = 2
#-------------------------------------------------------------------#
# 必须要修改,用于生成2007_train.txt、2007_val.txt的目标信息
# 与训练和预测所用的classes_path一致即可
# 如果生成的2007_train.txt里面没有目标信息
# 那么就是因为classes没有设定正确
# 仅在annotation_mode为0和2的时候有效
#-------------------------------------------------------------------#
classes_path = 'model_data/hand.txt'
#--------------------------------------------------------------------------------------------------------------------------------#
# trainval_percent用于指定(训练集+验证集)与测试集的比例,默认情况下 (训练集+验证集):测试集 = 9:1
# train_percent用于指定(训练集+验证集)中训练集与验证集的比例,默认情况下 训练集:验证集 = 9:1
# 仅在annotation_mode为0和1的时候有效
#--------------------------------------------------------------------------------------------------------------------------------#
trainval_percent = 0.9
train_percent = 0.9
#-------------------------------------------------------#
# 指向VOC数据集所在的文件夹
# 默认指向根目录下的VOC数据集
#-------------------------------------------------------#
VOCdevkit_path = 'VOCdevkit_shake_hands'
VOCdevkit_sets = [('2007', 'train'), ('2007', 'val')]
classes, _ = get_classes(classes_path)
#-------------------------------------------------------#
# 统计目标数量
#-------------------------------------------------------#
photo_nums = np.zeros(len(VOCdevkit_sets))
nums = np.zeros(len(classes))
def convert_annotation(year, image_id, list_file):
in_file = open(os.path.join(VOCdevkit_path, 'VOC%s\ Annotations\ %s.xml'%(year, image_id)), encoding='utf-8')
tree=ET.parse(in_file)
root = tree.getroot()
for obj in root.iter('object'):
difficult = 0
if obj.find('difficult')!=None:
difficult = obj.find('difficult').text
cls = obj.find('name').text
if cls not in classes or int(difficult)==1:
continue
cls_id = classes.index(cls)
xmlbox = obj.find('bndbox')
b = (int(float(xmlbox.find('xmin').text)), int(float(xmlbox.find('ymin').text)), int(float(xmlbox.find('xmax').text)), int(float(xmlbox.find('ymax').text)))
list_file.write(" " + ",".join([str(a) for a in b]) + ',' + str(cls_id))
nums[classes.index(cls)] = nums[classes.index(cls)] + 1
if __name__ == "__main__":
random.seed(0)
if " " in os.path.abspath(VOCdevkit_path):
raise ValueError("数据集存放的文件夹路径与图片名称中不可以存在空格,否则会影响正常的模型训练,请注意修改。")
if annotation_mode == 0 or annotation_mode == 1:
print("Generate txt in ImageSets.")
xmlfilepath = os.path.join(VOCdevkit_path, 'VOC2007\ Annotations')
saveBasePath = os.path.join(VOCdevkit_path, 'VOC2007\ ImageSets\ Main')
temp_xml = os.listdir(xmlfilepath)
total_xml = []
for xml in temp_xml:
if xml.endswith(".xml"):
total_xml.append(xml)
num = len(total_xml)
list = range(num)
tv = int(num*trainval_percent)
tr = int(tv*train_percent)
trainval= random.sample(list,tv)
train = random.sample(trainval,tr)
print("train and val size",tv)
print("train size",tr)
ftrainval = open(os.path.join(saveBasePath,'trainval.txt'), 'w')
ftest = open(os.path.join(saveBasePath,'test.txt'), 'w')
ftrain = open(os.path.join(saveBasePath,'train.txt'), 'w')
fval = open(os.path.join(saveBasePath,'val.txt'), 'w')
for i in list:
name=total_xml[i][:-4]+'\n'
if i in trainval:
ftrainval.write(name)
if i in train:
ftrain.write(name)
else:
fval.write(name)
else:
ftest.write(name)
ftrainval.close()
ftrain.close()
fval.close()
ftest.close()
print("Generate txt in ImageSets done.")
if annotation_mode == 0 or annotation_mode == 2:
print("Generate 2007_train.txt and 2007_val.txt for train.")
type_index = 0
for year, image_set in VOCdevkit_sets:
image_ids = open(os.path.join(VOCdevkit_path, 'VOC%s\ ImageSets\ Main\ %s.txt'%(year, image_set)), encoding='utf-8').read().strip().split()
list_file = open('%s_%s.txt'%(year, image_set), 'w', encoding='utf-8')
for image_id in image_ids:
list_file.write('%s/VOC%s/JPEGImages/%s.jpg'%(os.path.abspath(VOCdevkit_path), year, image_id))
convert_annotation(year, image_id, list_file)
list_file.write('\n')
photo_nums[type_index] = len(image_ids)
type_index += 1
list_file.close()
print("Generate 2007_train.txt and 2007_val.txt for train done.")
def printTable(List1, List2):
for i in range(len(List1[0])):
print("|", end=' ')
for j in range(len(List1)):
print(List1[j][i].rjust(int(List2[j])), end=' ')
print("|", end=' ')
print()
str_nums = [str(int(x)) for x in nums]
tableData = [
classes, str_nums
]
colWidths = [0]*len(tableData)
len1 = 0
for i in range(len(tableData)):
for j in range(len(tableData[i])):
if len(tableData[i][j]) > colWidths[i]:
colWidths[i] = len(tableData[i][j])
printTable(tableData, colWidths)
if photo_nums[0] <= 500:
print("训练集数量小于500,属于较小的数据量,请注意设置较大的训练世代(Epoch)以满足足够的梯度下降次数(Step)。")
if np.sum(nums) == 0:
print("在数据集中并未获得任何目标,请注意修改classes_path对应自己的数据集,并且保证标签名字正确,否则训练将会没有任何效果!")
print("在数据集中并未获得任何目标,请注意修改classes_path对应自己的数据集,并且保证标签名字正确,否则训练将会没有任何效果!")
print("在数据集中并未获得任何目标,请注意修改classes_path对应自己的数据集,并且保证标签名字正确,否则训练将会没有任何效果!")
print("(重要的事情说三遍)。")
③ 生成配置文件
在model_data目录中创建一个basketball.txt文件,并用于存储需要检测的类别名称。
2、修改对应的配置(是训练代码执行的是修改好的数据集)
① 修改
classes_path = 'model_data/basketball.txt'
将刚才新建的配置文件路径赋值给 classes_path
② 修改
train_annotation_path = '2007_train.txt'
val_annotation_path = '2007_val.txt'
这两个TXT格式的配置文件同样是由前述操作产出的,并且其中包含的信息是训练数据集与测试数据集相对应的列表
3、运行
python train.py
进行训练
最终将出现如下打印
PS D:\part_20220703\faster-rcnn-pytorch> python.exe .\train.py
Number of devices: 1
initialize network with normal type
Load weights model_data/voc_weights_resnet.pth.
Successful Load Key: ['extractor.0.weight', 'extractor.1.weight', 'extractor.1.bias', 'extractor.1.running_mean', 'extractor.1.running_var', 'extractor.1.num_batches_tracked', 'extractor.4.0.conv1.weight', 'extractor.4.0.bn1.weight', 'extractor.4.0.bn1.bias', 'extractor.4.0.bn1.running_mean', 'extractor.4.0.bn1.running_var', 'extractor.4.0.bn1.num_batches_tracked', 'extractor.4.0.conv2.weight', 'extractor.4.0.bn2.weight', 'extractor.4.0.bn2.bias', 'extractor.4.0.bn2.running_mean', 'extractor.4.0.bn2.running_var', 'e ……
Successful Load Key Num: 324
9. Fail To Load Key: ['head.cls_loc.weight', 'head.cls_loc.bias', 'head.score.weight', 'head.score.bias'] ……
Fail To Load Key num: 4
温馨提示,head部分没有载入是正常现象,Backbone部分没有载入是错误的。
Configurations:
----------------------------------------------------------------------
|keys|values|
----------------------------------------------------------------------
|classes_path|model_data/home.txt|
|model_path|model_data/voc_weights_resnet.pth|
|input_shape|[600, 600]|
|Init_Epoch|0|
|Freeze_Epoch|500|
|UnFreeze_Epoch|200|
|Freeze_batch_size|2|
|Unfreeze_batch_size|2|
|Freeze_Train|True|
|Init_lr|0.0001|
|Min_lr|1.0000000000000002e-06|
|optimizer_type|adam|
|momentum|0.9|
|lr_decay_type|cos|
|save_period|5|
|save_dir|logs|
|num_workers|4|
|num_train|45|
|num_val|5|
----------------------------------------------------------------------
[Warning] 使用adam优化器时,建议将训练总步长设置到15000以上。
[Warning] 本次运行的总训练数据量为45,Unfreeze_batch_size为2,共训练200个Epoch,计算出总训练步长为4400。
[Warning] 由于总训练步长为4400,小于建议总步长15000,建议设置总世代为682。
Start Train
Epoch 1/200: 100%|███████████████████████████████████████████████████████████████████████████████████████████████████████| 22/22 [23:08<00:00, 63.13s/it, lr=1e-5, roi_cls=2.75, roi_loc=1.77, rpn_cls=0.131, rpn_loc=0.192, total_loss=4.84]
Finish Train
Start Validation
Epoch 1/200: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [01:27<00:00, 43.87s/it, val_loss=5.09]
Finish Validation
Epoch:1/200
Total Loss: 4.840 || Val Loss: 5.085
Save best model to best_epoch_weights.pth
Start Train
Epoch 2/200: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████| 22/22 [24:33<00:00, 67.00s/it, lr=2e-5, roi_cls=1.41, roi_loc=1.98, rpn_cls=0.1, rpn_loc=0.18, total_loss=3.67]
Finish Train
Start Validation
Epoch 2/200: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [01:28<00:00, 44.30s/it, val_loss=2.83]
Finish Validation
Epoch:2/200
Total Loss: 3.666 || Val Loss: 2.827
Save best model to best_epoch_weights.pth
Start Train
Epoch 3/200: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████| 22/22 [24:49<00:00, 67.70s/it, lr=5e-5, roi_cls=0.216, roi_loc=1.85, rpn_cls=0.0605, rpn_loc=0.177, total_loss=2.31]
Finish Train
Start Validation
Epoch 3/200: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [01:29<00:00, 44.56s/it, val_loss=2.47]
Finish Validation
Epoch:3/200
Total Loss: 2.307 || Val Loss: 2.467
Save best model to best_epoch_weights.pth
Start Train
Epoch 4/200: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 22/22 [24:42<00:00, 67.37s/it, lr=0.0001, roi_cls=0.114, roi_loc=1.88, rpn_cls=0.0512, rpn_loc=0.0922, total_loss=2.14]
Finish Train
Start Validation
Epoch 4/200: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [01:25<00:00, 42.81s/it, val_loss=2.28]
Finish Validation
Epoch:4/200
Total Loss: 2.140 || Val Loss: 2.282
Save best model to best_epoch_weights.pth
Start Train
Epoch 5/200: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 22/22 [22:51<00:00, 62.33s/it, lr=0.0001, roi_cls=0.147, roi_loc=1.74, rpn_cls=0.0406, rpn_loc=0.0904, total_loss=2.02]
Finish Train
Start Validation
Epoch 5/200: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [01:25<00:00, 42.76s/it, val_loss=2.59]
Finish Validation
Get map.
100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 5/5 [06:41<00:00, 80.22s/it]
Calculate Map.
71.83% = chair AP || score_threhold=0.5 : F1=0.29 ; Recall=16.67% ; Precision=100.00%
mAP = 71.83%
Get map done.
Epoch:5/200
Total Loss: 2.021 || Val Loss: 2.593
Start Train
Epoch 6/200: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████| 22/22 [21:40<00:00, 59.11s/it, lr=0.0001, roi_cls=0.128, roi_loc=1.74, rpn_cls=0.0338, rpn_loc=0.0855, total_loss=1.98]
Finish Train
Start Validation
Epoch 6/200: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [01:19<00:00, 39.79s/it, val_loss=2.22]
Finish Validation
Epoch:6/200
Total Loss: 1.985 || Val Loss: 2.223
Save best model to best_epoch_weights.pth
Start Train
Epoch 7/200: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████| 22/22 [21:04<00:00, 57.50s/it, lr=9.99e-5, roi_cls=0.129, roi_loc=1.36, rpn_cls=0.0265, rpn_loc=0.0743, total_loss=1.59]
Finish Train
Start Validation
Epoch 7/200: 100%|██████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 2/2 [01:18<00:00, 39.34s/it, val_loss=2.05]
Finish Validation
Epoch:7/200
Total Loss: 1.593 || Val Loss: 2.052
即可正常进行训练啦
