??在前面的模型入門系列文章中，我們介紹了部署一個 PyTorch 模型到推理后端，如 ONNXRuntime，這其中可能遇到很多工程性的問題。
有些可以通過創建 ONNX 節點來解決，該節點仍然使用后端原生的實現進行推理。而有些無法導出到后端的算法，可以通過重寫代碼改變算法的實現過程，同樣可以導出到 ONNX ，達到一致的效果。以上兩種方式一般可以處理絕大多數的部署問題，同時也不需要向推理框架引入新的內容，是我們進行模型部署時候的優先選擇。
然而，仍然存在部分模型，模型中某些算子無法通過上述兩種方式繞過問題，這時候，如何對特定后端實現對應代碼就極為重要。這也是本文將介紹的第三種方式——自定義插件。
自定義插件是很多推理框架支持用戶自定義算子的方式，以 MMDeploy 為例，它是一個支持多種推理后端的算法庫。目前支持的后端有：

• ONNXRuntime
• TensorRT
• ncnn
• OpenVINO
• PPLNN

??其中，前三種后端均實現了一些自定義的算子。例如 ONNXRuntime 中的調制可變性卷積，ncnn 中的topk 算子，TensorRT 中的 MultiLevelRoiAlign 。
介紹如何給后端自定義算子是一件相對復雜的事情，所以本文只針對其中一種后端 TensorRT，介紹自定義算子。如果讀者對其他后端感興趣，可以去他們的代碼庫查看，一般地，各個推理框架均有詳細文檔介紹如何添加客制化的算子實現。

1. 在 MMDeploy 添加 TensorRT 插件

??仍然以前面教程二中的超分辨模型 SRCNN 為例。在教程二中，我們用 ONNXRuntime 作為后端，通過 PyTorch 的 symbolic 函數導出了一個支持動態 scale 的 ONNX 模型，這個模型可以直接用 ONNXRuntime 運行，這是因為 NewInterpolate 類導出的節點 Resize 就是 ONNXRuntime 支持的節點。下面我們嘗試直接將教程二導出的 srcnn3.onnx 轉換到 TensorRT。

from mmdeploy.backend.tensorrt.utils import from_onnx from_onnx( 'srcnn3.onnx', 'srcnn3', input_shapes=dict( input=dict( min_shape=[1, 3, 256, 256], opt_shape=[1, 3, 256, 256], max_shape=[1, 3, 256, 256]), factor=dict( min_shape=[4], opt_shape=[4], max_shape=[4]))) 

??沒有安裝過 MMDeploy 的小伙伴可以先參考 build.md 進行安裝，安裝完成后執行上述腳本，會有如下報錯：

RuntimeError: Failed to parse onnx, In node 1 (importResize): UNSUPPORTED_NODE: Assertion failed: mode != "cubic" && "This version of TensorRT does not support cubic interpolation!" 

報錯的原因有以下兩方面：

??'srcnn3.onnx’文件中的Resize 是 ONNX 原生節點。其插值方式之一 bicubic 并不被 TensorRT 支持（TensorRT 的 Resize Layer僅支持 nearest 和 bilinear 兩種插值方式）。日志的錯誤信息也明確提示了這點；
但即便將 “bicubic” 模式改為 “bilinear” ，轉換仍然失敗: RuntimeError: Failed to parse onnx, In node 1 (importResize): UNSUPPORTED_NODE: Assertion failed: scales.is_weights() && Resize scales must be initializer!"。這是因為 TensorRT 無法接受動態 scale 導致的。

2. 創建 ONNX 節點

??為解決上述問題，我們需要創建一個新的節點替換原生 Resize 節點，并且實現新節點對應的插件代碼。
新改節點名稱就叫 Test::DynamicTRTResize，這是種類 C++ 的寫法，Test 為域名，主要用于區分不同來源下的同名的節點，比如 ONNX:: 和 Test::。當然了，ONNX 本身也不存在 DynamicTRTResize 的節點名。

