[Python知识库]pb模型转uff模型(tensorflow2.x）

大多数的博客只是提到tensorflow1.x系列下的转换。大概步骤就是安装tensorrt，同时安装tensorrt下的几个python的wl文件。可参见博主之前的博客:

1.tensorrt的安装Ubuntu配置TensorRT及验证_竹叶青lvye的博客-CSDN博客

2.tensorrt下几个whl文件的安装TensorRT加速方法介绍（python pytorch模型）_竹叶青lvye的博客-CSDN博客_tensorrt加速pytorch

安装uff-0.6.9-py2.py3-none-any.whl后，会提供一个现成的可执行文件convert-to-uff

?可看到因为此文件所在目录已经在环境变量里，所以终端直接输入此命令，是可以识别到的。

?然后调用如下语句命令方式（具体结合自己的路径配置，博主这边已经cd到了weights.pb目录下了）

convert-to-uff -o test.uff --input_file weights.pb

在tensorflow2.x 下则会出现报错：

in from_tensorflow_frozen_model
??? with tf.gfile.GFile(frozen_file, "rb") as frozen_pb:
AttributeError: module 'tensorflow' has no attribute 'gfile'

此时可以修改conversion_helpers.py文件中代码

?为如下：

    graphdef = GraphDef()
    with tf.compat.v1.gfile.GFile(frozen_file, "rb") as frozen_pb:
        graphdef.ParseFromString(frozen_pb.read())
    return from_tensorflow(graphdef, output_nodes, preprocessor, **kwargs)

验证：

再次输入如下命令语句：

?

?成功转换

附：如何定位到问题点的

1.可以打开上面which语句查找到的convert-to-uff文件看看里面有什么

2. 从上面红框中可去再找到这个py文件

3. 打开之后可以看到其代码，看到其接收命令行参数后，调用了uff包里的from_tensorflow_frozen_model函数

4.我们可以在pycharm里看看这个函数在干嘛的

?然后去看下这个代码的实现，就定位到了conversion_helpers.py文件里，然后进行修改就是，毕竟是tensorrt友情提供的一种转换方法（还是为了众人更多的去尝试用tensorrt进行加速），人家也不知道你要用什么tensorflow的环境，所以自己小改动下就ok了！