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

 可看到因为此文件所在目录已经在环境变量里,所以终端直接输入此命令,是可以识别到的。pb模型转uff模型(tensorflow2.x)_第1张图片

 然后调用如下语句命令方式(具体结合自己的路径配置,博主这边已经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文件中代码

pb模型转uff模型(tensorflow2.x)_第2张图片

 为如下:

    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)

验证:

再次输入如下命令语句:

pb模型转uff模型(tensorflow2.x)_第3张图片

 pb模型转uff模型(tensorflow2.x)_第4张图片

 成功转换

附:如何定位到问题点的

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

pb模型转uff模型(tensorflow2.x)_第5张图片

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

pb模型转uff模型(tensorflow2.x)_第6张图片

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

pb模型转uff模型(tensorflow2.x)_第7张图片

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

pb模型转uff模型(tensorflow2.x)_第8张图片

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

你可能感兴趣的:(Python与深度学习,pb,uff,tensorrt,tensorflow,tensorflow2.x)