import torch 
from torch import nn 
from torch.nn.functional import interpolate 
import torch.onnx 
import cv2 
import numpy as np 
import os, requests 
# Download checkpoint and test image 
urls = ['https://download.openmmlab.com/mmediting/restorers/srcnn/srcnn_x4k915_1x16_1000k_div2k_20200608-4186f232.pth', 'https://raw.githubusercontent.com/open-mmlab/mmediting/master/tests/data/face/000001.png'] 
names = ['srcnn.pth', 'face.png'] 
for url, name in zip(urls, names): if not os.path.exists(name): open(name, 'wb').write(requests.get(url).content) 
class DynamicTRTResize(torch.autograd.Function): def __init__(self) -> None: super().__init__() @staticmethod def symbolic(g, input, size_tensor, align_corners = False): """Symbolic function for creating onnx op.""" return g.op( 'Test::DynamicTRTResize', input, size_tensor, align_corners_i=align_corners) @staticmethod def forward(g, input, size_tensor, align_corners = False): """Run forward.""" size = [size_tensor.size(-2), size_tensor.size(-1)] return interpolate( input, size=size, mode='bicubic', align_corners=align_corners) 
class StrangeSuperResolutionNet(nn.Module): def __init__(self): super().__init__() self.conv1 = nn.Conv2d(3, 64, kernel_size=9, padding=4) self.conv2 = nn.Conv2d(64, 32, kernel_size=1, padding=0) self.conv3 = nn.Conv2d(32, 3, kernel_size=5, padding=2) self.relu = nn.ReLU() def forward(self, x, size_tensor): x = DynamicTRTResize.apply(x, size_tensor) out = self.relu(self.conv1(x)) out = self.relu(self.conv2(out)) out = self.conv3(out) return out 
def init_torch_model(): torch_model = StrangeSuperResolutionNet() state_dict = torch.load('srcnn.pth')['state_dict'] # Adapt the checkpoint for old_key in list(state_dict.keys()): new_key = '.'.join(old_key.split('.')[1:]) state_dict[new_key] = state_dict.pop(old_key) torch_model.load_state_dict(state_dict) torch_model.eval() return torch_model 
model = init_torch_model() 
factor = torch.rand([1, 1, 512, 512], dtype=torch.float) 
input_img = cv2.imread('face.png').astype(np.float32) 
# HWC to NCHW 
input_img = np.transpose(input_img, [2, 0, 1]) 
input_img = np.expand_dims(input_img, 0) 
# Inference 
torch_output = model(torch.from_numpy(input_img), factor).detach().numpy() 
# NCHW to HWC 
torch_output = np.squeeze(torch_output, 0) 
torch_output = np.clip(torch_output, 0, 255) 
torch_output = np.transpose(torch_output, [1, 2, 0]).astype(np.uint8) 
# Show image 
cv2.imwrite("face_torch.png", torch_output) 
x = torch.randn(1, 3, 256, 256) 
dynamic_axes={ 'input': { 0: 'batch', 2: 'height', 3: 'width' }, 'factor': { 0: 'batch1', 2: 'height1', 3: 'width1' }, 'output': { 0: 'batch2', 2: 'height2', 3: 'width2' }, } 
with torch.no_grad(): torch.onnx.export( model, (x, factor), "srcnn3.onnx", opset_version=11, input_names=['input', 'factor'], output_names=['output'], dynamic_axes=dynamic_axes) 

??執行上述腳本，我們成功導出了一個 ONNX 模型 srcnn.onnx。用 netron 打開這個模型可視化如下：

直接將該模型轉換成 TensorRT 模型也是不可行的，這是因為 TensorRT 還無法解析 DynamicTRTResize 節點。而想要解析該節點，我們必須為 TensorRT 添加 c++ 代碼，實現該插件。

3. C++ 實現

??因為 MMDeploy 中已經實現了 Bicubic Interpolate 算子，所以我們可以復用其中的 CUDA 部分代碼，只針對 TensorRT 實現支持動態 scale 的插件即可。對 CUDA 編程感興趣的小伙伴可以參考 CUDA 的官方教程。
因為 csrc/backend_ops/tensorrt/bicubic_interpolate 中有我們需要的 CUDA 代碼，所以我們可以直接在該文件夾加添加 TensorRT 相關的 trt_dynamic_resize.hpp 和 trt_dynamic_resize.cpp 文件，在這兩個文件中分別聲明和實現插件就可以了。我們也可以新建文件夾 csrc/backend_ops/tensorrt/dynamic_resize，將這兩個文件直接放到這個文件夾下。
對 TensorRT 7+，要實現這樣一個自定義插件，我們需要寫兩個類。

• DynamicTRTResize，繼承自 nvinfer1::IPluginV2DynamicExt，完成插件的具體實現。
• DynamicTRTResizeCreator，繼承自 nvinfer1::IPluginCreator，是插件的工廠類，用于創建 DynamicTRTResize 插件的實例。

??在 MMDeploy 中，由于有若干插件需要實現，所以我們在mmdeploy/csrc/backend_ops/tensorrt/common/trt_plugin_base.hpp 中實現了 TRTPluginBase 和 TRTPluginCreatorBase 兩個類，用于管理一些所有插件共有的屬性方法。
其中，TRTPluginBase 繼承自 nvinfer1::IPluginV2DynamicExt，而TRTPluginCreatorBase繼承自nvinfer1::IPluginCreator。這樣，用戶實現插件時只需繼承這兩個新的類即可。所以我們只需在 dynamic_resize 文件夾下的 .hpp 文件中，引用 trt_plugin_base.hpp 頭文件，繼承邏輯如下：

