Tensorflow pb模型转uff模型方法及遇到KeyError20和expected Const问题解决

article/2025/8/16 6:58:07

项目所需要将Tensorflow 生成的pb模型转为uff模型，方法很简单，但是遇到的问题着实很多，这里主要记录下问题及解决方法，总之，最后是成功生成uff格式的模型的，有需要的可以参考。

pb模型转为uff模型的方法及步骤，网上很多，最后使用的方法如下：

1、电脑没装win10+ubuntu的双系统，为了简便，用vmware创建了个ubuntu的虚拟机，ubuntu系统是ubuntu-16.04.3-desktop-amd64；

2、进入虚拟机的Ubuntu系统，打开Terminal，安装了个Anaconda，使用如下命令创建Python虚拟环境，并安装python2.7和numpy：

conda create --name tensorflow-cpu python=2.7 numpy

3、创建完毕后，激活tensorflow-cpu虚拟环境：

conda activate tensorflow-cpu

4、安装cpu版本的tensorflow，这里没安装gpu版本的，cpu版本的转换uff够用了，用的1.10.0版本的也可以，别的版本的需要自己测试了：

pip install --upgrade pip
pip install tensorflow==1.10.0

5、去下载TensorRT-5.0.2.6.Ubuntu-16.04.4.x86_64-gnu.cuda-9.0.cudnn7.3.tar.gz：

下载网址：https://developer.nvidia.com/nvidia-tensorrt-5x-download

下载完后解压，主要得到三个安装whl文件：

6、上面下载的文件拷贝到Ubuntu系统，分别进入相应的目录，安装这三个whl文件即可：

pip install tensorrt-5.0.2.6-py2.py3-none-any.whl

pip install uff-0.5.5-py2.py3-none-any.whl

pip install graphsurgeon-0.3.2-py2.py3-none-any.whl

7、将训练好的pb模型文件拷贝到Ubuntu系统，进入到模型目录，使用如下命令转换模型文件：

convert-to-uff -o ssd_inception_v2_coco.uff --input_file frozen_inference_graph.pb

结果如下图所示，生成所需要的uff模型：

上述方法很简单，但是遇到的问题就让人惆怅了。

问题1：运行convert-to-uff转换命令后，运行不成功，提示KeyError:20的错误：

原因分析：前期探索过程中使用的uff-0.6.5和graphsurgeon-0.4.1版本，就会报这个错，试了好久，后来换成uff-0.5.5和graphsurgeon-0.3.2，就不报这个错误了。

问题2：运行convert-to-uff转换命令后，提示“uff.model.exceptions.UffException: Transpose permutation has op Sub, expected Const. Only constant permuations are supported in UFF.”的错误：

原因分析：Tensorrt并不支持Tensorflow的所有操作，但是tf.transpose操作对于Tensorrt5.0是支持的，但是并不支持tf.transpose里面还有tf.subtract操作，所以需要对代码里面的tf.subtract操作进行Constant folding optimizer处理，将运算转换为常量；也就是网上NVIDIA给出的解决方法：

刚看到这个constfold optimizer，根本不知道是什么啊，没办法，只能各种搜，最终google到了一个可以适用于这个错误的解决方法，就是在训练模型的代码里，找到tf.transpose的位置，加入如下代码，解决此问题，其实就是Tensorflow的一种优化处理，Constant folding optimizer，用来提升代码处理速度的：

import contextlib
import tensorflow as tf@contextlib.contextmanager
def options(options):old_opts = tf.config.optimizer.get_experimental_options()tf.config.optimizer.set_experimental_options(options)try:yieldfinally:tf.config.optimizer.set_experimental_options(old_opts)……
……
……with options({'constant_folding': True}):……tf.transpose(……)……

此外，也搜索到了官方给的名为fold constants的工具，fold_constants.py文件，可以自己试试官方的方法，文件代码如下（偶然的机会在google上搜到的，链接没保存，找不到了）：

#!/usr/bin/env python3
from tensorflow.core.protobuf import config_pb2, rewriter_config_pb2, meta_graph_pb2
from tensorflow.core.framework import graph_pb2
from tensorflow.python.saved_model import tag_constants
from tensorflow.python.framework import importer, ops
from tensorflow.python.grappler import tf_optimizer
from tensorflow.python.training import saver# For automatic output node deduction.
import graphsurgeon as gs
import argparsedef constfold(graphdef):graph = ops.Graph()with graph.as_default():outputs = [node.name for node in gs.StaticGraph(graphdef).graph_outputs]output_collection = meta_graph_pb2.CollectionDef()output_list = output_collection.node_list.valuefor output in outputs:output_list.append(output.encode("utf-8"))importer.import_graph_def(graphdef, name="")metagraph = saver.export_meta_graph(graph_def=graph.as_graph_def(add_shapes=True), graph=graph)metagraph.collection_def["train_op"].CopyFrom(output_collection)rewriter_config = rewriter_config_pb2.RewriterConfig()rewriter_config.optimizers.extend(["constfold"])rewriter_config.meta_optimizer_iterations = (rewriter_config_pb2.RewriterConfig.ONE)session_config = config_pb2.ConfigProto()session_config.graph_options.rewrite_options.CopyFrom(rewriter_config)return tf_optimizer.OptimizeGraph(session_config, metagraph, graph_id=b"tf_graph")def main():parser = argparse.ArgumentParser("Folds constants in the provided frozen model")parser.add_argument("model_file", help="The frozen model in which to fold constants.")parser.add_argument("-o", "--output-file", help="The path at which to write the new graph, with constants folded.")args, _ = parser.parse_known_args()with open(args.model_file, 'rb') as f:graphdef = graph_pb2.GraphDef()graphdef.ParseFromString(f.read())folded_graph = constfold(graphdef)if args.output_file:with open(args.output_file, "wb") as f:f.write(folded_graph.SerializeToString())if __name__ == '__main__':main()

问题3：加了上面的代码后，可能会出现“AttributeError: module 'tensorflow' has no attribute 'config'”的问题：

原因分析：跑训练模型的代码使用tensorflow版本必须大于等于1.14.0，小于此版本会报这个错，tensorflow gpu 1.14.0以后的版本需要cuda10.0以上版本才可以；

完！方法很简单，但是遇到的问题让人烦。

最后，如果上述方法还行不通的话，本人建议使用英伟达开源的模型转换工具；

官方链接为：https://github.com/AastaNV/TRT_object_detection

按照网站的步骤，执行完main.py将会生成temp.uff文件，还能测试生成的uff文件可不可用。

这个方法的参考链接：https://www.it610.com/article/1283183189296168960.htm

https://www.minds.ai/post/deploying-ssd-mobilenet-v2-on-the-nvidia-jetson-and-nano-platforms

有问题，可一起讨论。