Tensorflow 使用预训练模型训练的完整流程

liuian 2025-04-11 01:00 77 浏览

前面已经介绍了深度学习框架Tensorflow 的图像的标注和训练数据的准备工作，本文介绍一下使用预训练模型完成训练并导出训练的模型。

1.选择预训练模型

1.1 下载预训练模型

首先需要在Tensorflow 官网下载模型models-master，下载后在本地解压，找到
models/research/object_detection/g3doc/tf1_detection_zoo.md，用文本的方式打开可以看到很多目标检测模型的下载链接，根据链接地址可下载预训练模型（以目标检测为例）。

下载预训练模型，

1.2 COCO数据集介绍

模型名称：
ssd_mobilenet_v1_coco_2018_01_28 。v1 表示版本用于tf1，coco 是数据集名称。顺便提一下，COCO数据集是微软发布的一个可以用来进行图像识别训练的数据集，图像中的目标都经过精确的segmentation进行位置定位，COCO数据集包括90类目标。Object Detection API默认提供了5个预训练模型，都是使用COCO数据集训练的，分别为

SSD + MobileNet
Inception V2 + SSD
ResNet101 + R-CNN
ResNet101 + Faster R-CNN
Inception-ResNet V2 + Faster R-CNN

有了这些储备知识，在下一步我们修改预训练模型配置时，有个准备。

2.参数设置

2.1 解压模型文件

2.2 修改模型配置文件

用文本的方式打开 pipline.config，修改如下内容，

1）分类的数目

num_classes: 3

----分类数目需要与你样本标注的类型数量保持一致

2）一次训练所选取的样本数

batch_size: 8

----batch size 指的是一次训练所选取的样本数,根据电脑硬件配置做适当更改，显存较小可以适当减小该值

3）预训练模型ckpt文件的位置

fine_tune_checkpoint: "/home/XXX/modelfiles/model.ckpt"

----modelfiles 就是预训练模型的别名，里面记录了保存的最新的checkpoint文件以及其它checkpoint文件列表，包括了model.ckpt.data;model.ckpt.index;model.ckpt.meta

4）训练数据和标签文件保存路径

train_input_reader: {

tf_record_input_reader {

input_path: "/home/XXX/record/train.record"

}

label_map_path: "/home/XXX/label_map.pbtxt"

}

5）测试数据保存位置和相应标签文件路径

# shuffle表示是否随机选取测试图片

eval_input_reader: {

tf_record_input_reader {

input_path:"/home/XXX/record/test.record"

}

label_map_path: "/home/XXX/label_map.pbtxt"

shuffle: false

num_readers: 1

}

设置训练集和测试集数据存放路径，推荐绝对路径

6）训练的步数

num_steps: 200000

3.启动训练

配置文件设置完成后，准备启动训练了。进入model_main.py 所在的目录，

用如下格式的指令启动训练，

python model_main.py --logtostderr --pipeline_config_path=/XXX/XXX/ssd_mobilenet_v2_oid_v4.config --train_dir=/XXX/XXX/XXX/train_dir --model_dir=/XXX/XXX/XXX/train_dir/trained_ckpt

训练次数已经在配置文件中设置过了，指令中可以不必添加。

4.Tensorboard 训练过程可视化

训练启动后，可以再打开一个终端，用如下指令查看训练详情，

tensorboard --logdir=/XXX/XXX/train_dir/trained_ckpt --host=127.0.0.1

5.导出模型

模型收敛后可用ctrl+z 退出训练。之后用如下指令到处已训练好的模型，

python export_inference_graph.py --input_type image_tensor --pipeline_config_path ./XXX/ssd_mobilenet_v2_oid_v4.config --trained_checkpoint_prefix ./XXX/train_dir/trained_ckpt/model.ckpt-124467 --output_directory ./ssd_mobilenet_v2_oid_v4_2018_12_12_led/output_dir

----------124467 指的是模型训练次数

6.根据训练结果生成pbtxt文件

训练的模型需要配合模型类别描述文件(.pbtxt)一起使用完成预测。用tf_text_graph_ssd.py生成pbtxt(mn_dataset_led.pbtxt)文件的指令如下，

python tf_text_graph_ssd.py --input ./XXX/output_dir/frozen_inference_graph.pb --config ./XXX/ssd_mobilenet_v2_oid_v4.config --output ./XXX/output_dir/mn_dataset_led.pbtxt

导出模型后(frozen_inference_graph.pb+mn_dataset_led.pbtxt)，类别在项目中指定，最后再结合OpenCV 的DNN 模块功能，就可以完成模型的部署和应用。

End

tensorboard下载