本文提供给大家参考,若有不对的地方还请各位指教。
ok,学习Tensorflow物体检测之后,我们发现初始训练数据并不是我们自己提供的,那怎么才能换成我们自己的训练集呢?
下面是直接采用object-detection model 测试的结果,显示出了鼠标和杯子。下一步我们想加一些自己的数据给参与训练。
准备数据集:
我们可以自己生成TFRecord数据文件,推荐工具LabelImg。现在我们使用PASCAL VOC 2012数据集为例。下载地址
下载之后我们通过object-detection model 提供的create_pascal_tf_record.py文件将VOC2012数据转化为TFRecord文件。
在create_pascal_tf_record.py文件中修改以下地方:
python create_pascal_tf_record.py \
--label_map_path=data/pascal_label_map.pbtxt \
--data_dir=data/VOCdevkit \
--year=VOC2012 \
--set=train \
--output_path=data/pascal_train.record \
python create_pascal_tf_record.py \
--label_map_path=data/pascal_label_map.pbtxt \
--data_dir=data/VOCdevkit \
--year=VOC2012
--set=val \
--output_path=data/pascal_val.record
train 和 val 分别修改生成 pascal_train.record以及pascal_val.record文件。
下载模型
官方提供了不少预训练模型。我是地址
我推荐下载ssd_mobilenet_v1_coco模型,比较轻量级,之后测试训练速度相比其他比较快。
因为下载的是ssd_mobilenet_v1_coco模型,所以把object_detection/samples/configs/ssd_mobilenet_v1_pets.config文件复制到一个新建文件夹中,目录最好是在object-detection文件下。把之前下载的模型也复制到这个文件夹中,对
ssd_mobilenet_v1_pets.config进行修改:
num_classes:20
fine_tune_checkpoint:模型model.ckpt路径
train_input_reader:{
input_path:pascal_train.record路径
label_map_path:../data/pascal_label_map.pbtxt
}
eval_input_reader:{
input_path:pascal_val.record路径
label_map_path:../data/pascal_label_map.pbtxt
}
开始训练
配置object_detection/train.py文件(可以复制新建一份做出修改):
python ../../train.py \
--logtostderr \
--pipeline_config_path=ssd_mobilenet_v1.config 的路径
--train_dir=train_logs(随意一个路径,保存模型使用)
在object_detection train.py文件中,我们还需要把下载的model文件,内容可能包括checkpoint、frozen_inference_graph.pb等文件一共放置于pipeline_config_path同目录下。
运行train.py文件,ok已经开始打印日志文件了:
训练完成之后调用export_inference_graph.py把训练好的数据生成model就可以进行检测。
python export_inference_graph.py
--input_type image_tensor
--pipeline_config_path : 最开始复制的config文件,本文中举例的是ssd_mobilenet_v1_pets.config
--trained_checkpoint_prefix :训练好的mode.ckpt文件 ,例如 model.ckpt-n
--output_directory : 自己想保存的文件路径
执行python train.py报错:trainer.py“ModuleNotFoundError: No module named 'deployment'”,或者报错No module named 'nets'
错误原因,这两个model都是在TensorFlow/models/slim中,需要把models/slim加入到PYTHONPATH环境变量中。