TensorRT是一個高性能的深度學習推理（Inference）優(yōu)化器，可以為深度學習應用提供低延遲、高吞吐率的部署推理。TensorRT可用于對超大規(guī)模數(shù)據(jù)中心、嵌入式平臺或自動駕駛平臺進行推理加速。TensorRT現(xiàn)已能支持TensorFlow、Caffe、Mxnet、Pytorch等幾乎所有的深度學習框架，將TensorRT和NVIDIA的GPU結合起來，能在幾乎所有的框架中進行快速和高效的部署推理。（這段借鑒別人的文章）

但是，tensorRT真的對于所有的框架都一視同仁么，我們上一張圖就知道到底是什么情況了。

這是tensorRT7.1的c++接口文檔，針對于三種轉(zhuǎn)換格式的定義的類的數(shù)量都明顯不同。一般的轉(zhuǎn)換路徑有以下幾種：tensorflow->onnx->tensorRT;tensorflow->uff->tensorRT;tensorflow->caffe->tensorRT(沒有嘗試過)；pytorch->onnx->tensorRT。其他的轉(zhuǎn)換路徑還沒了解。所以大部分都是從onnx和uff這兩個中間載體轉(zhuǎn)換的。但是可以發(fā)現(xiàn)caffe提供的類個數(shù)最多。而onnx最少還沒有plugin的格式（onnx-tensorrt好像可以支持，還未仔細了解），uff只支持IPluginFactory和IPluginFactoryExt這兩種。

對于絕大部分的非分類網(wǎng)絡，肯定是需要一些自定義的插件來對tensorrt算子進行補充的。那么也就是說，官方文檔中只有uff和caffe這兩條路可以走通。本文中將詳細說明tensorflow->uff->tensorRT這個路徑下的模型轉(zhuǎn)換，自定義插件的開發(fā)，最終部署這三個方面。

1，tensorflow->uff

tensorflow freeze的pb模型轉(zhuǎn)換默認數(shù)據(jù)格式為NHWC，而在tensorRT中的數(shù)據(jù)格式默認為NCHW。但是不必擔心，pb模型通過tensorRT轉(zhuǎn)換工具轉(zhuǎn)換后的uff模型，在解析階段會自動解析為NCHW的算子。所以只需要將輸入的格式設置成NCHW即可。

轉(zhuǎn)換命令：

convert-to-uff frozen_inference_graph.pb -p config.py -t

轉(zhuǎn)換的配置文件config.py中定義了網(wǎng)絡中不支持的算子與自定義算子的對應關系：

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重新進行了映射。

經(jīng)過轉(zhuǎn)換之后會生成frozen_inference_graph.uff的文件，使用Netron查看，對應的input與upsample的算子類型已經(jīng)改變。

總結

以上是生活随笔為你收集整理的tensorrt轻松部署高性能dnn推理_部署环境之：tensorRT的插件的全部內(nèi)容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。