Tensorflow 使用slim框架下的分类模型进行分类

Tensorflow的slim框架可以写出像keras一样简单的代码来实现网络结构（虽然现在keras也已经集成在tf.contrib中了），而且models/slim提供了类似之前说过的object detection接口类似的image classification接口，可以很方便的进行fine-tuning利用自己的数据集训练自己所需的模型。

官方文档提供了比较详细的从数据准备，预训练模型的model zoo，fine-tuning，freeze model等一系列流程的步骤，但是缺少了inference的文档，不过tf所有模型的加载方式是通用的，所以调用方法和调用其他pb模型是一样的。

根据TF开发人员是说法Tensorflow对于模型读写的保存和调用的步骤一般如下：Build your graph --> write your graph --> import from written graph --> run compute etc。

以下我们使用slim提供的网络inception-resnet-v2作为例子：

1. export inference graph

import tensorflow as tf
import nets.inception_resnet_v2 as net
slim = tf.contrib.slim
# checkpoint path
checkpoint_path = "/your/path/to/inception_resnet_v2.ckpt" # ckpt file obtained during model training or fine-tuning
# set up and load session
sess = tf.Session()
arg_scope = net.inception_resnet_v2_arg_scope()
# initialize tensor suitable for model input
input_tensor = tf.placeholder(tf.float32, [None, 299, 299, 3])
with slim.arg_scope(arg_scope):
    logits, end_points = net.inception_resnet_v2(inputs=input_tensor)
# set up model saver
saver = tf.train.Saver()
saver.restore(sess, checkpoint_path)
with tf.gfile.GFile('/your/path/to/model_graph.pb', 'w') as f:   # save model to given pb file
    f.write(sess.graph_def.SerializeToString())
f.close()

2. freeze model

这里用tf提供的tensorflow/python/tools下的freeze_graph工具：

$ bazel build tensorflow/python/tools:freeze_graph
$ bazel-bin/tensorflow/python/tools/freeze_graph \
    --input_graph=/your/path/to/model_graph.pb  \  # obtained above
    --input_checkpoint=/your/path/to/inception_resnet_v2.ckpt \
    --input_binary=true
    --output_graph=/your/path/to/frozen_graph.pb \
    --output_node_names=InceptionResnetV2/Logits/Predictions   # output node name defined in inception resnet v2 net

(Optional) visualize frozen graph

LOG_DIR = ‘/tmp/graphdeflogdir’
model_filename = '/your/path/to/frozen_graph.pb'
with tf.Session() as sess:
    with tf.gfile.FastGFile(model_filename, 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        graph = tf.import_graph_def(graph_def, name='')
    writer = tf.summary.FileWriter(LOG_DIR, graph_def)
writer.close()

然后用tensorborad --logdir=LOG_DIR选择graph就可以查看到frozen后的网络结构。

3. inference

import cv2
import numpy as np
def preprocess_inception(image_np, central_fraction=0.875):
    image_height, image_width, image_channel = image_np.shape
    if central_fraction:
        bbox_start_h = int(image_height * (1 - central_fraction) / 2)
        bbox_end_h = int(image_height - bbox_start_h)
        bbox_start_w = int(image_width * (1 - central_fraction) / 2)
        bbox_end_w = int(image_width - bbox_start_w)
        image_np = image_np[bbox_start_h:bbox_end_h, bbox_start_w:bbox_end_w]
    # normalize
    image_np = 2 * (image_np / 255.) - 1
    return image_np
image_np = cv2.imread("test.jpg")
# preprocess image as inception resnet v2 does
image_np = preprcess_inception(image_np)
# resize to model input image size
image_np = cv2.resize(image_np, (299, 299))
# expand dims to shape [None, 299, 299, 3]
image_np = np.expand_dims(image_np, 0)
# load model
with tf.gfile.GFile('/your/path/to/frozen_graph.pb')
    graph_def = tf.GraphDef()
    graph_def.ParseFromString(f.read())
    graph = tf.import_graph_def(graph_def, name='')
    with tf.Session(graph=graph) as sess:
        input tensor = sess.graph.get_tensor_by_name("input:0")   # get input tensor
        output_tensor = sess.graph.get_tensor_by_name("InceptionResnetV2/Logits/Predictions:0")  # get output tensor
        logits = sess.run(output_tensor, feed_dict={input_tensor: image_np})
        print "Prediciton label index:", np.argmax(logits[0], 1)
        print "Top 3 Prediciton label index:", np.argsort(logits[0], 3)

参考: