NVIDIA TensorRT 支持多種類型的層，其功能不斷擴展；但是，在某些情況下，支持的層不能滿足模型的特定需求。

您可以通過實現自定義層（通常稱為插件）來擴展 TensorRT。

9.1. Adding Custom Layers Using The C++ API

您可以通過從 TensorRT 的插件基類之一派生來實現自定義層。Table 3. Base classes, ordered from least expressive to most expressive

為了在網絡中使用插件，您必須首先將其注冊到 TensorRT 的PluginRegistry （ C++ 、 Python ）。不是直接注冊插件，而是為插件注冊一個工廠類的實例，派生自PluginCreator （ C++ ， Python ）。插件創建者類還提供有關插件的其他信息：它的名稱、版本和插件字段參數。

您必須從插件的基類之一派生插件類。在支持具有不同類型/格式的輸入/輸出或具有動態形狀的網絡方面，它們具有不同的表達能力。下表總結了基類，按從最不具表現力到最具表現力的順序排列。

注意：如果插件用于一般用途，請提供 FP32 實現，以使其能夠在任何網絡上正常運行。

TensorRT 提供了一個宏REGISTER_TENSORRT_PLUGIN ，它將插件創建者靜態注冊到注冊表中。

TensorRT 庫包含可以加載到您的應用程序中的插件。有關開源插件的列表，請參閱GitHub：TensorRT 插件。

注意：

要在應用程序中使用 TensorRT 插件，必須加載libnvinfer_plugin.so庫，并且必須通過在應用程序代碼中調用initLibNvInferPlugins來注冊所有插件。

如果您有自己的插件庫，則可以包含一個類似的入口點，以便在唯一命名空間下的注冊表中注冊所有插件。這確保了在構建期間跨不同插件庫的插件名稱沒有沖突。

有關這些插件的更多信息，請參閱NvInferPlugin.h文件以供參考。

調用IPluginCreator：：createPlugin（）返回IPluginV2類型的插件對象。您可以使用addPluginV2（）將插件添加到 TensorRT 網絡，該插件使用給定插件創建網絡層。

例如，您可以將插件層添加到您的網絡，如下所示：

// Look up the plugin in the registry

auto creator = getPluginRegistry（）-》getPluginCreator（pluginName， pluginVersion）;

const PluginFieldCollection* pluginFC = creator-》getFieldNames（）;

// Populate the fields parameters for the plugin layer

PluginFieldCollection *pluginData = parseAndFillFields（pluginFC， layerFields）;

// Create the plugin object using the layerName and the plugin meta data

IPluginV2 *pluginObj = creator-》createPlugin（layerName， pluginData）;

// Add the plugin to the TensorRT network

auto layer = network.addPluginV2（&inputs［0］， int（inputs.size（））， pluginObj）;

… （build rest of the network and serialize engine）

// Destroy the plugin object

pluginObj-》destroy（）

… （destroy network， engine， builder）

… （free allocated pluginData）

注意：

pluginData必須在傳遞給createPlugin之前在堆上分配PluginField條目。

前面描述的 createPlugin 方法在堆上創建一個新的插件對象并返回一個指向它的指針。如前所示，確保銷毀 pluginObj 以避免內存泄漏。

IPluginV2類型插件的插件類型、插件版本和命名空間（如果存在） 。在反序列化期間，TensorRT 從插件注冊表中查找插件創建者并調用IPluginCreator：：deserializePlugin（） 。在反序列化過程中創建的插件對象由 TensorRT 引擎通過調用IPluginV2：：destroy（）方法在內部銷毀。

IPluginV2類型插件的插件類型、插件版本和命名空間（如果存在）， 在反序列化期間，TensorRT 從插件注冊表中查找插件創建者并調用IPluginCreator：：deserializePlugin（） 。在反序列化期間創建的插件對象由 TensorRT 引擎通過調用IPluginV2：：destroy（）方法在內部銷毀。

9.1.1. Example： Adding A Custom Layer With Dynamic Shape Support Using C++

要支持動態形狀，您的插件必須從IPluginV2DynamicExt派生。

關于這個任務

BarPlugin是一個有兩個輸入和兩個輸出的插件，其中： ? 第一個輸出是第二個輸入的拷貝 ? 第二個輸出是兩個輸入的串聯，沿著第一個維度，所有類型/格式必須相同并且是線性格式

BarPlugin需要按如下方式派生：

class BarPlugin ： public IPluginV2DynamicExt

{

。..override virtual methods inherited from IPluginV2DynamicExt.

};

繼承的方法都是純虛方法，所以如果你忘記了，編譯器會提醒你。

