基于tensorflow-slim实现分类任务

0. 环境准备

安装bazel 0.19.1 （本文测试版本)
Anaconda3 python 3.6.5 （本文测试版本)
CUDA10/cuDNN7.5.0 (本文测试版本)
显卡2070（本文测试版本）
tesorflow 1.13.1 （本文测试版本）

1. 基本配置

下载tensorflow源码tensorflow源码（切换到v1.12.0版本 — 本文测试版本），进入到tensorflow路径。执行以下命令编译graph_transforms：
bazel build tensorflow/tools/graph_transforms:transform_graph

创建文件夹Classfication用于存放训练的数据以及相关的脚本文件，这里先看下基本的目录结构以及相关文件信息：

下面先来解释下目录中的文件内容及作用。

checkpoints：用来存放预训练模型，这里我先现在了两个预训练模型
flower_photos：分类的数据集，其中是五种花的文件夹
train：该文件夹中将存储训练的record格式数据
val: 该文件夹中将存储验证的record格式数据
train_logs: 该文件夹用来保存训练时的模型文件
val_logs: 该文件夹用来保存验证时的log文件
feeze_model: 该文件夹用来保存模型图、冻结模型以及压缩模型
generate_label_map.py: 该脚本用来获图片路径与标签的关系,保存到txt文本中
split_traintxt_valtxt.py: 该脚本对generate_label_map.py生成的txt进行数据划分
createTFRecord.py: 该脚本将split_traintxt_valtxt.py划分的数据转成record格式，分别存储到train和val文件夹中

接下来看下三个脚本文件内容：
generate_label_map.py

import os
'''
生成list.txt来表示图片路径与标签的关系。
'''
class_names_to_ids = {'daisy': 0, 'dandelion': 1, 'roses': 2, 'sunflowers': 3, 'tulips': 4}
data_dir = './flower_photos'
output_path = 'list.txt'
with open(output_path, 'w') as f:
    for class_name in class_names_to_ids.keys():
        images_list = os.listdir(data_dir + os.sep + class_name)
        for image_name in images_list:
            f.write('{}/{} {}\n'.format(class_name, image_name, class_names_to_ids[class_name]))

split_traintxt_valtxt.py

import random
'''
随机生成训练集和验证集
'''
_NUM_VALIDATION = 500  # 这个验证集数量随意
_RANDOM_SEED = 0

list_path = 'list.txt'
train_list_path = 'list_train.txt'
val_list_path = 'list_val.txt'

with open(list_path, 'r') as f:
    lines = f.readlines()
random.seed(_RANDOM_SEED)
random.shuffle(lines)
with open(train_list_path, 'w') as f:
    for line in lines[_NUM_VALIDATION:]:
        f.write(line)
with open(val_list_path, 'w') as f:
    for line in lines[:_NUM_VALIDATION]:
        f.write(line)

createTFRecord.py

import sys
import os
'''
生成TFRecord数据
'''
#序号从0开始，表示最大优先级，sys.path.insert()加入的也是临时搜索路径，程序退出后失效.
sys.path.insert(0, '/xx/TF_Slim_classification' + os.sep)
from datasets import dataset_utils
import math
import os
import tensorflow as tf

def convert_dataset(list_path, data_dir, output_dir, _NUM_SHARDS=5):
    fd = open(list_path)
    lines = [line.split() for line in fd]
    fd.close()
    num_per_shard = int(math.ceil(len(lines) / float(_NUM_SHARDS)))
    with tf.Graph().as_default():
        decode_jpeg_data = tf.placeholder(dtype=tf.string)
        decode_jpeg = tf.image.decode_jpeg(decode_jpeg_data, channels=3)
        with tf.Session('') as sess:
            for shard_id in range(_NUM_SHARDS):
                output_path = os.path.join(output_dir,
                    'data_{:05}-of-{:05}.tfrecord'.format(shard_id, _NUM_SHARDS))
                tfrecord_writer = tf.python_io.TFRecordWriter(output_path)
                start_ndx = shard_id * num_per_shard
                end_ndx = min((shard_id + 1) * num_per_shard, len(lines))
                for i in range(start_ndx, end_ndx):
                    sys.stdout.write('\r>> Converting image {}/{} shard {}'.format(
                        i + 1, len(lines), shard_id))
                    sys.stdout.flush()
                    image_data = tf.gfile.FastGFile(os.path.join(data_dir, lines[i][0]), 'rb').read()
                    image = sess.run(decode_jpeg, feed_dict={decode_jpeg_data: image_data})
                    height, width = image.shape[0], image.shape[1]
                    example = dataset_utils.image_to_tfexample(
                        image_data, b'jpg', height, width, int(lines[i][1]))
                    tfrecord_writer.write(example.SerializeToString())
                tfrecord_writer.close()
    sys.stdout.write('\n')
    sys.stdout.flush()

os.system('mkdir -p train')
convert_dataset('list_train.txt', './flower_photos', 'train' + os.sep)
os.system('mkdir -p val')
convert_dataset('list_val.txt', './flower_photos', 'val' + os.sep)

注意三个脚本中相关文件的路径，都是按照上图中目录结构定的，自己根据实际情况做调整即可

2. 生成数据

依次执行下面三个脚本，然后就在train和val文件夹下生成用于训练和验证使用的record数据了。接着再创建一个labels.txt文件，内容如下：

这个文件是label标签文件，注意这里的label顺序要和generate_label_map.py脚本中的class_names_to_ids顺序一致！

3. 训练、验证、模型冻结、模型压缩

下载代码：TF_Slim_classification

说明：代码中sh文件中的“xx”替换为实际的路径，“xxx”替换为实际的训练epoch，“xxxx” 替换为实际需要训练的步数。
py脚本文件内容也很简单，这里就不详细说了，看下源码很好理解。

接下来就开始训练吧（默认网络为Inception_V4）

sh My_train_image_classifier.sh

训练过程中或者训练结束后可进行验证测试

sh My_eval_image_classifier.sh

训练或验证过程中可通过tensorboard观察训练或者验证情况：

tensorboard --logdir train_logs/
或
tensorboard --logdir val_logs/

训练结束后就需要导出模型图并进行冻结了，依次执行下面命令完成：

sh My_export_inference_graph.sh
sh My_frozen_inference_graph.sh

执行完后可在Classification/freeze_model路径下看到两个pb文件：

inception_v4_inf_graph.pb    模型图文件
frozen_graph_v4.pb           可用于推理的冻结模型文件

最后，对冻结模型进行量化压缩，具体的解释可参考官方说明 。进入tensorflow源码目录，执行以下命令即可：

	bazel-bin/tensorflow/tools/graph_transforms/transform_graph \
        --in_graph=/xxx/Classification/freeze_model/frozen_graph_v4.pb \
        --out_graph=/xxx/Classification/freeze_model/optimized_inception_graph.pb \
        --inputs='input' \
        --outputs='InceptionV4/Logits/Predictions' \
        --transforms='
	      add_default_attributes
	      strip_unused_nodes(type=float, shape="1,299,299,3")
	      remove_nodes(op=Identity, op=CheckNumerics)
	      fold_constants(ignore_errors=true)
	      fold_batch_norms
	      fold_old_batch_norms
	      quantize_weights
	      quantize_nodes
	      strip_unused_nodes
	      sort_by_execution_order'

这样在Classification/freeze_model路径下又得到一个名为optimized_inception_graph.pb 的压缩模型，这个模型和frozen_graph_v4.pb 差不多，区别还是见上面的官方解释。