18.3:tensorflow分类模型mobilenetv2训练(数据增强,保存模型,衰减学习率,tensorboard)，预测图像(单张,批量预测)，导出为pb完整示例

三、生成pb文件并预测。

他查阅资料发现，生成pb文件需要三个文件，也就是model_save中保存的。如：

此外，还需要知道输入节点，输出节点的名称。然而在他的训练的代码中并不好找到哪个是输入节点，在tensorboard的结构图中也没有找到。

但他又想了一个方法，从保存的模型中再加载一次模型，并建立一个输入，输出节点，然后再保存。于是他写了这个ckpt_to_pd.py文件：

save_new_ckpt函数可以加入新的输入节点(input_1)，输出节点(softmax_out)。

把model.ckpt-250000模型加载后，重新编辑加入输入输出节点后，再保存为modelnew.ckpt，会生成3个新文件。再使用这3个文件根据函数freeze_graph生成pb文件。

print_pb函数可以打印生成的pb文件中参数名字及尺寸。

#coding:utf-8

import tensorflow as tf
from tensorflow.python.framework import graph_util

import model


label_dict, label_dict_res = {}, {}
# 手动指定一个从类别到label的映射关系
with open("label.txt", 'r') as f:
    for line in f.readlines():
        folder, label = line.strip().split(':')[0], line.strip().split(':')[1]
        label_dict[folder] = label
        label_dict_res[label] = folder
print(label_dict)

IMG_W = 224
N_CLASSES = len(label_dict)

def print_pb(output_graph_path):
    tf.reset_default_graph()  # 重置计算图
    sess = tf.Session()
    sess.run(tf.global_variables_initializer())
    output_graph_def = tf.GraphDef()
    # 获得默认的图
    graph = tf.get_default_graph()
    with open(output_graph_path, "rb") as f:
        output_graph_def.ParseFromString(f.read())
        _ = tf.import_graph_def(output_graph_def, name="")
        # 得到当前图有几个操作节点
        print("%d ops in the final graph." % len(output_graph_def.node))
        #tensor_name = [tensor.name for tensor in output_graph_def.node]
        #print(tensor_name)
        print('---------------------------')
        for op in graph.get_operations():
            # print出tensor的name和值
            print(op.name, op.values())
    sess.close()


def save_new_ckpt(logs_train_dir, newckpt):
    x = tf.placeholder(tf.float32, shape=[None, IMG_W, IMG_W, 3], name="input_1")
    # predict
    logit = model.MobileNetV2(x, num_classes=N_CLASSES, is_training=False).output
    #logit = model.model4(x, N_CLASSES, is_trian=False)
    #logit = model.model2(x_4d, batch_size=1, n_classes=N_CLASSES)
    pred = tf.nn.softmax(logit, name="softmax_out")

    saver = tf.train.Saver()
    sess = tf.Session()
    saver.restore(sess, logs_train_dir)
    print("load model done...")
    saver.save(sess, newckpt)
    sess.close()
    print('save new model done...')


def freeze_graph(input_checkpoint, output_graph):
    # 指定输出的节点名称,该节点名称必须是原模型中存在的节点
    output_node_names = ["input_1", "softmax_out"]
    saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)
 
    with tf.Session() as sess:
        saver.restore(sess, input_checkpoint) #恢复图并得到数据
        output_graph_def = graph_util.convert_variables_to_constants(  # 模型持久化，将变量值固定
            sess=sess,
            input_graph_def=sess.graph_def, # 等于:sess.graph_def
            output_node_names=output_node_names)
 
        with tf.gfile.GFile(output_graph, "wb") as f: #保存模型
            f.write(output_graph_def.SerializeToString()) #序列化输出
        print("%d ops in the final graph." % len(output_graph_def.node)) #得到当前图有几个操作节点


if __name__ == "__main__":
    model_path = './model_save/model.ckpt-250000'
    new_model_path = './model_save/modelnew.ckpt'
    pb_model = "model_tf.pb"
    save_new_ckpt(model_path, new_model_path)
    freeze_graph(new_model_path, pb_model)
    print_pb(pb_model)

使用.pb模型进行预测：

有了model_tf.pb文件后，也知道了输入，输出节点了。之后就可以无需定义网络结构就可以直接进行预测了，使用比较方便。

调用pb文件并进行预测的代码为predict_pb.py：

#coding:utf-8
import os, cv2
import numpy as np
import tensorflow as tf
import glob

from tensorflow.python.platform import gfile


config = tf.ConfigProto()
config.gpu_options.allow_growth = True
os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # use gpu 0

