tensorflow使用python对pb模型做预估

灰太狼 2023-06-06 04:59 11阅读 0赞

tensorflow中ckpt模型转成pb模型的代码：参考链接[https://blog.csdn.net/dulingtingzi/article/details/90790282][https_blog.csdn.net_dulingtingzi_article_details_90790282]

但是为了使大家更容易明白，因为有些变量需要统一，这里针对下面的使用pb模型进行预估的代码，粘贴一下ckpt转pb模型：

import tensorflow as tf
    import os
    from tensorflow.contrib.layers import  flatten
    from tensorflow.python.framework import graph_util
    import tensorflow.contrib.slim as slim
    import numpy as np
    
    growth_rate = 6
    depth = 50
    compression = 0.5
    weight_decay = 0.0001
    nb_blocks = int((depth - 4) / 6)
    
    def dense_net(img_input, num_classes, nb_blocks, growth_rate, weight_decay, compression, flag):
       ##自定义densenet代码
        return densenet(growth_rate,img_input,num_classes,weight_decay,nb_blocks,compression,flag)
    
    def set_config():#设置GPU使用率# 控制使用率
        os.environ['CUDA_VISIBLE_DEVICES'] = '0'
        # 假如有16GB的显存并使用其中的8GB:
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.3)
        config = tf.ConfigProto(gpu_options=gpu_options)
        # session = tf.Session(config=config)
        return config
    #下面是你自定义的模型
    num_classes=2
    #is_training = tf.placeholder(tf.bool, name='placeholder_is_training')
    is_training = tf.constant(False, dtype=tf.bool)#下面的名字要和你一开始训练模型的时候是一致的
    inputs = tf.placeholder(tf.float32, shape=[None,30, 280, 3], name='placeholder_x')
    labels = tf.placeholder(tf.float32, shape=[None,num_classes], name='placeholder_y')
    pred=dense_net(inputs, num_classes,nb_blocks, growth_rate,weight_decay,compression,is_training)
    
    model_path="./version1/checkpoint/2_class.ckpt-1"#设置model的路径，因新版tensorflow会生成三个文件，只需写到数字前
    cfg=set_config()
    
    from tensorflow.python.saved_model import signature_constants, signature_def_utils, tag_constants, utils
    save_path = './version1/model_pb/test'
    with tf.Session(config=cfg) as sess:
        saver = tf.train.Saver()
        saver.restore(sess, model_path)
        print('ckpt loaded')
        #注意下面的inputs和outputs的名字要和后面的pb模型做预估保持一致
        model_signature = signature_def_utils.build_signature_def(inputs={"input": utils.build_tensor_info(inputs)},outputs={"pred": utils.build_tensor_info(pred)},method_name=signature_constants.PREDICT_METHOD_NAME)
        builder = tf.saved_model.builder.SavedModelBuilder(save_path)
        legacy_init_op = tf.group(tf.tables_initializer(), name='legacy_init_op')
        builder.add_meta_graph_and_variables(sess, [tag_constants.SERVING],clear_devices=True,signature_def_map={signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY:model_signature},legacy_init_op=legacy_init_op)
        builder.save()
        print('saved_model saved')

tensorflow使用pb模型做预估的代码：

import os
    import math
    import cv2
    import numpy as np
    import tensorflow as tf
    from tensorflow.python.saved_model import signature_constants, signature_def_utils, tag_constants, utils
    import matplotlib.pyplot as plt
    from time import time
    
    os.environ['CUDA_VISIBLE_DEVICES'] = '0'
    
    def preprocessing_crop_batch(images, height=30, width=280,depth=3):#按照batch去处理图像
        bs = len(images)
        GAUGE = height
        img_canvas = np.zeros([bs, height, width, depth], dtype=np.float32)
        #img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
        for i,img in enumerate(images):
           #####预处理代码块
        return img_canvas
    
    sess = tf.Session()
    m = tf.saved_model.loader.load(sess, tags=[tag_constants.SERVING], export_dir='./version1/model_pb/test/')
    graph = tf.get_default_graph()
    signature = m.signature_def
    signature_key = tf.saved_model.signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY
    input_tensor_name0 = signature[signature_key].inputs['input'].name#与之前转pb的时候的输入输出名字保持一致
    output_tensor_name = signature[signature_key].outputs['pred'].name
    x0 = tf.get_default_graph().get_tensor_by_name(input_tensor_name0)
    y0 = tf.get_default_graph().get_tensor_by_name(output_tensor_name)
    
    img_path = './images'
    imgs = os.listdir(img_path)
    imgs = list(map(lambda x : os.path.join(img_path, x), imgs))
    
    bs = 8
    img_batch = []
    for i in range(bs):
        img_batch.append(cv2.imread(imgs[i]))
    
    t0 = time()
    img_preprocess= preprocessing_crop_batch(img_batch)
    ans = sess.run(y0, {x0 : img_preprocess})
    ans = np.argmax(ans, axis=1)
    #ans = list(map(lambda x : x.decode(), ans))
    t1 = time()
    
    #plt.imshow(img[:,:,::-1])
    print(ans)
    print('seconds per frame %.2f ms' % ((t1-t0) * 1000 / bs))

[https_blog.csdn.net_dulingtingzi_article_details_90790282]: https://blog.csdn.net/dulingtingzi/article/details/90790282