cdxy.me
Footprints on Cyber Security and Python

SavedModel格式导出

In [1]:
from keras import backend as K
import tensorflow as tf   
from keras.models import load_model
tf.__version__
Using TensorFlow backend.
Out[1]:
'1.2.0'
  • 加载keras格式保存的模型
In [2]:
K.set_learning_phase(1)
In [3]:
model = load_model('build/lstm_extend')
In [4]:
model.input
Out[4]:
[<tf.Tensor 'main_input:0' shape=(?, 1024) dtype=int32>,
 <tf.Tensor 'feature_input:0' shape=(?, 306) dtype=float32>]
In [5]:
model.output
Out[5]:
<tf.Tensor 'main_output/Sigmoid:0' shape=(?, 1) dtype=float32>
In [6]:
legacy_init_op = tf.group(tf.tables_initializer())
In [7]:
with K.get_session() as sess:
#     sess.run(tf.global_variables_initializer())
#     sess.run(tf.local_variables_initializer())

    export_path = './saved_model/14'
    builder = tf.saved_model.builder.SavedModelBuilder(export_path)

    signature_inputs = {
        'main_input': tf.saved_model.utils.build_tensor_info(model.input[0]),
        'feature_input': tf.saved_model.utils.build_tensor_info(model.input[1]),
    }

    signature_outputs = {
        tf.saved_model.signature_constants.CLASSIFY_OUTPUT_CLASSES: tf.saved_model.utils.build_tensor_info(model.output)
    }

    classification_signature_def = tf.saved_model.signature_def_utils.build_signature_def(
        inputs=signature_inputs,
        outputs=signature_outputs,
        method_name=tf.saved_model.signature_constants.CLASSIFY_METHOD_NAME)

    builder.add_meta_graph_and_variables(
        sess,
        [tf.saved_model.tag_constants.SERVING],
        signature_def_map={
            'predict_webshell_php': classification_signature_def
        },
        legacy_init_op=legacy_init_op
    )

    builder.save()
INFO:tensorflow:No assets to save.
INFO:tensorflow:No assets to write.
INFO:tensorflow:SavedModel written to: ./saved_model/14/saved_model.pb
  • 测试已保存的模型
In [10]:
import numpy as np

f2 = np.array([0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1,0,0,0,0
,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,1,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0
,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0
,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0
,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0
,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0
,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0
,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0,0
,0,0,0,0,0,0,0,0,0,0]).reshape(1, -1)
f2
Out[10]:
array([[0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]])
In [11]:
f1=[0]*1020
f1.extend([14,2,56,8])
len(f1)
f1 = np.array(f1).reshape(1, -1)
f1
Out[11]:
array([[ 0,  0,  0, ...,  2, 56,  8]])
In [12]:
import tensorflow as tf  

export_dir = '/root/product/train/saved_model/14/'

with tf.Session(graph=tf.Graph()) as sess:
    meta_graph_def = tf.saved_model.loader.load(sess, [tf.saved_model.tag_constants.SERVING], export_dir)
    signature = meta_graph_def.signature_def

    x1_tensor_name = signature['predict_webshell_php'].inputs['main_input'].name
    x2_tensor_name = signature['predict_webshell_php'].inputs['feature_input'].name
    y_tensor_name = signature['predict_webshell_php'].outputs[tf.saved_model.signature_constants.CLASSIFY_OUTPUT_CLASSES].name

    x1 = sess.graph.get_tensor_by_name(x1_tensor_name)
    x2 = sess.graph.get_tensor_by_name(x2_tensor_name)
    y = sess.graph.get_tensor_by_name(y_tensor_name)
    print x1,x2,y

    print sess.run(y, feed_dict={x1:f1,x2:f2}) # 预测值
INFO:tensorflow:Restoring parameters from /root/product/train/saved_model/14/variables/variables
Tensor("main_input:0", shape=(?, 1024), dtype=int32) Tensor("feature_input:0", shape=(?, 306), dtype=float32) Tensor("main_output/Sigmoid:0", shape=(?, 1), dtype=float32)
[[ 0.99993169]]