class DynamicTRTResize : public TRTPluginBase{} 
class DynamicTRTResizeCreator : public TRTPluginCreatorBase{} 

??在 trt_dynamic_resize.hpp 中，我們聲明如下內容：

#ifndef TRT_DYNAMIC_RESIZE_HPP 
#define TRT_DYNAMIC_RESIZE_HPP 
#include <cublas_v2.h> 
#include <memory> 
#include <string> 
#include <vector> 
#include "trt_plugin_base.hpp" 
namespace mmdeploy { 
class DynamicTRTResize : public TRTPluginBase { public: DynamicTRTResize(const std::string &name, bool align_corners); DynamicTRTResize(const std::string name, const void *data, size_t length); DynamicTRTResize() = delete; // IPluginV2DynamicExt Methods nvinfer1::IPluginV2DynamicExt *clone() const TRT_NOEXCEPT override; nvinfer1::DimsExprs getOutputDimensions(int outputIndex, const nvinfer1::DimsExprs *inputs, int nbInputs, nvinfer1::IExprBuilder &exprBuilder) TRT_NOEXCEPT override; bool supportsFormatCombination(int pos, const nvinfer1::PluginTensorDesc *ioDesc, int nbInputs, int nbOutputs) TRT_NOEXCEPT override; void configurePlugin(const nvinfer1::DynamicPluginTensorDesc *in, int nbInputs, const nvinfer1::DynamicPluginTensorDesc *out, int nbOutputs) TRT_NOEXCEPT override; size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc *inputs, int nbInputs, const nvinfer1::PluginTensorDesc *outputs, int nbOutputs) const TRT_NOEXCEPT override; int enqueue(const nvinfer1::PluginTensorDesc *inputDesc, const nvinfer1::PluginTensorDesc *outputDesc, const void *const *inputs, void *const *outputs, void *workspace, cudaStream_t stream) TRT_NOEXCEPT override; // IPluginV2Ext Methods nvinfer1::DataType getOutputDataType(int index, const nvinfer1::DataType *inputTypes, int nbInputs) const TRT_NOEXCEPT override; // IPluginV2 Methods const char *getPluginType() const TRT_NOEXCEPT override; const char *getPluginVersion() const TRT_NOEXCEPT override; int getNbOutputs() const TRT_NOEXCEPT override; size_t getSerializationSize() const TRT_NOEXCEPT override; void serialize(void *buffer) const TRT_NOEXCEPT override; private: bool mAlignCorners; 
}; 
class DynamicTRTResizeCreator : public TRTPluginCreatorBase { public: DynamicTRTResizeCreator(); const char *getPluginName() const TRT_NOEXCEPT override; const char *getPluginVersion() const TRT_NOEXCEPT override; nvinfer1::IPluginV2 *createPlugin(const char *name, const nvinfer1::PluginFieldCollection *fc) TRT_NOEXCEPT override; nvinfer1::IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength) TRT_NOEXCEPT override; 
}; 
}  // namespace mmdeploy 
#endif  // TRT_DYNAMIC_RESIZE_HPP 

??在這樣一份頭文件中，DynamicTRTResize 類進行了如下的套娃繼承：
從上面的圖片和代碼中我們發現，在插件類 DynamicTRTResize 中，我們定義了私有變量 mAlignCorners，該變量表示是否 align corners。此外只要實現構造函數、析構函數和 TensoRT 中三個基類的方法即可。其中構造函數有二，分別用于創建插件和反序列化插件。而基類方法中：

• 基類 IPluginV2DynamicExt 的方法較為值得關注，getOutputDimensions 獲取輸出張量的形狀，enqueue 真正負責執行我們的算法，內部一般會調用 CUDA 核函數。本文實現的插件直接調用 MMDeploy 已定義在 csrc/backend_ops/tensorrt/bicubic_interpolate 的核函數 bicubic_interpolate。
• 基類 IPluginV2Ext 的方法，我們只要實現獲取輸出數據類型的 getOutputDataType 即可。
• 基類 IPluginV2 則是些獲取插件類型和版本號的方法，此外則是序列化輸入插件的參數的函數 serialize 和計算該參數的序列化后 buffer 大小的函數 getSerializationSize，以及獲取輸出張量個數的方法 getNbOutputs。還有部分公共方法被定義在 TRTPluginBase 類內了。

