Tensorflow2.x版本的模型保存,和opencv c++模型导入
由于一个小项目的需求,需要使用C++的接口调用python环境下使用Tensorflow 2.x版本训练好的模型。我想了两种方式:
- 使用 Tensorflow 2. x 的 C++ API 。(我觉得可行,在源码编译的最后一部,由于误操作系统崩了…,后续会继续尝试)
- 使用opencv dnn 模块提供的 API接口。(我使用的这种方式)
使用Tensorflow 训练模型
我使用的是2.1版本,2.0应该也可以
这里就用 fashion_mnist 数据集为例进行训练
导入的一些包:
import cv2
import os
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
import matplotlib.pyplot as plt
os.environ["TF_CPP_MIN_LOG_LEVEL"] = '2'
print('tf版本:',tf.__version__)
输出:tf版本: 2.1.0
导入数据集:
类别标注
每个训练和测试样本都按照以下类别进行了标注:
标注编号描述0T-shirt/top(T恤)1Trouser(裤子)2Pullover(套衫)3Dress(裙子)4Coat(外套)5Sandal(凉鞋)6Shirt(汗衫)7Sneaker(运动鞋)8Bag(包)9Ankle boot(踝靴)
(train_image, train_label), (test_image, test_label) = keras.datasets.fashion_mnist.load_data()
print("数据维度(训练集):", train_image.shape, train_label.shape)
print("数据维度(测试集):", test_image.shape, test_label.shape)
train_data = tf.data.Dataset.from_tensor_slices((train_image, train_label))
train_data = train_data.map(prepprocess).batch(128)
test_data = tf.data.Dataset.from_tensor_slices((test_image, test_label))
test_data = test_data.map(prepprocess).batch(128)
class_name = ['T-shirt', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle Boot']
# 数据可视化
plt.figure(figsize=(10, 10))
for i in range(25):
plt.subplot(5, 5, i + 1)
plt.xticks([])
plt.yticks([])
plt.grid(False)
plt.imshow(train_image[i], cmap=plt.cm.binary)
plt.xlabel(class_name[train_label[i]])
plt.show()
输出:
数据维度(训练集): (60000, 28, 28) (60000,)
数据维度(测试集): (10000, 28, 28) (10000,)
训练
这里直接使用官方文档上的那个网络,比较简单。
def prepprocess(x,y):
# 归一化
x = tf.cast(x, dtype=tf.float32) / 255.
x = tf.reshape(x, [28 * 28])
y = tf.cast(y, dtype=tf.int32)
y = tf.one_hot(y, depth=10)
return x, y
#训练
def train():
# 构建网络
network = keras.Sequential([
keras.layers.Dense(128,'relu'),
keras.layers.Dense(10)
])
network.build(input_shape=(None,28*28))
network.summary()
network.compile(optimizer=keras.optimizers.Adam(lr=0.01),
loss = tf.losses.CategoricalCrossentropy(from_logits=True),
# loss = keras.losses.SparseCategoricalCrossentropy(from_logits=True), # 用这个不用tf.one_hot()
metrics=['accuracy']
)
# 训练
history = network.fit(train_data,epochs=15,validation_data=test_data,validation_freq=1)
plt.plot(history.history['accuracy'],label='accuracy')
plt.plot(history.history['val_accuracy'],label='val_accuracy')
plt.xlabel('Epoch')
plt.ylabel('accuracy')
plt.ylim([0.5,1])
plt.legend(loc='lower right')
plt.show()
tf.saved_model.save(network,'/home/tjk/project/tf_doc/model_save/fashion_10/')
print("保存模型成功")
模型最后保存为pd格式。
模型的opencv导入
使用下面的API:
开始直接就用这个函数加载模型,然后出错了
然后我就在网上找找看有没有解决方法,开始看到的都是一些关于TF 1.x 版本的模型导入,甚至还有标题2.0,结果内容却是1.x的坑人版本。后来在下面的博客上找到了解决方法:
参考博客:
https://blog.csdn.net/ouening/article/details/104335552,博客中提到了解决方法的出处,
https://leimao.github.io/blog/Save-Load-Inference-From-TF2-Frozen-Graph/
按照给出的方法,使用dnn模块的相关接口去调用是无法直接读入模型的,需要对导出的模型进行转换,需要保存为frozen graph格式,
直接copy博客的中一段转换代码,主要是使用了tf.function,在TF的官方文档看到下面的一段话:
In TensorFlow 2.0, eager execution is turned on by default. The user interface is intuitive and flexible (running one-off operations is much easier and faster), but this can come at the expense of performance and deployability.
To get peak performance and to make your model deployable anywhere, use tf.function to make graphs out of your programs. Thanks to AutoGraph, a surprising amount of Python code just works with tf.function, but there are still pitfalls to be wary of.
我的理解就是,opencv 的dnn 接口支持的是以前的图模型,但是在tf2中,用的是eager execution
,所以导出的模型,需要使用tf.function
转化。具体的东西,还在看文档学习。
但是这里有个问题就是在,TF2 的官方文档中写了, 下面有几个API 会在后续的版本中遗弃
# Convert Keras model to ConcreteFunction
full_model = tf.function(lambda x: network(x))
full_model = full_model.get_concrete_function(
tf.TensorSpec(network.inputs[0].shape, network.inputs[0].dtype))
# Get frozen ConcreteFunction
frozen_func = convert_variables_to_constants_v2(full_model)
frozen_func.graph.as_graph_def()
layers = [op.name for op in frozen_func.graph.get_operations()]
print("-" * 50)
print("Frozen model layers: ")
for layer in layers:
print(layer)
print("-" * 50)
print("Frozen model inputs: ")
print(frozen_func.inputs)
print("Frozen model outputs: ")
print(frozen_func.outputs)
# Save frozen graph from frozen ConcreteFunction to hard drive
tf.io.write_graph(graph_or_graph_def=frozen_func.graph,
logdir="./frozen_models",
name="frozen_graph.pb",
as_text=False)
经过上面的转换,可以开始测试了。
一开始还是使用的python 版本的opencv进行的测试:
import cv2
from cv2 import dnn
import numpy as np
print(cv2.__version__)
class_name = ['T-shirt', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle Boot']
img = cv2.imread('/home/tjk/project/picture/aj1.png', cv2.IMREAD_GRAYSCALE)
cv2.imshow("2",img)
print(img.shape)
# print(1-img_cv21/255.)
blob = cv2.dnn.blobFromImage(1-img, # 这里对颜色进行反转,这是训练图片存储格式的问题,可以把数值打印出来看下一
scalefactor=1.0/225.,
size=(28, 28),
mean=(0, 0, 0),
swapRB=False,
crop=False)
print("[INFO]img shape: ", blob.shape)
net = dnn.readNetFromTensorflow('/home/tjk/project/tf_doc/frozen_models/frozen_graph.pb')
print("success!!")
net.setInput(blob)
out = net.forward()
out = out.flatten()
classId = np.argmax(out)
print("classId",classId)
print("预测结果为:",class_name[classId])
cv2.waitKey(0)
使用的图片是这个:AJ1,
给出了预测结果,为Sneaker.
python 这边运行成功了,下面就是c++了
#include <iostream>
#include <opencv2/opencv.hpp>
#include <opencv2/dnn.hpp>
using namespace cv;
using namespace std;
using namespace cv::dnn;
static string class_name[] = {"T-shirt", "Trouser", "Pullover",
"Dress", "Coat","Sandal", "Shirt", "Sneaker", "Bag", "Ankle Boot"};
int main()
{
Mat frame = imread("/home/tjk/project/picture/aj1.png",0);
imshow("1",frame);
// cout<<frame.channels()<<endl;
string path = "/home/tjk/project/tf_doc/frozen_models/frozen_graph.pb";
Net net = readNetFromTensorflow(path);
printf("模型加载成功");
Mat frame_32F;
frame.convertTo(frame_32F,CV_32FC1);
// cout<<1-frame_32F/255.0<<endl;
Mat blob = blobFromImage(1-frame_32F/255.0,
1.0,
Size(28,28),
Scalar(0,0,0));
// cout<<(blob.channels());
net.setInput(blob);
Mat out = net.forward();
// cout<<out.cols<<endl;
Point maxclass;
minMaxLoc(out, NULL, NULL, NULL, &maxclass);
cout <<"预测结果为:"<<class_name[maxclass.x] << endl;
waitKey(0);
}
结果:与之前一样。