本文介绍如何在C++环境中部署Keras或TensorFlow模型。
一、对于Keras,
第一步,使用Keras搭建、训练、保存模型。
model.save('./your_keras_model.h5')
第二步,冻结Keras模型。
from keras.models import load_model import tensorflow as tf from tensorflow.python.framework import graph_io from keras import backend as K def freeze_session(session, keep_var_names=None, output_names=None, clear_devices=True): from tensorflow.python.framework.graph_util import convert_variables_to_constants graph = session.graph with graph.as_default(): freeze_var_names = list(set(v.op.name for v in tf.global_variables()).difference(keep_var_names or [])) output_names = output_names or [] output_names += [v.op.name for v in tf.global_variables()] input_graph_def = graph.as_graph_def() if clear_devices: for node in input_graph_def.node: node.device = "" frozen_graph = convert_variables_to_constants(session, input_graph_def, output_names, freeze_var_names) return frozen_graph K.set_learning_phase(0) keras_model = load_model('./your_keras_model.h5') print('Inputs are:', keras_model.inputs) print('Outputs are:', keras_model.outputs) frozen_graph = freeze_session(K.get_session(), output_names=[out.op.name for out in model.outputs]) graph_io.write_graph(frozen_graph, "./", "your_frozen_model.pb", as_text=False)
二、对于TensorFlow,
1、使用TensorFlow搭建、训练、保存模型。
saver = tf.train.Saver() saver.save(sess, "./your_tf_model.ckpt")
2、冻结TensorFlow模型。
python freeze_graph.py --input_checkpoint=./your_tf_model.ckpt --output_graph=./your_frozen_model.pb --output_node_names=output_node
三、使用TensorFlow的C/C++接口调用冻结的模型。这里,我们向模型中输入一张经过opencv处理的图片。
#include "tensorflow/core/public/session.h" #include "tensorflow/core/platform/env.h" #include "opencv2/opencv.hpp" #include <iostream> using namespace tensorflow; int main(int argc, char* argv[]){ // tell the network that it is not training phaseTensor = Tensor(DT_BOOL, TensorShape()); auto phaseTensorPointer = phaseTensor.tensor<bool, 0>(); phaseTensorPointer(0) = false; // read the input image cv::Mat img = imread('./your_input_image.png', 0); input_image_height = img.size().height; input_image_width = img.size().width; input_image_channels = img.channels(); imageTensor = Tensor(DT_FLOAT, TensorShape({1, input_image_height, input_image_width, input_image_channels})); // convert the image to a tensor float * imageTensorPointer = imageTensor.flat<float>().data(); cv::Mat imageTensorMatWarpper(input_image_height, input_image_width, CV_32FC3, imageTensorPointer); img.convertTo(imageTensorMatWarpper, CV_32FC3); // construct the input string input_node_name1 = "input tesnor name1"; string input_node_name2 = "input tensor name2"; std::vector<std::pair<string, Tensor>> inputs; inputs = {{input_node_name1, phaseTensor}, {input_node_name2, imageTensor},}; // start a new session Session* session; Status status = NewSession(SessionOptions(), &session); if (!status.ok()) { cout << "NewSession failed! " << status.error_message() << std::endl; } // read the frozen graph GraphDef graph_def; status = ReadBinaryProto(Env::Default(), "./your_frozen_model.pb", &graph_def); if (!status.ok()) { cout << "ReadBinaryProto failed! " << status.error_message() << std::endl; } // initialize the session graph status = session->Create(graph_def); if (!status.ok()) { cout << "Create failed! " << status.error_message() << std::endl; } // define the output string output_node_name1 = "output tensor name1"; std::vector<tensorflow::Tensor> outputs; // run the graph tensorflow::Status status = session->Run(inputs, {output_node_name1}, {}, &outputs); if (!status.ok()) { cout << "Run failed! " << status.error_message() << std::endl; } // obtain the output Tensor output = std::move(outputs[0]); tensorflow::StringPiece tmpBuff = output.tensor_data(); const float* final_output = reinterpret_cast<const float*>(tmpBuff.data()); //for classification problems, the output_data is a tensor of shape [batch_size, class_num] /* auto scores = outputs[0].flat<float>(); */ session->Close(); return 0; }