【Tensorflow】使用SSD-Mobilenet训练模型
1. 准备工作
-
从github上克隆tensorflowmodels
国内不翻墙的情况下,各种办法尝试git clone的速度都特别的慢,没办法最后用迅雷下载了zip压缩包到本地解压缩使用; -
数据集需要自己制作,关于制作过程在最后会介绍,此处先使用VOC2012数据集;
-
下载原有模型:SSD-Mobilenet,在该网页同样存在其他模型可供选择。
2. 环境配置
- 都是在该目录下完成;
- 使用protobuf来配置模型和训练参数,所以API正常使用必须先安装protobuf库:
protoc object_detection/protos/*.proto --python_out=.
如果出现报错如下:
解决办法:点击该处链接
- 然后将models和slim(tf高级框架)加入python环境变量:
export PYTHONPATH="${PYTHONPATH}:/home/Github/models:/home/Github/models/research/slim/"
- 继续在research/目录下执行:
sudo python3 setup.py build
sudo python3 setup.py install
3. 配置和训练
- 在
object_detection/
目录下创建目录ssd_model
sudo mkdir ssd_model
- 把下载好的数据集解压进去,数据集路径为
解压代码:
tar -xvf VOCtrainval_11-May-2012.tar
- 执行配置文件
./object_detection/dataset_tools/create_pascal_tf_record.py
:
该配置文件的作用在于转换原始的PASCAL dataset到TFRecord为了目标检测,就是一个准换成合适的数据集的过程:
python3 ./object_detection/dataset_tools/create_pascal_tf_record.py --label_map_path=object_detection/data/pascal_label_map.pbtxt --data_dir=object_detection/ssd_model/VOCdevkit/ --year=VOC2012 --set=train --output_path=object_detection/ssd_model/pascal_train.record
python3 ./object_detection/dataset_tools/create_pascal_tf_record.py --label_map_path=object_detection/data/pascal_label_map.pbtxt --data_dir=object_detection/ssd_model/VOCdevkit/ --year=VOC2012 --set=val --output_path=object_detection/ssd_model/pascal_val.record
之后会在ssd_model/
目录下生成pascal_train.record和pascal_val.record两个文件。
- 下一步复制训练
pet
数据用到的配置文件,我们在这个基础上修改配置,训练我们的数据:
在object_detection/samples/configs/
目录下有很多的配置文件样例可以用来使用,并非局限于本文使用的:
cp object_detection/data/pascal_label_map.pbtxt object_detection/ssd_model/
cp object_detection/samples/configs/ssd_mobilenet_v1_pets.config object_detection/ssd_model/
-
打开
pascal_label_map.pbtxt
看一下,这个文件里面是类似Json格式的label集,列出了数据集里有哪些label。Pascal这个数据集label
共有20个。 -
之后打开配置文件
ssd_mobilenet_v1_pets.config
,把num_classes
改为20
配置默认训练次数num_steps: 200000
,我们根据自己需要改,注意这个训练是很慢的,差不多以天为单位,所以可以适当改小点,我改成了2000。 -
然后更改一些文件路径:
train_input_reader: {
tf_record_input_reader {
input_path: "/home/×××/×××/workspace/mobilenetssd/models-master/research/object_detection/ssd_model/pascal_train.record"
}
label_map_path: "/home/×××/×××/workspace/mobilenetssd/models-master/research/object_detection/ssd_model/pascal_label_map.pbtxt"
}
eval_config: {
metrics_set: "coco_detection_metrics"
num_examples: 1100
}
eval_input_reader: {
tf_record_input_reader {
input_path: "/home/×××/×××/workspace/mobilenetssd/models-master/research/object_detection/ssd_model/pascal_val.record"
}
label_map_path: "/home/×××/×××/workspace/mobilenetssd/models-master/research/object_detection/ssd_model/pascal_label_map.pbtxt"
shuffle: false
num_readers: 1
}
- 把之前下载的
ssd_mobilenet
解压到/object_detection/ssd_model/ssd_mobilenet
下:
tar -xzvf ssd_mobilenet_v1_coco_2018_01_28.tar.gz
- 把路径填进配置文件:
fine_tune_checkpoint: "/home/jari/guoshi/workspace/mobilenetssd/models-master/research/object_detection/ssd_model/ssd_mobilenet_v1_coco_2018_01_28/model.ckpt"
- 准备好上述文件后就可以直接调用model_main.py文件进行训练,同时在
object_detection/
下新建目录training
:
python3 object_detection/model_main.py \
--pipeline_config_path object_detection/ssd_model/ssd_mobilenet_v1_pets.config \
--model_dir=object_detection/training \
--num_train_steps=50000 \
--num_eval_steps=2000 \
--alsologtostderr
之后疯狂报错,报错和解决途径如下:
AttributeError: module 'tensorflow._api.v1.compat' has no attribute 'v2'
解决:
更改如下代码:
def global_pool(input_tensor, pool_op=tf.compat.v2.nn.avg_pool2d):
def global_pool(input_tensor, pool_op=tf.nn.avg_pool):
AttributeError: module 'tensorflow._api.v1.compat' has no attribute 'v1'
解决:
weights_init = tf.compat.v1.truncated_normal_initializer(stddev=stddev)
改为:
weights_init = tf.truncated_normal_initializer(stddev=stddev)
- 如果报找不到nets,就把research文件夹内的slim下的nets复制到object_detection文件夹
- 针对只训练了一次就没有任何报错的停止了,并且training目录下什么都没有,修改了训练的指令如上,应该是tensorflow更改了版本的问题;
- 可以利用tensorboard来观察训练过程:
~/xxx/xxx/xxx/models-master/research/object_detection$ tensorboard --logdir=training
通过链接打开就可以观察到:
4. 模型导出
- 经过漫长的等待,可以看到在
/object_detection/training
目录下生成了模型。然后创建文件夹ssd_model/model
:
python3 object_detection/export_inference_graph.py \
--input_type image_tensor \
--pipeline_config_path object_detection/ssd_model/ssd_mobilenet_v1_pets.config \
--trained_checkpoint_prefix object_detection/training/model.ckpt-18131 \
--output_directory object_detection/ssd_model/model/
生成frozen_inference_graph.pb文件,最重要:
- 测试模型:
在/object_detection
目录下,新建一个test.py
文件,内容如下:
import cv2
import numpy as np
import tensorflow as tf
from object_detection.utils import label_map_util
from object_detection.utils import visualization_utils as vis_util
class TOD(object):
def __init__(self):
self.PATH_TO_CKPT = '/home/***/***/***/mobilenetssd/models-master/research/object_detection/ssd_model/model/frozen_inference_graph.pb'
self.PATH_TO_LABELS = '/home/***/***/***/mobilenetssd/models-master/research/object_detection/ssd_model/pascal_label_map.pbtxt'
self.NUM_CLASSES = 20
self.detection_graph = self._load_model()
self.category_index = self._load_label_map()
def _load_model(self):
detection_graph = tf.Graph()
with detection_graph.as_default():
od_graph_def = tf.GraphDef()
with tf.gfile.GFile(self.PATH_TO_CKPT, 'rb') as fid:
serialized_graph = fid.read()
od_graph_def.ParseFromString(serialized_graph)
tf.import_graph_def(od_graph_def, name='')
return detection_graph
def _load_label_map(self):
label_map = label_map_util.load_labelmap(self.PATH_TO_LABELS)
categories = label_map_util.convert_label_map_to_categories(label_map,
max_num_classes=self.NUM_CLASSES,
use_display_name=True)
category_index = label_map_util.create_category_index(categories)
return category_index
def detect(self, image):
with self.detection_graph.as_default():
with tf.Session(graph=self.detection_graph) as sess:
# Expand dimensions since the model expects images to have shape: [1, None, None, 3]
image_np_expanded = np.expand_dims(image, axis=0)
image_tensor = self.detection_graph.get_tensor_by_name('image_tensor:0')
boxes = self.detection_graph.get_tensor_by_name('detection_boxes:0')
scores = self.detection_graph.get_tensor_by_name('detection_scores:0')
classes = self.detection_graph.get_tensor_by_name('detection_classes:0')
num_detections = self.detection_graph.get_tensor_by_name('num_detections:0')
# Actual detection.
(boxes, scores, classes, num_detections) = sess.run(
[boxes, scores, classes, num_detections],
feed_dict={image_tensor: image_np_expanded})
# Visualization of the results of a detection.
vis_util.visualize_boxes_and_labels_on_image_array(
image,
np.squeeze(boxes),
np.squeeze(classes).astype(np.int32),
np.squeeze(scores),
self.category_index,
use_normalized_coordinates=True,
line_thickness=8)
cv2.namedWindow("detection", cv2.WINDOW_NORMAL)
cv2.imshow("detection", image)
cv2.waitKey(0)
if __name__ == '__main__':
image = cv2.imread('/home/***/***/workspace/mobilenetssd/models-master/research/object_detection/ssd_model/img/2.jpg')
detecotr = TOD()
detecotr.detect(image)
识别结果如下(不是十分准确因为只训练了一万多次):
最后:自己做个数据集训练验证
- 首先需要先要标注图像相应标签,这里可以使用labelImg工具。每标注一张样本,即生成一个xml的标注文件。然后,把这些标注的xml文件,按训练集与验证集分别放置到两个目录下;
注意:linux下labelImg的下载可参考如下教程教程 - 在Datitran提供了xml_to_csv.py脚本。这里只要指定标注的目录名即可。接下来,需要我们把对应的csv格式转换成.record格式。
(1)首先编辑xml_to_csv.py,修改main函数:
def main():
#image_path = os.path.join(os.getcwd(), 'annotations')
image_path = os.path.join('/home/jari/guoshi/workspace/mobilenetssd/models-master/research/object_detection/ssd_model/data/train')
xml_df = xml_to_csv(image_path)
#xml_df.to_csv('raccoon_labels.csv', index=None)
xml_df.to_csv('gan_train_labels.csv', index=None)
print('Successfully converted xml to csv.')
执行:
python3 xml_to_csv.py
会生成train的csv,同样,修改代码,生成test的csv.
(2)进行csv到record的转换:
首先修改generate_tfrecord.py
,把main函数的path改成我们图片路径,然后把if row_label == 'raccoon'
:改成我们的label,比如gan.之后执行下面代码:
python3 generate_tfrecord.py --csv_input=gan_train_labels.csv --output_path=gan_train.record
python3 generate_tfrecord.py --csv_input=gan_test_labels.csv --output_path=gan_test.record
然后会生成对应的record文件:
- 回到之前训练SSD的目录,创建自己的label文件
gan_label_map.pbtxt
:
item {
id: 1
name: 'gan'
}
- 修改训练配置文件:
num_classes: 1 #20
再把所有PATH_TO_BE_CONFIGURED
的地方改掉,就可以用前面的命令进行训练和验证,如果想只进行训练采用object_detection/legacy/train.py
文件,而不是此处的model_main.py文件:
python3 object_detection/model_main.py \
--pipeline_config_path object_detection/ssd_model/ssd_mobilenet_v1_pets.config \
--model_dir=object_detection/training \
--num_train_steps=20000 \
--num_eval_steps=1000 \
--alsologtostderr
训练集很小所以loss下降的很快,但是对于测试集就不友好了。
.checkpoint
是保存各种权重,变量和计算图文件的总的记录文件;model.ckpt.data-00000-of-00001
和model.ckpt.index
是保存各种权重,变量的文件;.meta
是保存计算图结构的文件;.pbtxt
是保存输入种类的记录;.pb
文件是权重和图结构的总和,是一个固定的整体;