Plumber可以解析的模型和模型适配例子



  • 如何判断Plumber可以解析哪些模型

    目前Plumber支持解析的模型是基于TensorFlow训练出来的CNN模型,在这个基础上,TensorFlow训练的模型一般保存为Checkpoints,具体请看 https://www.tensorflow.org/guide/checkpoints。我们称之为检查点文件,这类检查点文件并不能直接用于Plumber解析,因为这种模型格式依赖于创建模型的代码,我们需要做些处理让它适配一下。而其他方式保存的模型目前我们还不能解析。

    以下是例子:

    我们是用MNIST数据集为例。

    定义网络

    def mnist_model_fn(inputs, mode):
        """Model function for CNN."""
    
        with tf.variable_scope("Lenet"):
    
            input_layer = tf.reshape(inputs, [-1, 28, 28, 1])
            
            with tf.variable_scope("conv1"):
                conv1 = tf.layers.conv2d(inputs=input_layer, filters=32, kernel_size=[5, 5], padding="same", activation=tf.nn.relu)
    
            with tf.variable_scope("pool1"):
                pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2, 2], strides=2)
            
            with tf.variable_scope("conv2"):
                conv2 = tf.layers.conv2d(inputs=pool1, filters=64, kernel_size=[5, 5], padding="same", activation=tf.nn.relu)
            
            with tf.variable_scope("pool2"):
                pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=2)
    
            pool2_flat = tf.reshape(pool2, [-1, 7 * 7 * 64])
    
            with tf.variable_scope("fc1"):
                dense = tf.layers.dense(inputs=pool2_flat, units=1024, activation=tf.nn.relu)
            
            dropout = tf.layers.dropout(inputs=dense, rate=0.4, training=mode == tf.estimator.ModeKeys.TRAIN)
            
            with tf.variable_scope("fc2"):
                logits = tf.layers.dense(inputs=dropout, units=10)
    
        return logits
    

    定义网络的损失和评估函数,使用的tf.estimator.EstimatorSpec

    def cnn_model_fn(features, labels, mode):
        
        logits = mnist_model_fn(features['x'], mode)
    
        predictions = {
            # Generate predictions (for PREDICT and EVAL mode)
            "classes": tf.argmax(input=logits, axis=1),
            # Add `softmax_tensor` to the graph. It is used for PREDICT and by the
            # `logging_hook`.
            "probabilities": tf.nn.softmax(logits, name="softmax_tensor")
        }
    
        if mode == tf.estimator.ModeKeys.PREDICT:
            return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
    
        # Calculate Loss (for both TRAIN and EVAL modes)
        loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)
    
        # Configure the Training Op (for TRAIN mode)
        if mode == tf.estimator.ModeKeys.TRAIN:
            optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
            train_op = optimizer.minimize(
                        loss=loss,
                        global_step=tf.train.get_global_step())
            return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
    
        # Add evaluation metrics (for EVAL mode)
        eval_metric_ops = {
                "accuracy": tf.metrics.accuracy(
                            labels=labels,
                            predictions=predictions["classes"])
                }
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)
    
    

    训练网络代码

    这里涉及到下载训练数据的,有时候下载比较慢,我们可以直接在浏览器下载好之后,然后使用第二种方式读取数据。

    MODEL_DIR是保存检查点文件的目录,到时换成你自己的路径。目前我设置的是训练迭代2000步,根据需要修改。

    
    MODEL_DIR = '~/tmp'
    
    def main(unused_argv):
        # Load training and eval data
        # 在线下载
        mnist = tf.contrib.learn.datasets.load_dataset("mnist")
        # 可以读入下载好的文件
        #mnist = input_data.read_data_sets("./MNIST-data/", one_hot=True)
        train_data = mnist.train.images # Returns np.array
        train_labels = np.asarray(mnist.train.labels, dtype=np.int32)
        eval_data = mnist.test.images # Returns np.array
        eval_labels = np.asarray(mnist.test.labels, dtype=np.int32)
    
        # Create the Estimator
        mnist_classifier = tf.estimator.Estimator(model_fn=cnn_model_fn, model_dir=MODEL_DIR)
    
        # Set up logging for predictions
        tensors_to_log = {"probabilities": "softmax_tensor"}
        logging_hook = tf.train.LoggingTensorHook(tensors=tensors_to_log, every_n_iter=50)
    
        # Train the model
        train_input_fn = tf.estimator.inputs.numpy_input_fn(
                        x={"x": train_data},
                        y=train_labels,
                        batch_size=32,
                        num_epochs=None,
                        shuffle=True)
        mnist_classifier.train(
                        input_fn=train_input_fn,
                        steps=2000,
                        hooks=[logging_hook])
    
        # Evaluate the model and print results
        eval_input_fn = tf.estimator.inputs.numpy_input_fn(
                        x={"x": eval_data},
                        y=eval_labels,
                        num_epochs=1,
                        shuffle=False)
        eval_results = mnist_classifier.evaluate(input_fn=eval_input_fn)
    
        print(eval_results)
    

    最后是我们的适配代码

    它通过加载之前训练好之后的检查点文件,然后进行图推理,只是改变了输入的形状,这一步是至关重要的,因为很多模型训练时它的输入一个以队列,或者其他的方式,它都是针对大量的训练数据的,而Plumber解析依赖的输入是tf.placeholder这样的输入,并且如果有batch size,应该把它置为1。因为我们不需要输入批量的数据了。

    def restore(ckpt_filename, output_path):
      """模型适配
      Args:
      ckpt_filename: 检查点文件路径,如model.ckpt-2000形式
      output_path:新生成的检查点文件,如mnist.ckpt形式
      """
      # 我们重新构造输入,并且设置为batch size 为1
      x = tf.placeholder(tf.float32, shape=[1, 28, 28, 1], name="img")
        
      logits = mnist_model_fn(x, False)
    
      with tf.Session() as sess:
          saver = tf.train.Saver()
          saver.restore(sess, ckpt_filename)
          saver.save(sess, output_path)
    

    模型使用建议

    因为Plumber目前只支持CNN网络模型,并且目前网上已经训练好的检查点文件需要根据实际情况来操作,不能直接在Plumber上使用,所以我们建议如果你想尝试,可以尝试有模型定义代码的模型,然后可以根据我们这个例子操作得到适用于Plumber的模型。

    完整代码

    from __future__ import absolute_import
    from __future__ import division
    from __future__ import print_function
    
    import numpy as np
    import tensorflow as tf
    from tensorflow.examples.tutorials.mnist import input_data
    
    
    tf.logging.set_verbosity(tf.logging.INFO)
    
    MODEL_DIR = "~/mnist_letnet"
    
    def mnist_model_fn(inputs, mode):
        """Model function for CNN."""
    
        with tf.variable_scope("Lenet"):
    
            input_layer = tf.reshape(inputs, [-1, 28, 28, 1])
            
            with tf.variable_scope("conv1"):
                conv1 = tf.layers.conv2d(inputs=input_layer, filters=32, kernel_size=[5, 5], padding="same", activation=tf.nn.relu)
    
            with tf.variable_scope("pool1"):
                pool1 = tf.layers.max_pooling2d(inputs=conv1, pool_size=[2, 2], strides=2)
            
            with tf.variable_scope("conv2"):
                conv2 = tf.layers.conv2d(inputs=pool1, filters=64, kernel_size=[5, 5], padding="same", activation=tf.nn.relu)
            
            with tf.variable_scope("pool2"):
                pool2 = tf.layers.max_pooling2d(inputs=conv2, pool_size=[2, 2], strides=2)
    
            pool2_flat = tf.reshape(pool2, [-1, 7 * 7 * 64])
    
            with tf.variable_scope("fc1"):
                dense = tf.layers.dense(inputs=pool2_flat, units=1024, activation=tf.nn.relu)
            
            dropout = tf.layers.dropout(inputs=dense, rate=0.4, training=mode == tf.estimator.ModeKeys.TRAIN)
            
            with tf.variable_scope("fc2"):
                logits = tf.layers.dense(inputs=dropout, units=10)
    
        return logits
    
    def cnn_model_fn(features, labels, mode):
        
        logits = mnist_model_fn(features['x'], mode)
    
        predictions = {
            # Generate predictions (for PREDICT and EVAL mode)
            "classes": tf.argmax(input=logits, axis=1),
            # Add `softmax_tensor` to the graph. It is used for PREDICT and by the
            # `logging_hook`.
            "probabilities": tf.nn.softmax(logits, name="softmax_tensor")
        }
    
        if mode == tf.estimator.ModeKeys.PREDICT:
            return tf.estimator.EstimatorSpec(mode=mode, predictions=predictions)
    
        # Calculate Loss (for both TRAIN and EVAL modes)
        loss = tf.losses.sparse_softmax_cross_entropy(labels=labels, logits=logits)
    
        # Configure the Training Op (for TRAIN mode)
        if mode == tf.estimator.ModeKeys.TRAIN:
            optimizer = tf.train.GradientDescentOptimizer(learning_rate=0.001)
            train_op = optimizer.minimize(
                        loss=loss,
                        global_step=tf.train.get_global_step())
            return tf.estimator.EstimatorSpec(mode=mode, loss=loss, train_op=train_op)
    
        # Add evaluation metrics (for EVAL mode)
        eval_metric_ops = {
                "accuracy": tf.metrics.accuracy(
                            labels=labels,
                            predictions=predictions["classes"])
                }
        return tf.estimator.EstimatorSpec(mode=mode, loss=loss, eval_metric_ops=eval_metric_ops)
    
    
    def main(unused_argv):
        # Load training and eval data
        mnist = tf.contrib.learn.datasets.load_dataset("mnist")
        #mnist = input_data.read_data_sets("./MNIST-data/", one_hot=True)
        train_data = mnist.train.images # Returns np.array
        train_labels = np.asarray(mnist.train.labels, dtype=np.int32)
        eval_data = mnist.test.images # Returns np.array
        eval_labels = np.asarray(mnist.test.labels, dtype=np.int32)
    
        # Create the Estimator
        mnist_classifier = tf.estimator.Estimator(model_fn=cnn_model_fn, model_dir=MODEL_DIR)
    
        # Set up logging for predictions
        tensors_to_log = {"probabilities": "softmax_tensor"}
        logging_hook = tf.train.LoggingTensorHook(tensors=tensors_to_log, every_n_iter=50)
    
        # Train the model
        train_input_fn = tf.estimator.inputs.numpy_input_fn(
                        x={"x": train_data},
                        y=train_labels,
                        batch_size=32,
                        num_epochs=None,
                        shuffle=True)
        mnist_classifier.train(
                        input_fn=train_input_fn,
                        steps=2000,
                        hooks=[logging_hook])
    
        # Evaluate the model and print results
        eval_input_fn = tf.estimator.inputs.numpy_input_fn(
                        x={"x": eval_data},
                        y=eval_labels,
                        num_epochs=1,
                        shuffle=False)
        eval_results = mnist_classifier.evaluate(input_fn=eval_input_fn)
    
        print(eval_results)
    
    def restore(ckpt_filename, output_path):
        x = tf.placeholder(tf.float32, shape=[1, 28, 28, 1], name="img")
        
        logits = mnist_model_fn(x, False)
    
        with tf.Session() as sess:
            saver = tf.train.Saver()
            saver.restore(sess, ckpt_filename)
            saver.save(sess, output_path)
    
    if __name__ == "__main__":
        tf.app.run()
        #restore("~/mnist_letnet/model.ckpt-2000", "~/mnist_letnet/new/mnist.ckpt')
    

登录后回复