keras模型转pb模型

import tensorflow as tf
from tensorflow.keras.models import load_model
import numpy as np

if __name__ == '__main__':
    norm_size=192
    graph = tf.Graph()
    with graph.as_default():
        session = tf.Session()
        with tf.Session().as_default():
            input = tf.placeholder(tf.float32, [None, 192, 192, 3], name='inputs')
            model = load_model('../model/192-v69.m', compile=False)
            output = tf.identity(model(input), name='outputs')
            ops = [tf.assign(v, w) for v, w
                   in zip(tf.global_variables(), model.get_weights())]
            session.run(ops)
            frozen = tf.graph_util.convert_variables_to_constants(
                session, graph.as_graph_def(), ['outputs'])
            # input: 'conv2d_1_input:0'
            with open('../model/192-v69.pb', 'wb') as f:
                f.write(frozen.SerializeToString())

    graph = tf.Graph()
    with graph.as_default():
        with open('../model/192-v69.pb', 'rb') as f:
            pb = tf.GraphDef.FromString(f.read())
        tf.import_graph_def(pb, name='')
        session = tf.Session()
        func1 = session.make_callable('outputs:0', ['inputs:0'])
    kg = tf.Graph()
    with kg.as_default():
        session2 = tf.Session()
        with session2.as_default():
            keras = load_model('../model/192-v69.m', compile=False)
            # keras.predict(np.zeros([1,192,192,3]))
            func2 = keras.predict

    testdata = np.random.uniform(0., 1., [1,192,192,3]).astype(np.float32)
    testoutput1 = func1(testdata)
    with kg.as_default():
        with session2.as_default():
            testoutput2 = func2(testdata)
    print(testoutput1)
    print(testoutput2)
    assert np.allclose(testoutput1, testoutput2)

 

你可能感兴趣的:(TensorFlow)