label_dict, label_dict_res = {}, {}
with open("label.txt", 'r') as f:
    for line in f.readlines():
        folder, label = line.strip().split(':')[0], line.strip().split(':')[1]
        label_dict[folder] = label
        label_dict_res[label] = folder
print(label_dict)

N_CLASSES = len(label_dict)
IMG_W = 224
IMG_H = IMG_W

def init_tf():
    global sess, pred, x
    sess = tf.Session(config=config)
    with gfile.FastGFile('./model_tf.pb', 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        sess.graph.as_default()
        tf.import_graph_def(graph_def, name='')

    # 获取输入tensor
    x = tf.get_default_graph().get_tensor_by_name("input_1:0")
    print("input:", x)
    # 获取预测tensor
    pred = tf.get_default_graph().get_tensor_by_name("softmax_out:0")  #mobilenet_v2
    print('load model done...')

def evaluate_image(img_dir):
    # read and process image
    image = cv2.imread(img_dir)
    image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
    # pre-process the image for classification
    image = cv2.resize(image, (IMG_W, IMG_W))
    #image = image.astype("float") / 255.0
    im_mean = np.mean(image)
    stddev = max(np.std(image), 1.0/np.sqrt(IMG_W*IMG_W*3))
    image = (image - im_mean) / stddev #代替tf.image.per_image_standardization
    #image_array = np.array(image)
    image = np.expand_dims(image, axis=0)

    prediction = sess.run(pred, feed_dict={x: image})
    max_index = np.argmax(prediction)
    pred_label = label_dict_res[str(max_index)]
    print("%s, predict: %s(index:%d), prob: %f" %(img_dir, pred_label, max_index, prediction[0][max_index]))
    

if __name__ == '__main__':
    init_tf()
    data_path = "/media/DATA2/sku_val"
    label = os.listdir(data_path)
    for l in label:
        if os.path.isfile(os.path.join(data_path, l)):
            continue
        for img in glob.glob(os.path.join(data_path, l, "*.jpg")):
            evaluate_image(img_dir=img)
    sess.close()

至于如何修改为批量预测，根据上一个博客18.2修改把。。。

利用混淆矩阵计算准确率和召回

#coding:utf-8
import os, cv2
import numpy as np
import shutil
import glob
import tqdm
import matplotlib.pyplot as plt
from tensorflow.python.platform import gfile
import tensorflow as tf


os.environ["CUDA_VISIBLE_DEVICES"] = "0"  # use gpu 0
config = tf.ConfigProto()
#config.gpu_options.allow_growth = True
config.gpu_options.per_process_gpu_memory_fraction = 0.2

label_dict, label_dict_res = {}, {}
with open("label.txt", 'r') as f:
    for line in f.readlines():
        folder, label = line.strip().split(':')[0], line.strip().split(':')[1]
        label_dict[folder] = label
        label_dict_res[label] = folder
print(label_dict)

N_CLASSES = len(label_dict)
IMG_W = 224
IMG_H = IMG_W


def mkdir(path, rm=False):
    if os.path.exists(path):
        if rm:
            shutil.rmtree(path)
            os.mkdir(path)
    else:
        os.mkdir(path)


def save_conf_mat(confusion_mat, classes):
    # 归一化
    confusion_mat_N = confusion_mat.copy()
    for i in range(len(classes)):
        confusion_mat_N[i, :] = confusion_mat[i, :] / (confusion_mat[i, :].sum() + 1e-9)

    # 获取颜色
    fig=plt.figure(figsize=(35, 35))
    cmap = plt.cm.get_cmap('Greys')  # 更多颜色: http://matplotlib.org/examples/color/colormaps_reference.html
    plt.imshow(confusion_mat_N, cmap=cmap)
    plt.colorbar()

    # 设置文字
    xlocations = np.array(range(len(classes)))
    plt.xticks(xlocations, list(classes), rotation=30)
    plt.yticks(xlocations, list(classes))
    plt.xlabel('Predict label')
    plt.ylabel('True label')
    plt.title('result')

    # 打印数字
    for i in range(confusion_mat_N.shape[0]):
        for j in range(confusion_mat_N.shape[1]):
            plt.text(x=j, y=i, s=int(confusion_mat[i, j]), va='center', ha='center', color='red', fontsize=10)
    # 保存
    plt.savefig('result.png')
    plt.close()
    print("save result.png")


def init_tf():
    global sess, pred, x
    sess = tf.Session(config=config)
    with gfile.FastGFile('./model.pb', 'rb') as f:
        graph_def = tf.GraphDef()
        graph_def.ParseFromString(f.read())
        sess.graph.as_default()
        tf.import_graph_def(graph_def, name='')

    # 获取输入tensor
    x = tf.get_default_graph().get_tensor_by_name("input_1:0")
    print("input:", x)
    # 获取预测tensor
    pred = tf.get_default_graph().get_tensor_by_name("softmax_out:0")  #mobilenet_v2
    print('load model done...')

def evaluate_image(image):
    # read and process image
    #image = cv2.imread(img_dir)
    image = cv2.cvtColor(image,cv2.COLOR_BGR2RGB)
    # pre-process the image for classification
    image = cv2.resize(image, (IMG_W, IMG_W))
    #image = image.astype("float") / 255.0
    im_mean = np.mean(image)
    stddev = max(np.std(image), 1.0/np.sqrt(IMG_W*IMG_W*3))
    image = (image - im_mean) / stddev #代替tf.image.per_image_standardization
    #image_array = np.array(image)
    image = np.expand_dims(image, axis=0)

    prediction = sess.run(pred, feed_dict={x: image})
    max_index = np.argmax(prediction)
    pred_label = label_dict_res[str(max_index)]
    print("%s, predict: %s(index:%d), prob: %f" %(img_dir, pred_label, max_index, prediction[0][max_index]))
    return pred_label, prediction[0][max_index]
    

if __name__ == '__main__':
    init_tf()
    output = "error_result"
    f = open("result.txt", "w")
    mkdir(output, rm=True)
    data_path = "/media/lishundong/DATA2/docker/data/sku_for_classify2/sku_val"
    #data_path = "/media/lishundong/DATA2/docker/data/sku_for_classify2/sku_val_small"
    #data_path = "/media/lishundong/DATA2/docker/data/sku_for_classify2/sku_train"
    label = os.listdir(data_path)
    cls_num = N_CLASSES
    conf_mat = np.zeros([cls_num, cls_num])
    for l in tqdm.tqdm(label):
        if os.path.isfile(os.path.join(data_path, l)):
            continue
        for img_dir in glob.glob(os.path.join(data_path, l, "*.jpg")):
            img = cv2.imread(img_dir)
            pre_label, prob = evaluate_image(img)
            if l != pre_label:
                save_badcase_dir = os.path.join(output, pre_label)
                mkdir(save_badcase_dir, rm=False)
                shutil.copyfile(img_dir, save_badcase_dir + "/" + os.path.basename(img_dir))
            index_l = int(label_dict[l])
            index_pre_label = int(label_dict[pre_label])
            conf_mat[index_l, index_pre_label] += 1.0
    for i in range(cls_num):
        cls = label_dict_res[str(i)]
        gt_num_this_cls = max(1e-9, np.sum(conf_mat[i, :]))
        predict_right = np.sum(conf_mat[i, i])
        predict_this_cls = max(1e-9, np.sum(conf_mat[:, i]))
        precise = predict_right / predict_this_cls
        recall = predict_right / gt_num_this_cls
        log = 'class:{:<5} total num:{:<5} right num:{:<5} Precision: {:.2%}({}/{})  Recall: {:.2%}({}/{})'.format(cls, int(gt_num_this_cls), int(predict_right), precise, int(predict_right), int(predict_this_cls), recall, int(predict_right), int(gt_num_this_cls))
        print(log)
        f.write(log + "\n")
    log = 'All accuracy:{:.2%}({}/{})'.format(np.trace(conf_mat) / np.sum(conf_mat), int(np.trace(conf_mat)), np.sum(conf_mat))
    print(log)
    f.write(log)
    save_conf_mat(conf_mat, [label_dict_res[str(index)] for index in range(cls_num)])
    sess.close()

显示：

class:10 total num:336 right num:335 Precision: 99.70%(335/336) Recall: 99.70%(335/336)
class:104 total num:703 right num:703 Precision: 100.00%(703/703) Recall: 100.00%(703/703)
class:105 total num:498 right num:497 Precision: 100.00%(497/497) Recall: 99.80%(497/498)
class:106 total num:558 right num:558 Precision: 99.11%(558/563) Recall: 100.00%(558/558)
class:11 total num:135 right num:135 Precision: 100.00%(135/135) Recall: 100.00%(135/135)
class:112 total num:324 right num:324 Precision: 99.69%(324/325) Recall: 100.00%(324/324)
class:113 total num:436 right num:436 Precision: 100.00%(436/436) Recall: 100.00%(436/436)
class:114 total num:499 right num:498 Precision: 100.00%(498/498) Recall: 99.80%(498/499)
......