tensorflow checkpoint转savedmodel

import tensorflow as tf

def convert_model():
    trained_checkpoint_prefix = '/home/tiwang/code/jupyter-notebook/DIEN/dien/dnn_best_model/ckpt_noshuffDIEN3'
    export_dir = '/home/tiwang/code/jupyter-notebook/DIEN/dien/saved_model'
    graph = tf.Graph()
    config = tf.compat.v1.ConfigProto(allow_soft_placement=True, log_device_placement=True)
    with tf.compat.v1.Session(graph=graph, config=config) as sess:
        # restore from checkpoint
        loader = tf.compat.v1.train.import_meta_graph(trained_checkpoint_prefix+'.meta')
        loader.restore(sess, trained_checkpoint_prefix)
        
        # export checkpoint to savedmodel
        builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(export_dir)
        builder.add_meta_graph_and_variables(sess, [tf.saved_model.TRAINING, tf.saved_model.SERVING], strip_default_attrs=True)
        builder.save()

if __name__ == '__main__':
    # view_model_structure()
    convert_model()