tensorflow导出冻结图模型

爱被打了一巴掌 2023-02-14 15:55 1阅读 0赞

### 导出推理图 ###

python export_inference_graph.py \
      --alsologtostderr \
      --dataset_name=autohome \  #自定义数据集的名称
      --model_name=mobilenet_v2_140 \ #模型名字
      --image_size=224 \
      --output_file=angle/mobilenet_v2_140_224_inference.pb #待输出的推理图

### 冻结模型文件 ###

python -m tensorflow.python.tools.freeze_graph \
       --input_graph=angle/mobilenet_v2_140_224_inference.pb \ #推理图
       --input_binary=true \
       --output_node_names=MobilenetV2/Predictions/Reshape_1 \ #根据不同模型确定各自的输出节点
       --input_checkpoint=angle/model/model.ckpt-10560 \  #训练的模型文件
       --output_graph=angle/model/frozen_angle8.pb   #导出的冻结图

### 使用冻结模型 ###

import tensorflow as tf
    import os
    import numpy as np
    import os, glob, cv2
    import sys, argparse
    
    # First, pass the path of the image
    dir_path = os.path.dirname(os.path.realpath(__file__))
    image_path = sys.argv[1]
    filename = dir_path + '/' + image_path
    image_size = 224
    num_channels = 3
    images = []
    
    ##  Reading the image using OpenCV
    image = cv2.imread(filename)
    
    ##   Resizing the image to our desired size and preprocessing will be done exactly as done during training
    image = cv2.resize(image, (image_size, image_size), cv2.INTER_LINEAR)
    images.append(image)
    images = np.array(images, dtype=np.uint8)
    images = images.astype('float32')
    images = np.multiply(images, 1.0 / 255.0)
    ##   The input to the network is of shape [None image_size image_size num_channels].
    ## Hence we reshape.
    
    x_batch = images.reshape(1, image_size, image_size, num_channels)
    
    frozen_graph = "./angle/model/frozen_angle8.pb"
    with tf.gfile.GFile(frozen_graph, "rb") as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
    
    with tf.Graph().as_default() as graph:
        tf.import_graph_def(graph_def,
                            input_map=None,
                            return_elements=None,
                            name=""
                            )
    ## NOW the complete graph with values has been restored
    y_pred = graph.get_tensor_by_name("MobilenetV2/Predictions/Reshape_1:0")
    ## Let's feed the images to the input placeholders
    x = graph.get_tensor_by_name("input:0")
    y_test_images = np.zeros((1, 2))
    sess = tf.Session(graph=graph)
    ### Creating the feed_dict that is required to be fed to calculate y_pred
    feed_dict_testing = {x: x_batch}
    result = sess.run(y_pred, feed_dict=feed_dict_testing)
    print(result)

测试模型结果。

python svc_on_freeze.py head45_l.jpg  #当前模型识别车辆的8个角度信息，加上背景信息供9个值
    
    [[0.00551218 0.00247121 0.8145236  0.04009005 0.01283834 0.01787371
      0.07377602 0.00181285 0.03110213]]

### 参考文献： ###

1.  [TensorFlow-Slim image classification model library][]
2.  [Freeze Tensorflow models and serve on web][]

[TensorFlow-Slim image classification model library]: https://github.com/tensorflow/models/tree/master/research/slim
[Freeze Tensorflow models and serve on web]: https://cv-tricks.com/how-to/freeze-tensorflow-models/