OpenCV DNN模块——从TensorFlow模型导出到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_pathsaver = 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 TransformGraphwith 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模式下正常运行。
总结
输入和输出的节点名称必须知道，自己设定或者按照我说的办法做。
输入和输出的节点可以是多个，根据需要自己设定。
第一次保存模型文件最好是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实用技术、计算机视觉最新突破，