SavedModel - TFLite - SignatureDef - TensorInfo -获取中间层输出

import tensorflow as tf

# 定义一个简单的模型
class SimpleModel(tf.keras.Model):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.dense1 = tf.keras.layers.Dense(10, activation='relu')
        self.dense2 = tf.keras.layers.Dense(1)

    def call(self, inputs):
        x = self.dense1(inputs)
        return self.dense2(x)

# 创建模型实例
model = SimpleModel()

# 保存模型为SavedModel格式
tf.saved_model.save(model, 'saved_model')

import tensorflow as tf
from tensorflow.python.saved_model import signature_constants
from tensorflow.python.saved_model import tag_constants

# 加载保存的模型
loaded = tf.saved_model.load('saved_model')

# 获取模型的输入和中间层输出张量
input_tensor = loaded.signatures[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY].inputs['input_1']
dense1_output_tensor = loaded.signatures[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY].variables['dense1/Relu:0']

# 定义一个新的SignatureDef来获取中间层输出
signature_def = tf.compat.v1.saved_model.signature_def_utils.build_signature_def(
    inputs={'input_1': input_tensor},
    outputs={'dense1_output': dense1_output_tensor},
    method_name=signature_constants.PREDICT_METHOD_NAME
)

# 保存新的SignatureDef
with tf.compat.v1.Session(graph=tf.Graph()) as sess:
    tf.compat.v1.saved_model.loader.load(sess, [tag_constants.SERVING], 'saved_model')
    builder = tf.compat.v1.saved_model.builder.SavedModelBuilder('saved_model_with_intermediate')
    builder.add_meta_graph_and_variables(
        sess,
        [tag_constants.SERVING],
        signature_def_map={
            'serving_default': signature_def
        }
    )
    builder.save()

import tensorflow as tf

# 加载新的SavedModel
loaded = tf.saved_model.load('saved_model_with_intermediate')

# 获取SignatureDef
infer = loaded.signatures['serving_default']

# 准备输入数据
input_data = tf.constant([[1.0, 2.0, 3.0, 4.0]])

# 调用SignatureDef获取中间层输出
result = infer(tf.constant([[1.0, 2.0, 3.0, 4.0]]))['dense1_output']

print(result)