大多数的博客只是提到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了!