Tensorflow SavedModel 模型的保存和加载

前言

最近参加了天池上的Apache Flink极客挑战赛——垃圾图片分类比赛，里面涉及到了Java调用tensorflow的SavedModel格式的模型进行预测，于是专门对此内容进行了调研。这里记录了SavedModel模型的优势，结构以及保存和加载的方法。

SavedModel的优势

Tensorflow训练的模型可以保存为ckpt格式，但是这种格式的模型文件在跨语言方面不是很灵活。而SaveModel与语言无关，比如可以使用python语言训练模型，然后在Java中非常方便的加载模型。

SavedModel的结构

以SavedModel格式保存模型时，tensorflow将创建一个SavedModel目录，该目录由以下子目录和文件组成：

assets/
assets.extra/
variables/
    variables.data-?????-of-?????
    variables.index
saved_model.pb|saved_model.pbtxt

其中，各目录和文件的说明如下：

• assets/是包含辅助（外部）文件（如词汇表）的子文件夹。资产被复制到SavedModel位置，并且可以在加载特定的MetaGraphDef时读取。
• assets.extra是一个子文件夹，高级库和用户可以将自己的资源添加进去，这些资源将与模型共存但不由图形加载。此子文件夹不由SavedModel库管理。
• variables/是包含tf.train.saver输出的子文件夹。
• saved_model.pb或saved_model.pbtxt是SavedModel协议缓冲区。它将图形定义作为MetaGraphDef协议缓冲区。

MetaGraph是一个数据流图，加上其相关的变量、assets和签名。MetaGraphDef是MetaGraph的Protocol Buffer表示。
assets/和assets.extra目录是可选的。

SavedModel模型的保存

这里以手写识别模型为例。

from tensorflow.examples.tutorials.mnist import input_data
import tensorflow as tf
from tensorflow.saved_model.signature_def_utils import predict_signature_def
from tensorflow.saved_model import tag_constants

mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)

sess = tf.InteractiveSession()
x = tf.placeholder(tf.float32, [None, 784], name="Input") # 为输入op添加命名"Input"
W = tf.Variable(tf.zeros([784, 10]))
b = tf.Variable(tf.zeros([10]))

y = tf.nn.softmax(tf.matmul(x, W) + b)
y_ = tf.placeholder(tf.float32, [None, 10])
cross_entropy = tf.reduce_mean(-tf.reduce_sum(y_ * tf.log(y), 1))
tf.identity(y, name="Output") # 为输出op命名为"Output"

train_step = tf.train.GradientDescentOptimizer(0.5).minimize(cross_entropy)
tf.global_variables_initializer().run()

for i in range(1000):
    batch_xs, batch_ys = mnist.train.next_batch(100)
    train_step.run({x: batch_xs, y_: batch_ys})

correct_prediction = tf.equal(tf.argmax(y, 1), tf.argmax(y_, 1))
accuracy = tf.reduce_mean(tf.cast(correct_prediction, tf.float32))

print(accuracy.eval({x: mnist.test.images, y_: mnist.test.labels}))

# 将模型保存到文件
# 简单方法：
tf.saved_model.simple_save(sess,
                           "./model_simple",
                           inputs={"Input": x},
                           outputs={"Output": y})
# 复杂方法
builder = tf.saved_model.builder.SavedModelBuilder("./model_complex")
signature = predict_signature_def(inputs={'Input': x},
                                  outputs={'Output': y})
builder.add_meta_graph_and_variables(sess=sess,
                                     tags=[tag_constants.SERVING],
                                     signature_def_map={'predict': signature})
builder.save()

代码解析：

• x = tf.placeholder(tf.float32, [None, 784], name="Input") # 为输入op添加命名"Input" 这里是为输入op进行命名，当然也可以不命名，系统会默认给一个名称"Placeholder"，当我们需要引用多个op的时候，给每个op一个命名，确实方便我们后面的使用。
• tf.identity(y, name="Output") # 为输出op命名为"Output" 使用tf.identity为输出tensor命名。
• 代码中给出了两种方法进行模型保存。复杂方法较简单方法的最大优势在于——可以自己定义tag，在签名的定义上更加灵活。

tag的作用: 一个模型可以包含不同的MetaGraphDef，比如你想保存graph的CPU版本和GPU版本，或者你想区分训练和发布版本。这个时候tag就可以用来区分不同的MetaGraphDef，加载的时候能够根据tag来加载模型的不同计算图。在simple_save方法中，系统会给一个默认的tag: “serve”，也可以用tag_constants.SERVING这个常量。

SavedModel模型的加载

Python

import numpy as np
mnist = input_data.read_data_sets("MNIST_data/", one_hot=True)

with tf.Session(graph=tf.Graph()) as sess:
    tf.saved_model.loader.load(sess, ["serve"], "./model_simple")
    graph = tf.get_default_graph()

    input = np.expand_dims(mnist.test.images[0], 0)
    x = sess.graph.get_tensor_by_name('Input:0')
    y = sess.graph.get_tensor_by_name('Output:0')
    batch_xs, batch_ys = mnist.test.next_batch(1)
    scores = sess.run(y,
                      feed_dict={x: batch_xs})
    print("predict: %d, actual: %d" % (np.argmax(scores, 1), np.argmax(batch_ys, 1)))

tf.saved_model.loader.load的第二个参数是定义的tag值，要和保存时定义保持一致；第三个参数是模型保存的路径。

Java

Java需要添加maven依赖：

        <dependency>
            <groupId>org.tensorflow</groupId>
            <artifactId>tensorflow</artifactId>
            <version>1.11.0</version>
        </dependency>
        <dependency>
            <groupId>org.tensorflow</groupId>
            <artifactId>proto</artifactId>
            <version>1.11.0</version>
        </dependency>

package com.garbage;

import org.tensorflow.SavedModelBundle;
import org.tensorflow.framework.ConfigProto;
import org.tensorflow.framework.MetaGraphDef;
import org.tensorflow.framework.SignatureDef;

/**
 * @Author: Jeremy
 * @Date: 2019/9/17 11:29
 */
public class SavedModelLoader {
    public static void main(String[] args) throws Exception{
        ConfigProto configProto = ConfigProto.newBuilder()
                .setAllowSoftPlacement(true)
                .build();
        SavedModelBundle model = SavedModelBundle.loader("YOUR_MODEL_PATH")
                .withConfigProto(configProto.toByteArray())
                .withTags("serve")
                .load();
        SignatureDef modelSig = MetaGraphDef.parseFrom(model.metaGraphDef()).getSignatureDefOrThrow("serving_default");
        int numInputs = modelSig.getInputsCount();
        String inputTensorName = modelSig.getInputsMap().get("Input").getName();
        String outputTensorName = modelSig.getOutputsMap().get("Output").getName();
        System.out.println(String.format("numInputs: %d, inputTensorName: %s, outputTensor: %s", numInputs, inputTensorName, outputTensorName));
    }
}

输出：

numInputs: 1, inputTensorName: Input_9:0, outputTensor: Softmax_10:0

使用Java预测的部分不是本文的重点，以后有机会再写。

总结

本篇博文总结了Tensorflow SaveModel模型，介绍了SaveModel模型的优势，目录结构，使用python进行模型保存，使用python和java加载模型。