OpenCV DNN模块——从TensorFlow模型导出到OpenCV部署详解

Posted 2021-04-25 OpenCV中文网

重磅干货，第一时间送达

---

本文来自OpenCV中文网粉丝小庄头发很多投稿。

(原文链接：https://blog.csdn.net/weixin_39928773/article/details/103910850)

引言

对于机器视觉从事者或者研究者来说，把训练好模型部署到项目中是关键的一步。现如今各大相机厂商都会提供相机的二次开发包，供给使用者进行使用和开发。据博主所知，目前大部分的相机开发包并不支持Python语言，而主流的深度学习框架都是基于Python语言，训练好的模型难以部署到自己的软件中。

举个例子，博主一般使用C/C++语言对相机进行二次开发及编写工业软件，使用Python语言的TensorFlow框架训练模型，博主一般采用两种方式在C/C++中调用训练好的模型：

• C/C++和Python混合编程，网上这部分的教程很多，博主这边就不过多说明。这个方法虽然不难，但是在使用过程中需要常常需要释放变量内存，并且程序在出错时难以排查。

• 使用OpenCV的DNN模块加载TensorFlow训练好的模型，这个方式博主十分推荐，无需混合语言编程，避免难以察觉的错误，并且在OpenCV支持使用OpenVINO对CPU加速以及在4.2版本以后支持CUDA加速推理，部署方法在博主这两篇博客中有介绍。

本文将对使用TensorFlow训练模型、导出模型、模型转化以及OpenCV模型导入方法及过程详细介绍。

使用过程

1.模型训练

模型训练这块博主以自己写的UNet网络为例，见下面的代码。

input = tf.placeholder(shape=[None, 640, 640, 3], name="input", dtype=tf.float32)gt = tf.placeholder(shape=[None, 640, 640, 1], dtype=tf.uint8)
weight_regularizer = contrib.layers.l2_regularizer(0.0005)with slim.arg_scope([slim.conv2d, slim.conv2d_transpose], weights_regularizer = weight_regularizer, biases_regularizer = weight_regularizer, biases_initializer = tf.constant_initializer(0.0), weights_initializer = 'he_normal', reuse=None): logits = Unet(input, layer_dims, name='color')
logits = slim.conv2d(inputs=logits, kernel_size=[3,3], num_outputs=num_classes, activation_fn=None, scope='logits')output = tf.nn.softmax(logits, axis = -1, name='output')
##print(output.name)

需要注意的点是对于输入和输出，必须自己设定名字，如博主这边设定的名字是"input"和’'output"。如果说是使用别人的网络模型，不知道输出的名字，有一个办法就是手动输出模型最后的节点。如：print(output.name)。对于输入，第一维度的shape必须是None，否则OpenCV导入的时候有可能报错。然后开始训练模型。

2.导出模型

模型训练好之后保存成ckpt文件。由于这块地方容易出错，博主强烈不建议直接保存成.pb文件，先保存成ckpt文件，然后转成pb文件。转化代码如下：

import tensorflow as tfimport os.pathimport argparsefrom tensorflow.python.framework import graph_utildef freeze_graph(model_folder, output_node_names): checkpoint = tf.train.get_checkpoint_state(model_folder)  input_checkpoint = checkpoint.model_checkpoint_path
 saver = tf.train.import_meta_graph(input_checkpoint + '.meta', clear_devices=True)  graph = tf.get_default_graph()  input_graph_def = graph.as_graph_def()
 with tf.Session() as sess: saver.restore(sess, input_checkpoint)  for node in input_graph_def.node: if node.op == 'RefSwitch': node.op = 'Switch' for index in range(len(node.input)): if 'moving_' in node.input[index]: node.input[index] = node.input[index] + '/read' elif node.op == 'AssignSub': node.op = 'Sub' if 'use_locking' in node.attr: del node.attr['use_locking']
 output_graph_def = graph_util.convert_variables_to_constants(  sess, input_graph_def, output_node_names  ) with tf.gfile.GFile('frozen_model.pb', "wb") as f:  f.write(output_graph_def.SerializeToString())
if __name__ == '__main__':    freeze_graph('checkpoints/', ['output'])

如果模型中有残差结构或者BN层，一定要使用上面的代码去转换成pb文件，否则后续读入文件时将会出错。

3.模型优化

保存成pb文件之后，可以在进行一次优化，把模型中某些层合并，删除掉其中冗余的层，见如下代码。

import tensorflow as tffrom tensorflow.python.tools import optimize_for_inference_libfrom tensorflow.tools.graph_transforms import TransformGraph
with tf.gfile.FastGFile('frozen_model.pb', 'rb') as f: graph_def = tf.GraphDef() graph_def.ParseFromString(f.read()) graph_def = optimize_for_inference_lib.optimize_for_inference(graph_def, ['input'], ['ResizeBilinear'], tf.float32.as_datatype_enum) graph_def = TransformGraph(graph_def, ['input'], ['ResizeBilinear'], ['remove_nodes(op=PlaceholderWithDefault)', 'sort_by_execution_order']) with tf.gfile.FastGFile('final_model.pb', 'wb') as f: f.write(graph_def.SerializeToString()) # tf.train.write_graph(graph_def, "", 'final_model.pbtxt', as_text=True)

有的模型还需要pbtxt模型结构文件，但是博主这边没有用到。

4.OpenCV模型导入

OpenCV模型导入这块相对简单，但是要自己进行相应的解码。OpenCV官方提供了语义分割，目标检测以及分类的解码例程，可以查看下方参考资料[1]。下面是导入代码（C++代码，Python同理）：

Net net2 = readNetFromTensorflow("final_model.pb"); //载入模型
net2.setPreferableBackend(DNN_BACKEND_CUDA);net2.setPreferableTarget(DNN_TARGET_CUDA);//设置推理后台
Mat image = imread("color.png");
vector<Mat> images(1, image);Mat inputBlob2 = blobFromImages(images, 1 / 255.F, Size(640, 640), Scalar(), true, false);
net2.setInput(inputBlob2); //输入数据Mat score;net2.forward(score); //前向传播Mat segm;colorizeSegmentation(score, segm);   //结果可视化

到这边就完成了，OpenCV4.2似乎有个Bug，在Debug模式下有时候会运行失败，而Release模式下正常运行。

总结

1. 输入和输出的节点名称必须知道，自己设定或者按照我说的办法做。

   
   

2. 输入和输出的节点可以是多个，根据需要自己设定。

   
   

3. 第一次保存模型文件最好是ckpt文件，将保存好的ckpt文件在转化成pb文件，如果模型中有残差结构或者BN层，一定要这样做，否则会有报错。

   
   

参考资料：

[1] https://github.com/opencv/opencv/tree/master/samples/dnn

[2] https://github.com/tflearn/tflearn/issues/964

---

OpenCV技术交流群

关注最新的OpenCV技术，欢迎加入开发者群群，扫码添加CV君拉你入群，

（请务必注明：opencv）

（不会时时在线，如果没能及时通过验证还请见谅）

---

OpenCV中文网专注分享：

OpenCV实用技术、计算机视觉最新突破，