tensorflow 在训练完一个模型之后,如何将其应用到我的产品中呢?
参考源码:github:(https://github.com/azuredsky/tensorflow-tutorial-1)
暂且记录下关键步骤:
1.将训练好的模型转固化为二进制文本
工具使用freeze_graph.py工具
python freeze_graph.py --input_graph=../model/nn_model.pbtxt --input_checkpoint=../ckpt/nn_model.ckpt --output_graph=../model/nn_model_frozen.pb --output_node_names=output_node |
参数1:./model/nn_model.pbtxt
通过tf.train.write_graph(session.graph_def, FLAGS.model_dir, "nn_model.pbtxt", as_text=True)生成实现,
将图中的计算节点保存在模型nn_model.pbtxt中
参数2:../ckpt/nn_model.ckpt 训练好的参数模型
参数3:../model/nn_model_frozen.pb 设置输出文件名
参数4:output_node,在图中输出节点的名字python 代码:
for op in tf.get_default_graph().get_operations():
print(op.name)
output_node 输出节点需要为在python实际sess.run的节点,否则图中不会保存。
在输出nn_model_frozen.pb就可以通过python 或者C++直接调用了
python示例:
def load_graph(fz_gh_fn):
with tf.gfile.GFile(fz_gh_fn,"rb") as f:
graph_def = tf.GraphDef()
graph_def.ParseFromString(f.read())
with tf.Graph().as_default() as graph:
tf.import_graph_def(
graph_def,
input_map = None,
return_elements = None,
name = "prefix"
)
return graph
定义了load_graph 就可以加载模型运行了
graph = load_graph(args.fz_model_fn)
x = graph.get_tensor_by_name('prefix/inputs:0')
y = graph.get_tensor_by_name('prefix/output_node:0')
img = Image.open('./test_image/8_0031.bmp')
flatten_img = np.reshape(img, [1, img_H, img_W, channels])
with tf.Session(graph=graph) as sess:
y_out=sess.run(y,feed_dict={x:flatten_img})
print(y_out)
C++ 示例:
首先通过自己手动编译tensorflow 动态链接库
其次加载模型运行:
int main(int argc, char* argv[]) {
// Initialize a tensorflow session
Session* session;
Status status = NewSession(SessionOptions(), &session);
if (!status.ok()) {
std::cout << status.ToString() << "\n";
return 1;
}
// Read in the protobuf graph we exported
// (The path seems to be relative to the cwd. Keep this in mind
// when using `bazel run` since the cwd isn't where you call
// `bazel run` but from inside a temp folder.)
GraphDef graph_def;
status = ReadBinaryProto(Env::Default(), "nn_model_frozen.pb", &graph_def);
if (!status.ok()) {
std::cout << status.ToString() << "\n";
return 1;
}
// Add the graph to the session
status = session->Create(graph_def);
if (!status.ok()) {
std::cout << status.ToString() << "\n";
return 1;
}
// TensorName pre-defined in python file, Need to extract values from tensors
std::string input_tensor_name = "inputs:0";
std::string output_tensor_name = "output_node:0";
Tensor x(DT_FLOAT, TensorShape({1, 32, 24, 3})); // New Tensor shape [1, ndim]
Mat img;
img = imread("4_0179.bmp",1);
imshow("img", img);
waitKey(10);
auto x_map = x.tensor<float, 4>(); // == x.scalar<float>()
cvtColor(img, img, CV_RGB2BGR);
img.convertTo(img,CV_32FC3);
float *data = x_map.data();
memcpy(x_map.data(), (float*)img.data, 1*24*32*3*sizeof(float));
// Setup inputs and outputs:
std::vector<std::pair<string, Tensor>> inputs = {
{ input_tensor_name, x }};
// The session will initialize the outputs
std::vector<Tensor> outputs;
// Run the session, evaluating our "y" operation from the graph
status = session->Run(inputs, { output_tensor_name }, {}, &outputs);
if (!status.ok()) {
std::cout << status.ToString() <<__LINE__<< "\n";
return 1;
}
// Grab the first output (we only evaluated one graph node: "c")
// and convert the node to a scalar representation.
auto output_y = outputs[0].scalar<int>();
// (There are similar methods for vectors and matrices here:
// https://github.com/tensorflow/tensorflow/blob/master/tensorflow/core/public/tensor.h)
// Print the results
std::cout << outputs[0].DebugString() << "\n"; // Tensor<type: float shape: [] values: 32>
std::cout << output_y() << "\n"; // 32
// Free any resources used by the session
session->Close();
return 0;
}
在其中注意tensorflow 使用bgr格式,所以加载自己的图片时需要将其转化为bgr格式。