??在插件工廠類 DynamicTRTResizeCreator 中，我們需要聲明獲取插件名稱和版本的方法 getPluginName 和 getPluginVersion。同時我們還需要聲明創建插件和反序列化插件的方法 createPlugin 和 deserializePlugin，前者調用 DynamicTRTResize 中創建插件的方法，后者調用反序列化插件的方法。
接下來，我們就實現上述聲明吧。在 .cpp 文件中我們實現的代碼如下：

// Copyright (c) OpenMMLab. All rights reserved 
#include "trt_dynamic_resize.hpp" 
#include <assert.h> 
#include <chrono> 
#include "trt_plugin_helper.hpp" 
#include "trt_serialize.hpp" 
// to get the reference to kernel function bicubic_interpolate，which will be used in enqueue 
#include "../bicubic_interpolate/trt_bicubic_interpolate_kernel.hpp" 
using namespace nvinfer1; 
namespace mmdeploy { 
namespace { 
static const char *PLUGIN_VERSION{"1"}; 
static const char *PLUGIN_NAME{"DynamicTRTResize"};//plagin name == ONNX node name，triggered in building engine 
}  // namespace 
DynamicTRTResize::DynamicTRTResize(const std::string &name, bool align_corners) : TRTPluginBase(name), mAlignCorners(align_corners) {} 
DynamicTRTResize::DynamicTRTResize(const std::string name, const void *data, size_t length) : TRTPluginBase(name) { deserialize_value(&data, &length, &mAlignCorners); 
} 
nvinfer1::IPluginV2DynamicExt *DynamicTRTResize::clone() const TRT_NOEXCEPT { DynamicTRTResize *plugin = new DynamicTRTResize(mLayerName, mAlignCorners); plugin->setPluginNamespace(getPluginNamespace()); return plugin; 
} 
nvinfer1::DimsExprs DynamicTRTResize::getOutputDimensions( int outputIndex, const nvinfer1::DimsExprs *inputs, int nbInputs, nvinfer1::IExprBuilder &exprBuilder) TRT_NOEXCEPT { nvinfer1::DimsExprs ret; ret.nbDims = 4; // input two tensors: input and size_tensor, the later is for shape inference only ret.d[0] = inputs[0].d[0]; ret.d[1] = inputs[0].d[1]; ret.d[2] = inputs[1].d[2]; ret.d[3] = inputs[1].d[3]; return ret; 
} 
bool DynamicTRTResize::supportsFormatCombination(int pos, const nvinfer1::PluginTensorDesc *ioDesc, int nbInputs, int nbOutputs) TRT_NOEXCEPT { if (pos == 0) { return (ioDesc[pos].type == nvinfer1::DataType::kFLOAT && ioDesc[pos].format == nvinfer1::TensorFormat::kLINEAR); } else { return ioDesc[pos].type == ioDesc[0].type && ioDesc[pos].format == ioDesc[0].format; } 
} 
void DynamicTRTResize::configurePlugin(const nvinfer1::DynamicPluginTensorDesc *inputs, int nbInputs, const nvinfer1::DynamicPluginTensorDesc *outputs, int nbOutputs) TRT_NOEXCEPT {} 
size_t DynamicTRTResize::getWorkspaceSize(const nvinfer1::PluginTensorDesc *inputs, int nbInputs, const nvinfer1::PluginTensorDesc *outputs, int nbOutputs) const TRT_NOEXCEPT { return 0; 
} 
int DynamicTRTResize::enqueue(const nvinfer1::PluginTensorDesc *inputDesc, const nvinfer1::PluginTensorDesc *outputDesc, const void *const *inputs, void *const *outputs, void *workSpace, cudaStream_t stream) TRT_NOEXCEPT { int batch = inputDesc[0].dims.d[0]; int channels = inputDesc[0].dims.d[1]; int height = inputDesc[0].dims.d[2]; int width = inputDesc[0].dims.d[3]; int height_out = outputDesc[0].dims.d[2]; int width_out = outputDesc[0].dims.d[3]; const void *x = inputs[0]; void *output = outputs[0]; // TODO: add fp16 support auto data_type = inputDesc[0].type; switch (data_type) { case nvinfer1::DataType::kFLOAT: bicubic_interpolate<float>((float *)x, (float *)output, batch, channels, height, width, height_out, width_out, mAlignCorners, stream); break; default: return 1; break; } return 0; 
} 
nvinfer1::DataType DynamicTRTResize::getOutputDataType(int index, const nvinfer1::DataType *inputTypes, int nbInputs) const TRT_NOEXCEPT { return inputTypes[0]; 
} 
// IPluginV2 Methods 
const char *DynamicTRTResize::getPluginType() const TRT_NOEXCEPT { return PLUGIN_NAME; } 
const char *DynamicTRTResize::getPluginVersion() const TRT_NOEXCEPT { return PLUGIN_VERSION; } 
int DynamicTRTResize::getNbOutputs() const TRT_NOEXCEPT { return 1; } 
size_t DynamicTRTResize::getSerializationSize() const TRT_NOEXCEPT { return serialized_size(mAlignCorners); 
} 
void DynamicTRTResize::serialize(void *buffer) const TRT_NOEXCEPT { serialize_value(&buffer, mAlignCorners); 
} 
////////////////////// creator ///////////////////////////// 
DynamicTRTResizeCreator::DynamicTRTResizeCreator() { mPluginAttributes.clear(); mPluginAttributes.emplace_back(nvinfer1::PluginField("align_corners")); mFC.nbFields = mPluginAttributes.size(); mFC.fields = mPluginAttributes.data(); 
} 
const char *DynamicTRTResizeCreator::getPluginName() const TRT_NOEXCEPT { return PLUGIN_NAME; } 
const char *DynamicTRTResizeCreator::getPluginVersion() const TRT_NOEXCEPT { return PLUGIN_VERSION; 
} 
nvinfer1::IPluginV2 *DynamicTRTResizeCreator::createPlugin( const char *name, const nvinfer1::PluginFieldCollection *fc) TRT_NOEXCEPT { nvinfer1::Dims size{2, {1, 1}}; bool align_corners = 1; for (int i = 0; i < fc->nbFields; i++) { if (fc->fields[i].data == nullptr) { continue; } std::string field_name(fc->fields[i].name); if (field_name.compare("align_corners") == 0) { align_corners = static_cast<const int *>(fc->fields[i].data)[0]; } } // create the instance of DynamicTRTResize DynamicTRTResize *plugin = new DynamicTRTResize(name, align_corners); plugin->setPluginNamespace(getPluginNamespace()); return plugin; 
} 
nvinfer1::IPluginV2 *DynamicTRTResizeCreator::deserializePlugin( const char *name, const void *serialData, size_t serialLength) TRT_NOEXCEPT { auto plugin = new DynamicTRTResize(name, serialData, serialLength); plugin->setPluginNamespace(getPluginNamespace()); return plugin; 
} 
REGISTER_TENSORRT_PLUGIN(DynamicTRTResizeCreator);//register the plugin 
}  // namespace mmdeploy 

