您的位置：首页 > 其它

训练自己的物体检测模型

2019-05-12 18:02 246 查看

前言

在上一篇博客中，介绍了怎样利用SSD实现物体检测，我们使用的是官方的模型。在本文中，将继续介绍在此基础上如何实现迁移学习，训练自己的数据集。

1. 准备数据集

数据集来源一般有两个：使用labelImg自己标注，或者使用网上开源的数据集。本文使用的是开源的VOC 2007，总共有20个分类。数据集包括5001张图片，以及对应的
```
.xml
```
文件，分别存放在
```
image
```
和
```
annotation
```
文件夹中。
本文通过
```
.csv
```
文件生成
```
.record
```
文件，如果习惯直接使用
```
.xml
```
文件生成
```
.record
```
文件，可以跳过这一步。转换程序
```
xml_to_csv.py
```
已准备好，直接运行即可在当前路径下生成
```
labels.csv
```
，顺便说一下，本项目的
```
.py
```
程序都配备了一个同名的
```
.bat
```
启动程序，直接双击后者也可以，转换效果如图：

注意： 在准备数据阶段，可以同步准备测试集。在
```
xml_to_csv.py
```
中需要修改如下两个参数，对应测试集的
```
.xml
```
路径和输出的
```
.csv
```
文件名：
生成
```
.record
```
文件：
程序也已准备好：
```
generate_tfrecord.py
```
和
```
generate_tfrecord.bat
```
，我们只需要指定
```
.bat
```
文件中的三个参数即可运行。分别是：
```
csv_input:
```
标签文件csv
```
images_dir:
```
图片存放路径
```
output_path:
```
输出的文件

注意： 在
```
generate_tfrecord.py
```
有一段代码将20个类别转换为数字，如果改变数据集，这里是需要做相应调整的：
准备
```
.pbtxt
```
文件
这个文件用来解码数字和类别，要和之前
```
generate_tfrecord.py
```
中的对应关系一致：

2. 准备预训练好的模型

下载模型
tensorflow官方提供了很多训练好的模型，我们可以选择一个下载：

下载后解压，
```
.pb
```
文件是在之前的博客中用到过的，现在我们要继续训练，需要的是
```
.ckpt
```
文件，将他们放到
```
train
```
目录下
设置配置文件
在
```
./object_detection/samples/configs/
```
下有很多配置文件，对应各个模型，此处我们选择
```
ssd_mobilenet_v1_coco.config
```
，将其复制到当前路径，然后修改几个参数：

num_classes: 20
batch_size: 16 # 计算力强的话可以增大
num_steps: 2500 # 训练步数，达到就停止，即使再次运行
fine_tune_checkpoint: “./training/model.ckpt”
运行
```
train.bat
```
，第一个参数固定，第二个是配置文件，第三个是模型输出路径
运行结果

此时，我们还可以改大配置文件中的
```
num_steps
```
，继续训练