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

引言

对于机器视觉从事者或者研究者来说，把训练好模型部署到项目中是关键的一步。现如今各大相机厂商都会提供相机的二次开发包，供给使用者进行使用和开发。据博主所知，目前大部分的相机开发包并不支持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 tf
import os.path
import argparse
from tensorflow.python.framework import graph_util
def 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 tf
from tensorflow.python.tools import optimize_for_inference_lib
from 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官方提供了语义分割，目标检测以及分类的解码例程，可以点击这里看下。下面是导入代码（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层，一定要这样做，否则会有报错。

参考 https://github.com/tflearn/tflearn/issues/964