??然后，我們就對 MMDeploy 重新 build 一次 TensorRT 的動態庫 build/lib/libmmdeploy_tensorrt_ops.so。一般編譯成功就表示已經注冊算子了，但是我們需要進行一些測試以保證結果正確。

4.測試

??我們用 TensorRT 的 python api 查看一下目前的插件列表：

import tensorrt as trt 
from mmdeploy.backend.tensorrt import load_tensorrt_plugin 
load_tensorrt_plugin() 
def get_plugin_names(): return [pc.name for pc in trt.get_plugin_registry().plugin_creator_list] 
print(get_plugin_names()) 

??可以發現 ‘DynamicTRTResize’ 在插件列表中。然后我們對這個插件進行功能測試，看推理結果是否和 PyTroch 結果一致，并且是否可以動態控制輸出尺寸。

from mmdeploy.backend.tensorrt import create_trt_engine, save_trt_engine 
engine = create_trt_engine( 'srcnn3.onnx', input_shapes=dict(input = dict( min_shape=[1, 3, 256, 256], opt_shape=[1, 3, 256, 256], max_shape=[1, 3, 256, 256]), factor = dict(min_shape = [1, 1, 256, 256], opt_shape = [1, 1, 512, 512], max_shape = [1, 1, 1024, 1024]))) 
save_trt_engine(engine, 'srcnn3.engine') 
from mmdeploy.backend.tensorrt import TRTWrapper 
trt_model = TRTWrapper('srcnn3.engine', ['output']) 
factor = torch.rand([1, 1, 768, 768], dtype=torch.float) 
trt_output = trt_model.forward(dict(input = x.cuda(), factor = factor.cuda())) 
torch_output = model.forward(x, factor) 
assert np.allclose(trt_output['output'].cpu().numpy(), torch_output.cpu().detach(), rtol = 1e-3, atol = 1e-5) 

??對比 TensorRT 的輸出結果和 PyTorch 的輸出結果是否一致，程序如果不報錯即可說明推理正確。此外，測試時我們使用和導出時不一樣的尺寸，結果也和 PyTorch 一致，說明可以支持動態的尺寸。