04. TensorRT Custom Layer 만들기

TensorRT는 많은 종류의 layer들을 지원하고 지속적으로 확장하고 있습니다.

하지만 지원하는 layer가 특정한 상황에 맞지 않을 수도 있습니다.

그렇기 때문에 TensorRT는 Custom layer를 만들 수 있고 이를 위해 plugin을 제공합니다.

 

plugin에는 위와 같이 3종류가 있고 각각의 특성이 있습니다.

plugin을 사용하기 위해서는 PluginCreator를 사용하면 됩니다.

PluginCreator안에는 Name, Version, plugin field prameters 가 들어 있습니다.

Custom layer를 만들기 위해서는 libnvinfer_plugin.so library가 있어야 합니다.

 

 

 

---

IPluginV2 내부 함수 살펴보기

plugin layer를 연결하고 input, output data의 구조를 셋팅하는 IPluginV2 입니다.

IPluginV2 에는 굉장히 많은 함수가 있습니다.

Custom layer를 만들때 대부분의 method들은 모두 공통적으로 사용됩니다.

 

https://docs.nvidia.com/deeplearning/tensorrt/api/c_api/classnvinfer1_1_1_i_plugin_v2_i_o_ext.html#a9168d66dfaf07ccfdd058d49a25b3c9a

TensorRT: nvinfer1::IPluginV2IOExt Class Reference

위는 공식 사이트에서 제공하는 docs로서 함수들이 설명되어 있는 사이트 입니다.

사이트를 토대로 간단하게 정리해보았습니다.

아래에 이상하거나 틀린 내용 있으면 댓글 부탁드립니다. ㅠ

 

 

 

⊙  getNbOutputs ☆

    - 원본 : Get the number of outputs from the layer.

    - output tensor의 개수.

 

 

⊙  getOutputDimensions ☆

    - 원본 : Get the dimension of an output tensor.

    - output tensor의 차원.

 

 

⊙ supportsFormat   ,   supportsFormatCombination

    - 원본 : Check format support.

    - plugin에서 data type을 지원하는지에 대한 bool값

 

 

⊙  getOutputDataType

    - 원본 : Return the DataType of the plugin output at the requested index.

    - output의 data type을 반환해 줍니다.

 

 

⊙  configurePlugin

    - 원본 : Configure the layer with input and output data types.

    - layer들의 input, output의 data type을 구성합니다.

 

 

⊙  initialize ☆

    - 원본 : Initialize the layer for execution. This is called when the engine is created.

    - 실행을 위해 layer를 초기화 시켜줍니다.

    - engine이 생성될 때, 호출됩니다.

 

 

⊙  enqueue ☆

    - 원본 : Execute the layer.

    - layer를 실행시켜 줍니다.

 

 

⊙  terminate

    - 원본 : Release resources acquired during plugin layer initialization.

                This is called when the engine is destroyed.

    - 엔진을 더이상 사용하지 않을때 호출되고 엔진을 없애줍니다.

 

 

⊙  clone

    - 원본 : Clone the plugin object.

    - plugin을 복사해서 반환합니다.

 

 

⊙  destroy

    - 원본 : Destroy the plugin object. 

    - plugin object를 없애줍니다.

 

 

⊙  set / getPluginNamespace 

    - plugin의 name을 설정해주거나 가져옵니다.

 

 

⊙  canBroadcastInputAcrossBatch ???? 이건 무슨 소린지... 잘 모르겠다...

    - 원본 : Return true if plugin can use input that is broadcast across batch without replication.

    - plugin이 복제 없이 broadcast 되는 input을 사용할 수 있으면 True값을 반환합니다.

 

 

⊙  isOutputBroadcastAcrossBatch ???? 이건 무슨 소린지... 잘 모르겠다...

    - 원본 : Return true if output tensor is broadcast across a batch.

    - output tensor가 broadcast가 된다면  True를 반환합니다.

 

 

---

IPluginCreator 내부 함수 살펴보기

https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#iplugincreator

Developer Guide :: NVIDIA Deep Learning TensorRT Documentation

위는 공식 사이트에서 제공하는 docs로서 함수들이 설명되어 있는 사이트 입니다.

사이트를 토대로 간단하게 정리해보았습니다.

아래에 이상하거나 틀린 내용 있으면 댓글 부탁드립니다. ㅠ

 

 

 

⊙  getPluginName

    - 원본 : Return the plugin name.

    - plugin의 이름을 반환.

 

 

⊙  getPluginVersion

    - 원본 : Return the plugin version.

    - plugin의 버전을 반환.

 

 

⊙ getFieldNames

    - 원본 : Return a list of fields that needs to be passed to createPlugin.

    - createPlugin을 만들기 위해 전달해야 하는 필드 리스트를 반환.

 

 

⊙  createPlugin

    - 원본 : Return a plugin object. Return nullptr in case of error.

    - plugin object를 반환하지만 error일 경우 nullptr을 반환합니다.

 

 

⊙  deserializePlugin

    - 원본 : It should return the plugin object to be used for inference.

    - 추론을 위한 plugin 객체를 반환.

 

 

⊙  set / getPluginNamespace

    - plugin creator의 namespace를 설정하거나 가져옵니다.

 

 

 

---

Custom layer 만들기

 

1. Plugin 만들어 주기

https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#example4_add_custlay_int8

Developer Guide :: NVIDIA Deep Learning TensorRT Documentation

위의 공식 문서에서 제공해주는 Custom pooling layer의 예시를 기준으로 작성하겠습니다.

 

class PoolPlugin : public IPluginV2IOExt
{
    ...override virtual methods inherited from IPluginV2IOExt.
};

 

poolPlugin에서는

• supportsFormatCombination
• configurePlugin
• enqueue

위의 세가지 method를 수정해주어야 합니다.

 

 

 

---

1. network에 Custom layer 추가하기

// Look up the plugin in the registry
IPluginV2 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 
IPluginV2Layer* layer = network.addPluginV2(&inputs[0], int(inputs.size()), pluginObj);
ITensor* data = layer.getOutput(0);
pluginObj->destroy() // Destroy the plugin object

 

 

 

---

 

 

 

 

※ 아래의 github는 공식 사이트에서 제공해주는 예시 코드들이며

Plugin을 만들때는 위에서 복사해서 사용하면 될것 같습니다.

https://github.com/NVIDIA/TensorRT/tree/main/plugin#tensorrt-plugins

GitHub - NVIDIA/TensorRT: TensorRT is a C++ library for high performance inference on NVIDIA GPUs and deep learning accelerators

 

 

※ TensorRT 공식 가이드 문서 : 

https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#overview

https://github.com/NVIDIA/TensorRT/tree/main/plugin#tensorrt-plugins

 

※ 참고 자료:

https://seobway.tistory.com/entry/TensorRT-1-Build-tensorrt-engine-tensorRT-723?category=850806

 

※ 참고자료 : ONNX 변환 및 사용 - 갓꾸준희님

https://eehoeskrap.tistory.com/414

 

※ 참고 github - tensorRT 예제 코드 정리 잘되어 있음!

https://github.com/wang-xinyu/tensorrtx