TensorRT的自定義運算元Plugin的實現

2021-10-06 18:43:56 字數 1867 閱讀 1516

這篇文章主要介紹了如何使用tensorrt實現自定義運算元。

note:

我使用的是tensorrt7.0,自定義運算元使用的ipluginv2ioext實現的。

模型框架是caffe,所以以下實現都只適用於caffe模型的解析,但理論上解析tf和onnx的改動不大。

實現細節不方便全部貼出,但是基本實現過程和結構都在下面了,照著寫寫沒啥問題了。

其實自定義運算元寫多了發現其實還挺好寫的,格式都差不多,主要區別是enqueue的前向計算邏輯可能寫起來複雜些。

整個實現過程基本上是:

繼承nvinfer1::ipluginv2ioext,並實現相應的虛函式。

繼承nvinfer1::iplugincreator並實現相應的虛函式。

繼承nvcaffeparser1::ipluginfactoryv2並實現相應的虛函式。

在解析網路之前呼叫register_tensorrt_plugin註冊upsamplecreator和呼叫parser->setpluginfactoryv2()以使用自定義層型別。

以upsample為例,tensorrt不支援caffe的upsample層,所以這裡實現了乙個自定義層型別,即plugin。需要實現:

upsample類,繼承自nvinfer1::ipluginv2ioext。

upsamplecreator類,繼承自nvinfer1::iplugincreator。

caffepluginfactory類,繼承自nvcaffeparser1::ipluginfactoryv2。

需要實現的函式詳見如下**段。

upsample類的實現:

class

upsample

:public nvinfer1::ipluginv2ioext

// 在這裡new乙個該自定義型別並返回

nvinfer1::ipluginv2ext*

clone()

const override;

// 設定命名空間,用來在網路中查詢和建立plugin

void

setpluginnamespace

(const

char

* lib_namespace) override;

// 返回plugin物件的命名空間

const

char

*getpluginnamespace()

const override;

bool

isoutputbroadcastacrossbatch

(int output_index,

const

bool

* input_is_broadcasted,

int num_inputs)

const override;

bool

canbroadcastinputacrossbatch

(int input_index)

const override;

}

下面是對應的creator類的實現

class

upsamplecreator

:public nvinfer1::iplugincreator

下面是對應的plugin factory類的實現

class

caffepluginfactory

:public nvcaffeparser1::ipluginfactoryv2

Tensor RT 4 自定義層 layers

tensor rt可以直接解析caffe模型。對於不支援的操作op,您可以使用介面手動新增它們 本部落格將提供了乙個廣播操作 broadcast op 的示例,此外,它還測試了pooling layer,它還新增了乙個test layer,它可以單獨列印結構中的引數並協助debug 提供的介面類內部...

js 運算元組函式 自定義

1.刪除陣列中指定的元素 參考例項 foreach function obj,insp obj是要處理的陣列,obj null 表示物件尚未存在 obj.constructor array 表示物件obj的屬性的建構函式不是陣列 constructor屬性始終指向建立當前物件的建構函式。兩個條件均滿...

pytorch 官方介面 自定義運算元註冊

pytorch官方提供了註冊自己的自定義運算元的介面,不需要像native functions.yaml那樣每次改原始碼,再編譯 在pytorch倉庫下找乙個地方,建立新的資料夾 我建議在 pytorch aten src aten core op registration下建立 建立的資料夾假設叫...