受動態形狀影響的四種方法是：

獲取輸出維度

支持格式組合

配置插件

隊列

getOutputDimensions的覆蓋根據輸入維度返回輸出維度的符號表達式。您可以使用傳遞給getOutputDimensions的IExprBuilder從輸入的表達式構建表達式。在示例中，不必為案例 1 構建新表達式，因為第二個輸出的維度與第一個輸入的維度相同。

DimsExprs BarPlugin：：getOutputDimensions（int outputIndex，

const DimsExprs* inputs， int nbInputs，

IExprBuilder& exprBuilder）

{

switch （outputIndex）

{

case 0：

{

// First dimension of output is sum of input

// first dimensions.

DimsExprs output（inputs［0］）;

output.d［0］ =

exprBuilder.operation（DimensionOperation：：kSUM，

inputs［0］.d［0］， inputs［1］.d［0］）;

return output;

}

case 1：

return inputs［0］;

default：

throw std：：invalid_argument（“invalid output”）;

}

supportsFormatCombination的覆蓋必須指示是否允許格式組合。接口將輸入/輸出統一索引為“connections”，從第一個輸入的 0 開始，然后依次為其余輸入，然后為輸出編號。在示例中，輸入是connections 0 和 1，輸出是connections 2 和 3。

TensorRT 使用supportsFormatCombination來詢問給定的格式/類型組合是否適用于連接，給定的格式/類型用于索引較少的連接。因此，覆蓋可以假設較少索引的連接已經過審查，并專注于與索引pos的連接。

bool BarPlugin：：supportsFormatCombination（int pos， const PluginTensorDesc* inOut， int nbInputs， int nbOutputs） override

{

assert（0 《= pos && pos 《 4）;

const auto* in = inOut;

const auto* out = inOut + nbInputs;

switch （pos）

{

case 0： return in［0］.format == TensorFormat：：kLINEAR;

case 1： return in［1］.type == in［0］.type &&

in［0］.format == TensorFormat：：kLINEAR;

case 2： return out［0］.type == in［0］.type &&

out［0］.format == TensorFormat：：kLINEAR;

case 3： return out［1］.type == in［0］.type &&

out［1］.format == TensorFormat：：kLINEAR;

}

throw std：：invalid_argument（“invalid connection number”）;

}

這里的局部變量in和out允許通過輸入或輸出編號而不是連接編號檢查inOut 。

重要提示：覆蓋檢查索引小于pos的連接的格式/類型，但絕不能檢查索引大于pos的連接的格式/類型。該示例使用case 3來檢查連接 3 和連接 0，而不是使用case 0來檢查連接 0 和連接 3。

TensorRT 使用configurePlugin在運行時設置插件。這個插件不需要configurePlugin來做任何事情，所以它是一個空操作：

void BarPlugin：：configurePlugin（

const DynamicPluginTensorDesc* in， int nbInputs，

const DynamicPluginTensorDesc* out， int nbOutputs） override

{

}

如果插件需要知道它可能遇到的最小或最大尺寸，它可以檢查字段DynamicPluginTensorDesc：：min或DynamicPluginTensorDesc：：max的任何輸入或輸出。格式和構建時維度信息可以在DynamicPluginTensorDesc：：desc中找到。任何運行時維度都顯示為 -1。實際維度提供給BarPlugin：：enqueue 。

最后，重寫BarPlugin：：enqueue必須完成這項工作。由于形狀是動態的，因此 enqueue 會收到一個PluginTensorDesc ，它描述了每個輸入和輸出的實際尺寸、類型和格式。

9.1.2. Example： Adding A Custom Layer With INT8 I/O Support Using C++

PoolPlugin是一個插件，用于演示如何為自定義池層擴展 INT8 I/O。推導如下：

class PoolPlugin ： public IPluginV2IOExt

{

。..override virtual methods inherited from IPluginV2IOExt.

};

大多數純虛方法對插件來說都是通用的。影響INT8 I/O的主要方法有：

支持格式組合

配置插件

enqueue

supportsFormatCombination的覆蓋必須指示允許哪個 INT8 I/O 組合。此接口的用法類似于示例：使用 C++ 添加具有動態形狀支持的自定義層。在本例中，支持的 I/O 張量格式為線性 CHW，數據類型為 FP32、FP16 或 INT8，但 I/O 張量必須具有相同的數據類型。

bool PoolPlugin：：supportsFormatCombination（int pos， const PluginTensorDesc* inOut， int nbInputs， int nbOutputs） const override

{

assert（nbInputs == 1 && nbOutputs == 1 && pos 《 nbInputs + nbOutputs）;

bool condition = inOut［pos］.format == TensorFormat：：kLINEAR;

condition &= （（inOut［pos］.type == DataType：：kFLOAT） ||

（inOut［pos］.type == DataType：：kHALF） ||

（inOut［pos］.type == DataType：：kINT8））;

condition &= inOut［pos］.type == inOut［0］.type;

return condition;

}

重要的：

如果 INT8 校準必須與具有 INT8 I/O 插件的網絡一起使用，則該插件必須支持 FP32 I/O，因為它被 FP32 校準圖使用。

如果不支持 FP32 I/O 變體或未使用 INT8 校準，則必須明確設置所有必需的 INT8 I/O 張量尺度。

校準無法確定插件內部張量的動態范圍。對量化數據進行操作的插件必須為內部張量計算自己的動態范圍。

TensorRT 調用configurePlugin方法通過PluginTensorDesc將信息傳遞給插件，這些信息存儲為成員變量，序列化和反序列化。

void PoolPlugin：：configurePlugin（const PluginTensorDesc* in， int nbInput， const PluginTensorDesc* out， int nbOutput）

{

。..

mPoolingParams.mC = mInputDims.d［0］;

mPoolingParams.mH = mInputDims.d［1］;

mPoolingParams.mW = mInputDims.d［2］;

mPoolingParams.mP = mOutputDims.d［1］;

mPoolingParams.mQ = mOutputDims.d［2］;

mInHostScale = in［0］.scale 》= 0.0f ？ in［0］.scale ： -1.0f;

mOutHostScale = out［0］.scale 》= 0.0f ？ out［0］.scale ： -1.0f;

}

每個張量的 INT8 I/O 比例可以從PluginTensorDesc：：scale獲得。 最后，重寫UffPoolPluginV2：：enqueue必須完成這項工作。它包括一組核心算法，可在運行時通過使用實際批量大小、輸入、輸出、cuDNN 流和配置的信息來執行自定義層。

int PoolPlugin：：enqueue（int batchSize， const void* const* inputs， void** outputs， void* workspace， cudaStream_t stream）

{

。..

CHECK（cudnnPoolingForward（mCudnn， mPoolingDesc， &kONE， mSrcDescriptor， input， &kZERO， mDstDescriptor， output））;

。..

return 0;

}

9.2. Adding Custom Layers Using The Python API

盡管 C++ API 是實現自定義層的首選語言，但由于可以訪問 CUDA 和 cuDNN 等庫，您還可以在 Python 應用程序中使用自定義層。

您可以使用 C++ API 創建自定義層，在 Python 中使用pybind11打包層，然后將插件加載到 Python 應用程序中。有關更多信息，請參閱在 Python 中創建網絡定義。

相同的自定義層實現可用于 C++ 和 Python。

9.2.1. Example： Adding A Custom Layer To A TensorRT Network Using Python

可以使用插件節點將自定義層添加到 Python 中的任何 TensorRT 網絡。

Python API 有一個名為add_plugin_v2的函數，可讓您將插件節點添加到網絡。以下示例說明了這一點。它創建了一個簡單的TensorRT網絡，并通過查找TensorRT插件注冊表來添加一個 Leaky ReLU 插件節點。

import tensorrt as trt

import numpy as np

TRT_LOGGER = trt.Logger（）

trt.init_libnvinfer_plugins（TRT_LOGGER， ‘’）

PLUGIN_CREATORS = trt.get_plugin_registry（）.plugin_creator_list

def get_trt_plugin（plugin_name）：

plugin = None

for plugin_creator in PLUGIN_CREATORS：

if plugin_creator.name == plugin_name：

lrelu_slope_field = trt.PluginField（“neg_slope”， np.array（［0.1］， dtype=np.float32）， trt.PluginFieldType.FLOAT32）

field_collection = trt.PluginFieldCollection（［lrelu_slope_field］）

plugin = plugin_creator.create_plugin（name=plugin_name， field_collection=field_collection）

return plugin

def main（）：

builder = trt.Builder（TRT_LOGGER）

network = builder.create_network（）

config = builder.create_builder_config（）

config.max_workspace_size = 2**20

input_layer = network.add_input（name=“input_layer”， dtype=trt.float32， shape=（1， 1））

lrelu = network.add_plugin_v2（inputs=［input_layer］， plugin=get_trt_plugin（“LReLU_TRT”））

lrelu.get_output（0）.name = “outputs”

network.mark_output（lrelu.get_output（0））

9.3. Using Custom Layers When Importing A Model With A Parser

ONNX 解析器會自動嘗試將無法識別的節點作為插件導入。如果在插件注冊表中找到與節點具有相同op_type的插件，則解析器將節點的屬性作為插件字段參數轉發給插件創建者，以創建插件。默認情況下，解析器使用“1”作為插件版本， “”作為插件命名空間。可以通過在相應的 ONNX 節點中設置plugin_version或plugin_namespace字符串屬性來覆蓋此行為。

在某些情況下，您可能希望在將 ONNX 圖導入 TensorRT 之前對其進行修改。例如，用插件節點替換一組操作。為此，您可以使用ONNX GraphSurgeon 實用程序。有關如何使用 ONNX-GraphSurgeon 替換子圖的詳細信息，請參閱此示例。

有關更多示例，請參閱onnx_packnet示例。

9.4. Plugin API Description

所有新插件都應從IPluginCreator和使用 C++ API 添加自定義層中描述的插件基類之一派生類。此外，新插件還應調用REGISTER_TENSORRT_PLUGIN（。..）宏以將插件注冊到 TensorRT 插件注冊表或創建等效于initLibNvInferPlugins（）的init函數。

9.4.1. Migrating Plugins From TensorRT 6.x Or 7.x To TensorRT 8.x.x

IPluginV2和IPluginV2Ext以分別向后兼容 TensorRT 5.1 和 6.0.x。但是，新插件應針對IPluginV2DynamicExt或IPluginV2IOExt接口，而舊插件應重構以使用這些接口。

IPluginV2DynamicExt中的新特性如下：

virtual DimsExprs getOutputDimensions（int outputIndex， const DimsExprs* inputs， int nbInputs， IExprBuilder& exprBuilder） = 0;

virtual bool supportsFormatCombination（int pos， const PluginTensorDesc* inOut， int nbInputs， int nbOutputs） = 0;

virtual void configurePlugin（const DynamicPluginTensorDesc* in， int nbInputs， const DynamicPluginTensorDesc* out， int nbOutputs） = 0;

virtual size_t getWorkspaceSize（const PluginTensorDesc* inputs， int nbInputs， const PluginTensorDesc* outputs， int nbOutputs） const = 0;

virtual int enqueue（const PluginTensorDesc* inputDesc， const PluginTensorDesc* outputDesc， const void* const* inputs， void* const* outputs， void* workspace， cudaStream_t stream） = 0;

IPluginV2IOExt中的新特性如下：

virtual void configurePlugin（const PluginTensorDesc* in， int nbInput， const PluginTensorDesc* out， int nbOutput） = 0;

virtual bool supportsFormatCombination（int pos， const PluginTensorDesc* inOut， int nbInputs， int nbOutputs） const = 0;

遷移到IPluginV2DynamicExt或IPluginV2IOExt的指南：

getOutputDimensions實現給定輸入的輸出張量維度的表達式。

supportsFormatCombination檢查插件是否支持指定輸入/輸出的格式和數據類型。

configurePlugin模仿IPluginV2Ext中等效的configurePlugin的行為，但接受張量描述符。

getWorkspaceSize和enqueue模仿IPluginV2Ext中等效 API 的行為，但接受張量描述符。

更多詳細信息，請參閱IPluginV2 API 說明中的 API 說明。

9.4.2. IPluginV2 API Description

以下部分介紹IPluginV2類的功能。要將插件層連接到相鄰層并設置輸入和輸出數據結構，構建器通過調用以下插件方法檢查輸出的數量及其維度。

getNbOutputs

用于指定輸出張量的數量。

getOutputDimensions

用于將輸出的維度指定為輸入維度的函數。

supportsFormat

用于檢查插件是否支持給定的數據格式。

getOutputDataType

用于獲取給定索引處輸出的數據類型。返回的數據類型必須具有插件支持的格式。

插件層可以支持四種數據格式，例如：

NCHW單精度 （FP32）、半精度 （FP16） 和整型 （INT32） 張量

NC / 2HW2和NHWC8半精度 （FP16） 張量

格式由PluginFormatType枚舉。

除了輸入和輸出張量之外，不計算所有數據并且需要內存空間的插件可以使用getWorkspaceSize方法指定額外的內存需求，該方法由構建器調用以確定和預分配暫存空間。

在構建和推理期間，可能會多次配置和執行插件層。在構建時，為了發現最佳配置，層被配置、初始化、執行和終止。為插件選擇最佳格式后，再次配置插件，然后在推理應用程序的生命周期內初始化一次并執行多次，最后在引擎銷毀時終止。這些步驟由構建器和引擎使用以下插件方法控制：

configurePlugin（配置插件）

傳達輸入和輸出的數量、所有輸入和輸出的維度和數據類型、所有輸入和輸出的廣播信息、選擇的插件格式和最大批量大小。此時，插件設置其內部狀態并為給定配置選擇最合適的算法和數據結構。

initialize（初始化）

此時配置是已知的，并且正在創建推理引擎，因此插件可以設置其內部數據結構并準備執行。

enqueue（排隊）

封裝插件的實際算法和內核調用，并提供運行時批處理大小、指向輸入、輸出和暫存空間的指針，以及用于內核執行的CUDA流。

terminate（終止）

引擎上下文被銷毀，插件持有的所有資源必須被釋放。

clone（克?。?/p>

每次創建包含此插件層的新構建器、網絡或引擎時都會調用它。它必須返回一個帶有正確參數的新插件對象。

destroy（銷毀）

用于銷毀插件對象和/或每次創建新插件對象時分配的其他內存。每當構建器或網絡或引擎被破壞時都會調用它。

“`set/getPluginNamespace`（設置/獲取插件命名空間）“

該方法用于設置該插件對象所屬的庫命名空間（默認可以是“”）。來自同一個插件庫的所有插件對象都應該具有相同的命名空間。

IPluginV2Ext支持可以處理廣播輸入和輸出的插件。此功能需要實現以下方法：

canBroadcastInputAcrossBatch

對每個輸入調用此方法，其張量在批次中進行語義廣播。如果canBroadcastInputAcrossBatch返回true （意味著插件可以支持廣播），則 TensorRT 不會復制輸入張量。插件應該在批處理中共享一個副本。如果它返回false ，則 TensorRT 會復制輸入張量，使其看起來像一個非廣播張量。

isOutputBroadcastAcrossBatch

這為每個輸出索引調用。該插件應在給定索引處返回 true 輸出，并在整個批次中廣播。

IPluginV2IOExt 這由構建器在initialize（）之前調用。它為層提供了基于 I/O PluginTensorDesc和最大批量大小進行算法選擇的機會。

注意：基于IPluginV2的插件在引擎級別共享，而不是在執行上下文級別共享，因此這些可能被多個線程同時使用的插件需要以線程安全的方式管理它們的資源。創建ExecutionContext時會克隆基于IPluginV2Ext和派生接口的插件，因此這不是必需的。

9.4.3. IPluginCreator API Description

IPluginCreator類中的以下方法用于從插件注冊表中查找和創建適當的插件：

getPluginName

IPluginExt：：getPluginType的返回值。

getPluginVersion 返回插件版本。對于所有內部 TensorRT 插件，默認為1 。

getFieldNames 要成功創建插件，需要了解插件的所有字段參數。此方法返回PluginFieldCollection結構，其中填充了PluginField條目以反映字段名稱和PluginFieldType （數據應指向nullptr ）。

createPlugin

此方法用于使用PluginFieldCollection參數創建插件。應填充PluginField條目的數據字段以指向每個插件字段條目的實際數據。

注意：傳遞給createPlugin函數的數據應該由調用者分配，并在程序被銷毀時最終由調用者釋放。 createPlugin函數返回的插件對象的所有權被傳遞給調用者，并且也必須被銷毀。

deserializePlugin

此方法由 TensorRT 引擎根據插件名稱和版本在內部調用。它應該返回要用于推理的插件對象。在該函數中創建的插件對象在引擎被銷毀時被 TensorRT 引擎銷毀。

set/getPluginNamespace 該方法用于設置此創建者實例所屬的命名空間（默認可以是“”）。

9.5. Best Practices For Custom Layers Plugin

調試自定義層問題

必須釋放插件中分配的內存以確保沒有內存泄漏。如果在initialize（）函數中獲取資源，則需要在terminate（）函數中釋放它們。應該釋放所有其他內存分配，最好在插件類析構函數中或在destroy（）方法中。使用 C++ API 添加自定義層詳細概述了這一點，并提供了一些使用插件時的最佳實踐說明。

關于作者

Ken He 是 NVIDIA 企業級開發者社區經理 & 高級講師，擁有多年的 GPU 和人工智能開發經驗。自 2017 年加入 NVIDIA 開發者社區以來，完成過上百場培訓，幫助上萬個開發者了解人工智能和 GPU 編程開發。在計算機視覺，高性能計算領域完成過多個獨立項目。并且，在機器人和無人機領域，有過豐富的研發經驗。對于圖像識別，目標的檢測與跟蹤完成過多種解決方案。曾經參與 GPU 版氣象模式GRAPES，是其主要研發者。

審核編輯：郭婷