在前面的模型入门系列文章中，我们介绍了部署一个 PyTorch 模型到推理后端，如 ONNXRuntime，这其中可能遇到很多工程性的问题。
有些可以通过创建 ONNX 节点来解决，该节点仍然使用后端原生的实现进行推理。而有些无法导出到后端的算法，可以通过重写代码改变算法的实现过程，同样可以导出到 ONNX ，达到一致的效果。以上两种方式一般可以处理绝大多数的部署问题，同时也不需要向推理框架引入新的内容，是我们进行模型部署时候的优先选择。
然而，仍然存在部分模型，模型中某些算子无法通过上述两种方式绕过问题，这时候，如何对特定后端实现对应代码就极为重要。这也是本文将介绍的第三种方式——自定义插件。
自定义插件是很多推理框架支持用户自定义算子的方式，以 MMDeploy 为例，它是一个支持多种推理后端的算法库。目前支持的后端有：
- ONNXRuntime
- TensorRT
- ncnn
- OpenVINO
- PPLNN
其中，前三种后端均实现了一些自定义的算子。例如 ONNXRuntime 中的调制可变性卷积，ncnn 中的topk 算子，TensorRT 中的 MultiLevelRoiAlign 。
介绍如何给后端自定义算子是一件相对复杂的事情，所以本文只针对其中一种后端 TensorRT，介绍自定义算子。如果读者对其他后端感兴趣，可以去他们的代码库查看，一般地，各个推理框架均有详细文档介绍如何添加客制化的算子实现。
在 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 导致的。
创建 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++ 代码，实现该插件。
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
</vector></string></memory></cublas_v2.h>
在这样一份头文件中，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
</const></float></chrono></assert.h>
然后，我们就对 MMDeploy 重新 build 一次 TensorRT 的动态库 build/lib/libmmdeploy_tensorrt_ops.so。一般编译成功就表示已经注册算子了，但是我们需要进行一些测试以保证结果正确。
测试
我们用 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 一致，说明可以支持动态的尺寸。
总结
本篇教程我们主要讲述如何在 MMDeploy 代码库中添加一个自定义的 TensorRT 插件，整个过程不涉及太多更复杂的 CUDA 编程，相信小伙伴们学完可以自己实现想要的插件。
https://github.com/open-mmlab/mmdeploygithub.com/open-mmlab/mmdeploy
至此，我们的模型部署入门系列教程已经更新了八期，那到这里可能就先暂时告一段落啦！感谢小伙伴的支持和喜爱，后续我们也将继续推出模型部署进阶教程，请大家敬请期待哦！所有跟模型部署相关的内容，都整理在这个专栏啦，欢迎大家关注~
模型部署那些事www.zhihu.com/column/c_1497987564452114432正在上传…重新上传取消
Original: https://blog.csdn.net/qq_39967751/article/details/126429831
Author: OpenMMLab
Title: 模型部署入门教程（八）：如何添加TensorRT自定义算子
原创文章受到原创版权保护。转载请注明出处：https://www.johngo689.com/711183/
转载文章受原作者版权保护。转载请注明原作者出处！