OpenCV读取tensorflow神经网络模型：SavedModel格式转为frozen graph的方法

news2024/11/15 7:05:50

本文介绍基于Python的tensorflow库，将tensorflow与keras训练好的SavedModel格式神经网络模型转换为frozen graph格式，从而可以用OpenCV库在C++ 等其他语言中将其打开的方法。

如果我们需要训练并使用一个神经网络模型，一般情况下都是首先借助Python语言中完善的神经网络模型API对其加以训练，训练完毕后在C++、Java等语言环境下高效、快速地使用它。最近，就需要在C++ 中打开、使用几个前期已经在Python的tensorflow库中训练好的神经网络模型。但是，由于训练模型时使用的是2.X版本的tensorflow库（且用的是keras的框架），所以训练模型后保存的是SavedModel格式的神经网络模型文件——就是包含3个.pb格式文件，以及assets及variables等2个文件夹那种形式的模型；如下图所示。

而在C++ 中读取神经网络模型，首先是可以借助tensorflow库的C++ API来实现，但是这种方法非常复杂——完整的TensorFlow C++ API部署起来非常困难——需要系统盘至少40 G到50 G的剩余空间、动辄0.5 h到1 h的编译时长，经常需要花费一周的时间才可以配置成功；所以如果仅仅是需要在C++ 中读取已经训练好的神经网络模型的话，没必要花费这么大功夫去配置TensorFlow C++ API。而同时，基于OpenCV库，我们则可以在简单、快速地配置完其环境后，就基于1个函数对训练好的tensorflow库神经网络模型加以读取、使用。这里如果大家需要配置C++ 环境的OpenCV库，可以参考文章C++计算机视觉库OpenCV在Visual Studio 2022的配置方法（https://blog.csdn.net/zhebushibiaoshifu/article/details/128260507）。

但是，还有一个问题——OpenCV库自身目前仅支持读取tensorflow库frozen graph格式的神经网络模型，不支持读取SavedModel格式的模型。因此，如果希望基于OpenCV库读取tensorflow中SavedModel格式的模型，就需要首先将其转换为frozen graph格式；那么，本文就介绍一下这个操作的具体方法，并给出2种实现这一转换功能的Python代码。

首先，本文神经网络模型格式转换的代码是基于Python环境中tensorflow库实现的，因此需要配置好这一个库（大家都已经需要转换神经网络模型的格式了，那Python环境中tensorflow库肯定早已经配置好了）；如果没有配置，可以参考文章Anaconda配置Python新版本tensorflow库（CPU、GPU通用）的方法（https://blog.csdn.net/zhebushibiaoshifu/article/details/129285815）。

第1种代码如下。

# -*- coding: utf-8 -*-
"""
Created on Sat Mar  9 14:31:18 2024

@author: fkxxgis
"""

import tensorflow as tf
from tensorflow.keras import models
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2

# model_save_model = tf.saved_model.load("F:/Data_Reflectance_Rec/model/model_blue/")
model_save_model = models.load_model("F:/Data_Reflectance_Rec/model/model_blue/")

signatures = model_save_model.signatures["serving_default"]
graph = tf.function(lambda x: model_save_model(x))
graph = graph.get_concrete_function(tf.TensorSpec(signatures.inputs[0].shape.as_list(), signatures.inputs[0].dtype.name))
frozen_variable = convert_variables_to_constants_v2(graph)
frozen_variable.graph.as_graph_def();

tf.io.write_graph(graph_or_graph_def = frozen_variable.graph, 
                  logdir = "F:/Data_Reflectance_Rec/model/model_blue_new", 
                  name = "frozen_graph.pb", 
                  as_text = False)
# tf.io.write_graph(graph_or_graph_def = frozen_variable.graph, 
#                   logdir = "F:/Data_Reflectance_Rec/model/model_blue_new", 
#                   name = "frozen_graph.pbtxt", 
#                   as_text = True)

其中，我们首先需要导入对应的Python模块和convert_variables_to_constants_v2()函数。

随后，加载我们待转换的、SavedModel格式的tensorflow神经网络模型。这里需要注意，我写了2句不同的代码来加载初始的模型——其中，如果用第1句代码加载模型，倒也可以不报错地运行完成上述代码，但是等到用C++ 环境的OpenCV库读取这个转换后的模型时，会出现Microsoft C++ 异常: cv::Exception字样的报错，如下图所示；而如果用第2句代码加载模型，就没有问题。之所以会这样，应该是因为我当初训练这个神经网络模型时，用的是tensorflow中keras模块的Model，所以导致加载模型时，就不能用传统的加载SavedModel格式模型的方法了（可能是这样）。

接下来，我们从初始模型中获取其签名。tensorflow库中的签名（Signature），是用于定义模型输入、输出的一种机制——其定义了模型接受的输入参数和返回的输出结果的名称、数据类型和形状等信息；这个默认签名为serving_default，我们这里获取这个默认的签名即可。

接下来，这个graph = tf.function(lambda x: model_save_model(x))表示将模型封装在tensorflow的图函数中；随后，get_concrete_function()获取具体函数并指定输入张量的形状和数据类型。说实话，这里的2行代码我也搞不太清楚具体详细含义是什么——但大体上，这些内容应该是tensorflow库1.X版本中的一些操作与名词（因为frozen graph格式的模型本来就是tensorflow库1.X版本中用的，而SavedModel格式则是2.X版本中常用的）。

再次，通过convert_variables_to_constants_v2()函数，将图中的变量转换为常量，并基于as_graph_def()定义1个冻结图。

最后，就可以通过tf.io.write_graph()函数，将冻结图写入指定的目录中，输出文件名为frozen_graph.pb，as_text = False表示以二进制格式保存这个模型（如果不加这个参数，就相当于成了.pbtxt文件了，导致后续用C++环境的OpenCV库还是读取不了这个模型）。代码末尾，还有一段注释的部分——如果取消注释，将以文本格式保存冻结图，也就是.pbtxt文件。因为我们只要.pb文件就够了，所以就不需要这段代码了。

执行上述代码，在结果文件夹中，我们将看到1个.pb格式的神经网络模型结果文件，如下图所示。

接下来，在C++、Python等语言的OpenCV库中，我们都可以基于cv::dnn::readNetFromTensorflow()这个函数，来读取我们的神经网络模型了。

除此之外，再给出另一个版本的转换代码；这个代码其实和前述代码的含义差不多，如果前述代码不能执行，大家可以再尝试尝试下面这个。

import tensorflow as tf
from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2

loaded = tf.saved_model.load('F:/Data_Reflectance_Rec/model/model_nir/')
infer = loaded.signatures['serving_default']

f = tf.function(infer).get_concrete_function(tf.TensorSpec(infer.inputs[0].shape.as_list(), dtype=tf.float32))
f2 = convert_variables_to_constants_v2(f)
graph_def = f2.graph.as_graph_def()

with tf.io.gfile.GFile('frozen_graph.pb', 'wb') as f:
    f.write(graph_def.SerializeToString())