三、生成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)
......
下一篇：使用nvidia的tensorrt加速.pb文件https://blog.csdn.net/u010397980/article/details/86382849
扫一扫
4338
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
