在 Object Detection API 的示例代码中包含了一个训练识别宠物的 Demo,包括数据集和相应的一些代码。虽然本课程中我们会自己准备数据和脚本来进行训练,但是在这之前还需要安装一些库、配置一下环境。在配置完成之后,运行一下这个训练宠物的 Demo,以便检查环境配置是否 OK,同时对训练过程先有个整体的了解,然后再准备自己的数据和训练脚本。
请确保已经安装好了 Python 2.7-(我用的是Py35)。
安装 Object Detection API
首先下载 Object Detection API 的代码:
git clone https://github.com/tensorflow/models.git
修改代码(具体看-这篇文章):
把research/object_detection/utils/learning_schedules.py
文件的 第167-169行由
# # 修改167-169 rate_index = tf.reduce_max(tf.where(tf.greater_equal(global_step, boundaries), range(num_boundaries), [0] * num_boundaries)) # # 成 rate_index = tf.reduce_max(tf.where(tf.greater_equal(global_step, boundaries), list(range(num_boundaries)), [0] * num_boundaries))
然后安装 TensorFlow(我这里使用 tensorflow-gpu 1.4.1报错,API已经更新):
pip install tensorflow-gpu==1.8.0
接着是一些依赖库:
pip install pillow
pip install lxml
pip install jupyter
pip install matplotlib
pip install contextlib2
Object Detection API 中的模型和训练参数是使用 protobuf 来序列化和反序列化的,所以在运行之前需要将相应的 protobuf 文件编译出来。
#进入 tensorflow/models/research/
protoc object_detection/protos/*.proto --python_out=.
成功编译以后可以在 object_detection/protos/
下找到生成 .py 和 .pyc 文件。
接下来将 Object Detection API 的库加入到 PYTHONPATH 中:
#进入 tensorflow/models/research/
export PYTHONPATH=$PYTHONPATH:`pwd`:`pwd`/slim
运行 Object Detection API 的脚本,以及之后自己写的脚本都会用到这些库,如果不想每次运行前都敲这个命令的话,可以把这条命令加入到 ~/.bashrc 中(需要将 pwd 展开为实际路径):
# When running locally, the tensorflow/models/research/ and slim directories
export PYTHONPATH=$PYTHONPATH:/home/raini/pro/tf_models/research:/home/raini/pro/tf_models/research/slim
最后运行一下测试脚本来检测安装是否正确:
#进入 tensorflow/models/research/
python object_detection/builders/model_builder_test.py
如果看到下面的输出,那么 Object Detection API 的安装就完成了。
下载数据集
数据集由图片和相应的标注文件组成:
wget http://www.robots.ox.ac.uk/~vgg/data/pets/data/images.tar.gz
wget http://www.robots.ox.ac.uk/~vgg/data/pets/data/annotations.tar.gz
tar -xvf annotations.tar.gz
tar -xvf images.tar.gz
完成以后目录应该看起来是这样的:
images:
annotations:
在 images 目录就是一些宠物的照片,而在 annotations 文件夹里面是对相应照片的标注,在 annotations 文件夹中的和 images 文件夹中照片文件名一致的 xml 文件就是标注文件,这些标注文件为 PASCAL VOC 格式,可以打开 Abyssinian_1.xml
看一下:
标注内容主要为图片的源信息,如高和宽、物体的名称及所在位置:(xmin、ymin、xmax、ymax)所标识的矩形框。
还记得需要一个物体类别的数字编号和物体类别实际名称的对应关系的文件吗?可以在这里找到:
object_detection/data/pet_label_map.pbtxt
文件内容看起来是这样的:
注意:所有物体类别的数字编号都是从 1 开始的,因为 0 是一个在数学计算中很特殊的值。
生成 TFRecord 文件
Object Detection API 的训练框架使用 TFRecord 格式的文件作为输入。所以这里需要将图片和标注转换为 TFRecord 格式的文件。
TFRecord 数据文件是一种将图像数据和标签统一存储的二进制文件,能更好的利用内存,在 TensorFlow 中快速的复制、移动、读取、存储等。
Demo 里面包含了生成对应 TFRecord 格式文件的脚本,运行:
# 进入 tensorflow/models/research/ python object_detection/create_pet_tf_record.py \ --label_map_path=object_detection/data/pet_label_map.pbtxt \ --data_dir=DATA_DIR \ --output_dir=DATA_DIR
这里需要将 DATA_DIR
替换为 images 和 annotations 所在的文件夹(父文件夹),例如:
python object_detection/dataset_tools/create_pet_tf_record.py \ --label_map_path=object_detection/data/pet_label_map.pbtxt \ --data_dir=/home/byz/pro/tf_models/byz/pet_dataset \ --output_dir=/home/byz/pro/tf_models/byz/pet_dataset/output_dir
得到输出文件大致如下:
pet_train.record
为训练集,pet_val.record
为测试集。
准备转移学习
还需要一个 Pre-trained 模型来进行转移学习,尽量的缩短学习的时间,预训练模型文件可在这里下载所需要的。
下载以后解压备用:
在转移学习中要用的文件是 model.ckpt.* 这三个文件。
准备配置文件
还需要一个配置文件来对训练的流程进行配置,如使用什么算法,选用什么优化器等。在 object_detection/samples/configs/
可以找到很多配置模板,在这里使用 object_detection/samples/configs/ssd_mobilenet_v1_pets.config
作为起始的配置文件,需要在这个模板上面稍作修改。
这个配置文件是一个 JSON 格式的文件,里面有很多配置项,先挑一些必须修改的或者重要的项目:
train_input_reader: {
tf_record_input_reader {
input_path: "PATH_OF_TRAIN_TFRECORD" <--("byz/pet_dataset/output_dir/pet_faces_train.record"
}
label_map_path: "PATH_OF_LABEL_MAP" <--("byz/pet_dataset/pet_label_map.pbtxt"
}
需要将PATH_OF_TRAIN_TFRECORD
替换为pet_train.record
的绝对路径,将PATH_OF_LABEL_MAP
替换为pet_label_map.pbtxt
的绝对路径。
eval_input_reader: {
tf_record_input_reader {
input_path: "PATH_OF_VAL_TFRECORD" <--(byz/pet_dataset/output_dir/pet_faces_val.record"
}
label_map_path: "PATH_OF_LABEL_MAP"
}
需要将PATH_OF_VAL_TFRECORD
替换为pet_val.record
的绝对路径,将PATH_OF_LABEL_MAP
替换为pet_label_map.pbtxt
的绝对路径:
train_config: { fine_tune_checkpoint: "CHECK_POINT_PATH" <--("/home/byz/pro/tf_models/byz/faster_rcnn_resnet101_coco_11_06_2017/model.ckpt"
from_detection_checkpoint: true
num_steps: 200000
}
如果将from_detection_checkpoint
设为true
的话,代表将从一个事先训练好的模型开始继续训练(转移学习),此时需要将CHECK_POINT_PATH
替换为 model.ckpt 的绝对路径(注意之前有三个文件,model.ckpt.index、model.ckpt.meta、model.ckpt.data-xxx
在配置时不需要加model.ckpt
之后的后缀),如:fine_tune_checkpoint: "/home/byz/pro/tf_models/byz/faster_rcnn_resnet101_coco_11_06_2017/model.ckpt"
num_steps 为训练迭代的步数,这里暂时不修改。
将改好以后的配置文件重命名为 pipeline.config
。
开始训练
准备好训练数据和配置文件以后,就可以开始进行训练了。通常会把训练会用到的文件放到一起(训练目录),这里建议把训练目录设置为这样:
注意:需要按照这个目录结构修改 pipeline.config 中的相应项。
然后执行训练脚本:
# 进入 tensorflow/models/research/
python object_detection/train.py \
--logtostderr \
--pipeline_config_path=${TRAIN_DIR}/model/pipeline.config} \
--train_dir=${TRAIN_DIR}/model/train
TRAIN_DIR
需要替换为训练目录的绝对路径-checkpoint输出。
CUDA_VISIBLE_DEVICES="" python object_detection/train.py \ --logtostderr \ --pipeline_config_path=/media/raini/deepin/byz/dataset/pet-record/pipeline.config \ --train_dir=/home/raini/pro/tf_models/byz/train_checkpoint
模型开始训练,可以在终端上看到以下输出:
每一行输出为:训练迭代步数/当前损失值/每步训练所花时间。
基本上可以看出,随着训练的进行,每一步的损失值是下降的,那是不是可以喝咖啡等待训练结束了呢?
不过,每一步执行的时间大概在 10 秒左右,那么按照我们的配置 200000 步需要 200000 X 10 秒 = 23 天左右,这显然是不能接受的。
看来用笔记本的 CPU 进行训练可能不是一个好主意,需要更强的计算力:GPU。
配置好了训练环境,也把一个训练 Demo 运行了起来,但是笔记本的 CPU 运算能力显然不足应付这个任务,那么接下来让我们在 GPU 上面运行训练。
报错:
(如今的代码仅支持tf-1.5版本以上,所以需要升级版本,但是我的cuda是8.0,愁死我了)
Instructions for updating: Please switch to tf.train.create_global_step WARNING:tensorflow:num_readers has been reduced to 0 to match input file shards. Traceback (most recent call last): File "object_detection/train.py", line 184, in <module> tf.app.run() File "/home/byz/anaconda3/envs/py35/lib/python3.5/site-packages/tensorflow/python/platform/app.py", line 48, in run _sys.exit(main(_sys.argv[:1] + flags_passthrough)) File "object_detection/train.py", line 180, in main graph_hook_fn=graph_rewriter_fn) File "/home/byz/pro/tf_models/models/research/object_detection/trainer.py", line 274, in train train_config.prefetch_queue_capacity, data_augmentation_options) File "/home/byz/pro/tf_models/models/research/object_detection/trainer.py", line 59, in create_input_queue tensor_dict = create_tensor_dict_fn() File "object_detection/train.py", line 121, in get_next dataset_builder.build(config)).get_next() File "/home/byz/pro/tf_models/models/research/object_detection/builders/dataset_builder.py", line 186, in build process_fn, config.input_path[:], input_reader_config) File "/home/byz/pro/tf_models/models/research/object_detection/utils/dataset_util.py", line 137, in read_dataset tf.contrib.data.parallel_interleave( AttributeError: module 'tensorflow.contrib.data' has no attribute 'parallel_interleave'
升级了一下公司服务器的CUDA:
ubuntu 安装cuda9.0+cudnn7.1-与cuda8.0共存
训练会一直进行直到kill它。
在train_checkpoint目录下会看到model.ckpt-***, 需要将其转换成.pb文件。
# 使用官方的脚本转换成.pb模型 # From tensorflow/models/research/ PIPELINE_CONFIG_PATH = TRAIN_PATH = EXPORT_DIR = python object_detection/export_inference_graph.py \ --input_type image_tensor \ --pipeline_config_path "/home/raini/pro/tf_models/byz/trian_deepfasion_all/raini/deepfasion_nas/train_checkpoint/2299/pipeline.config" \ --trained_checkpoint_prefix "/home/raini/pro/tf_models/byz/trian_deepfasion_all/raini/deepfasion_nas/train_checkpoint/2299/model.ckpt-2299" \ --output_directory "/home/raini/pro/tf_models/byz/trian_deepfasion_all/raini/deepfasion_nas/train_checkpoint/output_pb"
使用模型检测图片
import tensorflow as tf # Imports import time start = time.time() import numpy as np import os import six.moves.urllib as urllib import sys import tarfile import zipfile import cv2 from collections import defaultdict # from io import StringIO # from matplotlib import pyplot as plt from PIL import Image if tf.__version__ < '1.4.0': raise ImportError('Please upgrade your tensorflow installation to v1.4.* or later!') os.chdir('/home/raini/pro/tf_models/models/research/object_detection') # Env setup # This is needed to display the images. # %matplotlib inline # This is needed since the notebook is stored in the object_detection folder. sys.path.append("..") # Object detection imports from utils import label_map_util from utils import visualization_utils as vis_util # Model preparation # What model to download. # 这是我们刚才训练的模型 # MODEL_NAME = 'tv_vehicle_inference_graph' MODEL_NAME = '/home/raini/pro/tf_models/byz/trian_deepfasion_all/raini/deepfasion_nas/train_checkpoint/output_pb' # 对应的Frozen model位置 # Path to frozen detection graph. This is the actual model that is used for the object detection. PATH_TO_CKPT = MODEL_NAME + '/frozen_inference_graph.pb' # List of the strings that is used to add correct label for each box. # PATH_TO_LABELS = os.path.join('training', 'detection.pbtxt') PATH_TO_LABELS = '/home/raini/pro/tf_models/byz/trian_deepfasion_all/raini/deepfasion_nas/test_3_label.pbtxt' # 改成自己例子中的类别数,2 NUM_CLASSES = 3 ''''' #Download Model 自己的模型,不需要下载了 opener = urllib.request.URLopener() opener.retrieve(DOWNLOAD_BASE + MODEL_FILE, MODEL_FILE) tar_file = tarfile.open(MODEL_FILE) for file in tar_file.getmembers(): file_name = os.path.basename(file.name) if 'frozen_inference_graph.pb' in file_name: tar_file.extract(file, os.getcwd()) ''' # Load a (frozen) Tensorflow model into memory. detection_graph = tf.Graph() with detection_graph.as_default(): od_graph_def = tf.GraphDef() with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid: serialized_graph = fid.read() od_graph_def.ParseFromString(serialized_graph) tf.import_graph_def(od_graph_def, name='') # Loading label map label_map = label_map_util.load_labelmap(PATH_TO_LABELS) categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True) category_index = label_map_util.create_category_index(categories) # Helper code def load_image_into_numpy_array(image): (im_width, im_height) = image.size return np.array(image.getdata()).reshape( (im_height, im_width, 3)).astype(np.uint8) # Detection # If you want to test the code with your images, just add path to the images to the TEST_IMAGE_PATHS. # 测试图片位置 PATH_TO_TEST_IMAGES_DIR = "/home/raini/pro/tf_models/byz/trian_deepfasion_all/raini/deepfasion_nas/test_img" os.chdir(PATH_TO_TEST_IMAGES_DIR) TEST_IMAGE_PATHS = os.listdir(PATH_TO_TEST_IMAGES_DIR) # Size, in inches, of the output images. IMAGE_SIZE = (12, 8) output_path = ('/home/raini/pro/tf_models/byz/trian_deepfasion_all/raini/deepfasion_nas/test_output/') with detection_graph.as_default(): with tf.Session(graph=detection_graph) as sess: # Definite input and output Tensors for detection_graph image_tensor = detection_graph.get_tensor_by_name('image_tensor:0') # Each box represents a part of the image where a particular object was detected. detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0') # Each score represent how level of confidence for each of the objects. # Score is shown on the result image, together with the class label. detection_scores = detection_graph.get_tensor_by_name('detection_scores:0') detection_classes = detection_graph.get_tensor_by_name('detection_classes:0') num_detections = detection_graph.get_tensor_by_name('num_detections:0') for image_path in TEST_IMAGE_PATHS: image = Image.open(image_path) # the array based representation of the image will be used later in order to prepare the # result image with boxes and labels on it. image_np = load_image_into_numpy_array(image) # Expand dimensions since the model expects images to have shape: [1, None, None, 3] image_np_expanded = np.expand_dims(image_np, axis=0) # Actual detection. (boxes, scores, classes, num) = sess.run( [detection_boxes, detection_scores, detection_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, np.squeeze(boxes), np.squeeze(classes).astype(np.int32), np.squeeze(scores), category_index, use_normalized_coordinates=True, line_thickness=8) # 保存文件 cv2.imwrite(output_path + image_path.split('\\')[-1], image_np) end = time.time() print("Execution Time: ", end - start)
最后的预测结果(换了份数据集)