TensorRT是一个高性能的深度学习推理(Inference)优化器,可以为深度学习应用提供低延迟、高吞吐率的部署推理。TensorRT可用于对超大规模数据中心、嵌入式平台或自动驾驶平台进行推理加速。TensorRT现已能支持TensorFlow、Caffe、Mxnet、Pytorch等几乎所有的深度学习框架,将TensorRT和NVIDIA的GPU结合起来,能在几乎所有的框架中进行快速和高效的部署推理。(这段借鉴别人的文章)
但是,tensorRT真的对于所有的框架都一视同仁么,我们上一张图就知道到底是什么情况了。
这是tensorRT7.1的c++接口文档,针对于三种转换格式的定义的类的数量都明显不同。一般的转换路径有以下几种:tensorflow->onnx->tensorRT;tensorflow->uff->tensorRT;tensorflow->caffe->tensorRT(没有尝试过);pytorch->onnx->tensorRT。其他的转换路径还没了解。所以大部分都是从onnx和uff这两个中间载体转换的。但是可以发现caffe提供的类个数最多。而onnx最少还没有plugin的格式(onnx-tensorrt好像可以支持,还未仔细了解),uff只支持IPluginFactory和IPluginFactoryExt这两种。
对于绝大部分的非分类网络,肯定是需要一些自定义的插件来对tensorrt算子进行补充的。那么也就是说,官方文档中只有uff和caffe这两条路可以走通。本文中将详细说明tensorflow->uff->tensorRT这个路径下的模型转换,自定义插件的开发,最终部署这三个方面。
1,tensorflow->uff
tensorflow freeze的pb模型转换默认数据格式为NHWC,而在tensorRT中的数据格式默认为NCHW。但是不必担心,pb模型通过tensorRT转换工具转换后的uff模型,在解析阶段会自动解析为NCHW的算子。所以只需要将输入的格式设置成NCHW即可。
转换命令:
convert-to-uff frozen_inference_graph.pb -p config.py -t
转换的配置文件config.py中定义了网络中不支持的算子与自定义算子的对应关系:
import graphsurgeon as gs
import tensorflow as tfInput = gs.create_node("Input",op="Placeholder",dtype=tf.float32,shape=[1, 3, 224, 224])ResizeNearest = gs.create_node(name="ResizeNearest", op="BatchTilePlugin_TRT", dtype=tf.float32, axis=2)namespace_plugin_map = {"input_data": Input,"upsampled": ResizeNearest
}def preprocess(dynamic_graph):# Now create a new graph by collapsing namespacesdynamic_graph.collapse_namespaces(namespace_plugin_map, unique_inputs=True)
程序中重新定义了输入和upsample算子,并使用name_plugin_map重新进行了映射。
经过转换之后会生成frozen_inference_graph.uff的文件,使用Netron查看,对应的input与upsample的算子类型已经改变。