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了!

评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包

打赏作者

竹叶青lvye

你的鼓励将是我创作的最大动力

¥1 ¥2 ¥4 ¥6 ¥10 ¥20
扫码支付:¥1
获取中
扫码支付

您的余额不足,请更换扫码支付或充值

打赏作者

实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值