：）模型保存为单一个pb文件

模型保存为单一个pb文件

背景

参考连接： https://www.yuque.com/g/jesse-ztr2k/nkke46/ss4rlv/collaborator/join?token=XUVZNORisVWEWyst# 

注意有些时候需要添加一个pb文件。 而不是tensorflow 提供的save 方法生成的一个目录里面包含了若干pb文件。

load时候直接填写这个目录即可。 但是有些时候需要合成一个pb文件。

 

tf2生成pb 目录描述

　　1 目录结构

 -assets

-variables

-variables.data-00000-of-00001

-variables.index

-saved_model.pb

　　2 作用

　　　　其中 variables 记录模型参数 ， pb文件记录模型结构

tf2 都是保存的 权重和 结构分开的， 如果需要兼容tf V1的代码，即导入一个pb文件，就需要 1 ）保存常量计算图 2）frozen graph  pb格式。

tf1 生成pb脚本

环境准备：

tensorflow==1.15, tf-slim==1.1.0

https://github.com/tensorflow/models/tree/master/research/slim

注意 一定在tf v1 环境下生成pb

import cv2
import numpy as np
import tensorflow as tf
import os
from tensorflow.python.framework import graph_util

# 参考连接 https://blog.csdn.net/tensorflowforum/article/details/112352764 代码
# 参考连接 参数详解：https://blog.csdn.net/weixin_43529465/article/details/124721583
# https://blog.csdn.net/rain6789/article/details/78754516

class SingleCnn(tf.keras.Model):
    def __init__(self):
        super(SingleCnn, self).__init__()
        # filters=1 卷积核数目，相当于卷积核的channel
        self.conv = tf.keras.layers.Conv2D(filters=1,
                                           kernel_size=[1, 1],
                                           # valid表示不填充， same表示合理填充
                                           padding='valid',
                                        # data_format='channels_last',-> 表示HWC，输入可以定义批次
                                           data_format='channels_last',
                                           use_bias=False,
                                           kernel_initializer=tf.keras.initializers.he_uniform(seed=None),
                                           name="conv")

    def call(self, inputs):
        x = self.conv(inputs)
        return x
if __name__ == "__main__":
    # 构建场景输入数据

    # images=tf.random.uniform((1, 300, 300, 3))

    # 图像数据
    imagefile = r"catanddog\cat\5.JPG"
    img = cv2.imread(imagefile)
    img = cv2.resize(img, (64, 64))
    img = np.expand_dims(img, axis=0)
    print(img.shape, type(img), img.dtype)

    # 未量化的model不支持int32和int8
    # img = img.astype(np.int32)
    img = tf.convert_to_tensor(img, np.float32)
    print(img.shape, type(img), img.dtype)
    singlecnn = SingleCnn()

    output = singlecnn(img)
    print(output.shape, type(output))
    print(output[0][2:10][2:6])
    # =========== ckpt保存 with session的写法tf2 已不再使用 ===========
    # with tf.Session(graph=tf.Graph()) as sess:
    #     constant_graph = graph_util.convert_variables_to_constants(sess, sess.graph_def, ['op_to_store'])

    # 保存参考 https://zhuanlan.zhihu.com/p/146243327
    # save_format='tf' 代表保存pb
    # singlecnn.save('./pbmodel/singlecnn', save_format='tf')
    # tf.saved_model.save(singlecnn, './pbmodel/singlecnn')
    tf.keras.models.save_model(singlecnn, './pbmodel/singlecnn_0',
                               save_format="tf",
                               include_optimizer=False, save_traces=False)

    # 加载模型 验证可以加载
    new_model = tf.keras.models.load_model('./pbmodel/singlecnn_0', compile=False)
    # new_model = tf.saved_model.load('./pbmodel/singlecnn_0')
    # output_ = new_model(img)
    # # print(output_.shape, output_[0][2:6][2:6])
    # print(output_.shape)
    #
    # 查看结构
    new_model.summary()

    # print("----------------")
    # # 加载模型
    # saved_model = tf.saved_model.load('./pbmodel/singlecnn_0')
    # # 将模型转换为pb格式    还是目录方法。
    # converter = tf.saved_model.save(saved_model, "model.pb")

    def change_pb(pretrained_model):
        """tf v1 选用tf1 跑这个脚本生成pb"""
        from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
        # 重点
        # Convert Keras model to ConcreteFunction
        # MobileNet is a function
        full_model = tf.function(lambda x: pretrained_model(x))

        # 指定shape和dtype对tf function进行重新追踪
        full_model = full_model.get_concrete_function(
            tf.TensorSpec(pretrained_model.inputs[0].shape, pretrained_model.inputs[0].dtype))

        # Get frozen ConcreteFunction，将计算图中的变量及其取值通过常量的方式保存
        frozen_func = convert_variables_to_constants_v2(full_model)
        frozen_func.graph.as_graph_def()

        layers = [op.name for op in frozen_func.graph.get_operations()]
        print("-" * 50)
        print("Frozen model layers: ")
        for layer in layers:
            print(layer)

        print("-" * 50)
        print("Frozen model inputs: ")
        print(frozen_func.inputs)
        print("Frozen model outputs: ")
        print(frozen_func.outputs)

        # Save frozen graph from frozen ConcreteFunction to hard drive
        # as_text: If True, writes the graph as an ASCII proto; otherwise, The graph is written as a text proto
        tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
                          logdir="./frozen_models",
                          name="frozen_graph.pb",
                          as_text=True)


change_pb(new_model)

python download_and_convert_data.py --dataset_name=flowers --dataset_dir="tmp/dataset"