H5模型转pb模型转tflite模型

from keras.models import load_model
import tensorflow as tf
import os 
import os.path as osp
from keras import backend as K
os.environ["CUDA_VISIBLE_DEVICES"] = "-1" 


    
def h5_to_pb(h5_model,output_dir,model_name,out_prefix = "yourname_",log_tensorboard = False):
    if osp.exists(output_dir) == False:
        os.mkdir(output_dir)
    out_nodes = []
    
    for i in range(len(h5_model.outputs)):
        out_nodes.append(out_prefix + str(i + 1))
        tf.identity(h5_model.output[i],out_prefix + str(i + 1))
    sess = K.get_session()
    from tensorflow.python.framework import graph_util,graph_io
    init_graph = sess.graph.as_graph_def()
    main_graph = graph_util.convert_variables_to_constants(sess,init_graph,out_nodes)
    graph_io.write_graph(main_graph,output_dir,name = model_name,as_text = False)
    if log_tensorboard:
        from tensorflow.python.tools import import_pb_to_tensorboard
        import_pb_to_tensorboard.import_to_tensorboard(osp.join(output_dir,model_name),output_dir)
    
    
if __name__ == '__main__':
    
    K.set_learning_phase(0)
    
    h5_model= load_model("./yourmodel.h5", compile=False)
    input = h5_model.input
    output = h5_model.output
    
    print('input', input)
    print('output', output)
    
    inputs=["one"]
    outputs=["two"]
    
    h5_to_pb(h5_model, './tflite_model', 'yourname.pb')

    #转换pb模型到tflite模型
    converter = tf.lite.TFLiteConverter.from_frozen_graph('./yourname.pb', inputs, outputs)
    converter.post_training_quantize = True
    converter.target_ops = [tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS]
    converter.allow_custom_ops=True
    #converter.optimizations = [tf.lite.Optimize.OPTIMIZE_FOR_SIZE]
    #converter.supported_types = [tf.float16]
    
    
    tflite_model = converter.convert()
    open('./yourname.tflite', "wb").write(tflite_model)
    


有可能会报错，类似于：
ValueError: Input 0 of node batch_normalization_1/cond/ReadVariableOp/Switch was passed float from batch_normalization_1/gamma:0 incompatible with expected resource.

需要再loadmodel之前调用：
K.set_learning_phase(0)

 

关注它，获取更多干货：