python tensorflow将numpy dict文件转成pb文件

Posted 2021-05-09

def export_saved_model(network, sess, args):
    model_signature = saved_model.signature_def_utils.build_signature_def(
        inputs={
            "image_data": saved_model.utils.build_tensor_info(network.input_image),
            "image_info": saved_model.utils.build_tensor_info(network.input_image_info)
        },
        outputs={
            "cls_prob": saved_model.utils.build_tensor_info(network.net[0]),
            "bbox_pred": saved_model.utils.build_tensor_info(network.net[1]),
            "rois": saved_model.utils.build_tensor_info(network.net[2])
        },
        method_name=saved_model.signature_constants.PREDICT_METHOD_NAME
    )
    export_path = os.path.join(compat.as_bytes(args.saved_model_dir),
                               compat.as_bytes(args.saved_model_version))
    legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')
    builder = saved_model.builder.SavedModelBuilder(export_path)
    builder.add_meta_graph_and_variables(
        sess, [saved_model.tag_constants.SERVING],
        clear_devices=True,
        signature_def_map={
            saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:
                model_signature,
        },
        legacy_init_op=legacy_init_op
    )
    builder.save()