对于tensorrt 不支持的算子，可以通过自定义算子，通过plugin 的方式实现。对于pytorch, 将自定义算子导出到onnx 中，然后通过tensorrt 的plugin解析。

一、在pytorch 中自定义onnx 算子

1、自定义算子类继承 torch.autograd.Function 类，实现forward()和backward()方法，这样就变成一个可导函数可以在pytorch 模型网络中调用。如果只做推理可以不用实现backward。
2、实现symbolic 静态方法，这样在调用torch.onnx.export()时就能根据symbolic定义的规则，将自定义算子类转换成onnx 算子。
说明：symbolic是符号函数，通常在其内部**返回一个g.op()**对象。g.op() 把一个 PyTorch 算子映射成一个或多个自带的 ONNX 算子，或者是自定义的 ONNX 算子。

二、模型搭建

自定义算子调用类，继承nn.Module，调用自定义算子，这样就可以加入到模型网络中。最后在模型网络中调用自定义算子调用类即可。

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.onnx
import torch.autograd
import os
class MYSELUImpl(torch.autograd.Function):
    @staticmethod
    def symbolic(g, x, p):
        print("==================================call symbolic")
        return g.op("MYSELU", x, p,
                    g.op("Constant", value_t=torch.tensor([3, 2, 1], dtype=torch.float32)),
                    attr1_s="这是字符串属性",
                    attr2_i=[1, 2, 3],
                    attr3_f=222
                    )
    @staticmethod
    def forward(ctx, x, p):
        return 1.0 / (1.0 + torch.exp(-x)) * x

class MYSELU(nn.Module):
    def __init__(self, n):
        super().__init__()
        self.param = nn.parameter.Parameter(torch.arange(n).float())
    def forward(self, x):
        return MYSELUImpl.apply(x, self.param)

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv = nn.Conv2d(1, 1, 3, padding=1)
        self.myselu = MYSELU(3)
        #self.myselu = nn.Sigmoid()
        self.conv.weight.data.fill_(1)
        self.conv.bias.data.fill_(0)
    def forward(self, x):
        x = self.conv(x)
        x = self.myselu(x)
        return x

model = Model().eval()
input = torch.tensor([
    # batch 0
    [
        [1, 1, 1],
        [1, 1, 1],
        [1, 1, 1],
    ],
], dtype=torch.float32).view(1, 1, 3, 3)
print(input.shape)
output = model(input)
print(f"inference output = \n{output}")
# dummy = torch.zeros(1, 1, 3, 3)
# output = model(dummy)
print(f"inference output = \n{output}")
torch.onnx.export(
    model, # 这里的args，是指输入给model的参数，需要传递tuple，因此用括号
    (input,),
    "myselu.onnx", # 储存的文件路径
    verbose=True,# 打印详细信息
    input_names=["image"], # 为输入和输出节点指定名称，方便后面查看或者操作
    output_names=["output"],
    opset_version=11,# 这里的opset，指，各类算子以何种方式导出，对应于symbolic_opset11
    # 表示他有batch、height、width3个维度是动态的，在onnx中给其赋值为-1,通常，我们只设置batch为动态，其他的避免动态
    dynamic_axes={
        "image": {0: "batch", 2: "height", 3: "width"},
        "output": {0: "batch", 2: "height", 3: "width"},
    },
    # dynamic_axes={
    #     "image": {0: "batch"},
    #     "output": {0: "batch"},
    # },
    # 对于插件，需要禁用onnx检查
    # enable_onnx_checker=False
    operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK
)
print("Done.!")

以上就可以导出 onnx 模型，动态推理可自行设置。

三、onnx 转 tensorrt

onnx 转tensorrt 时需要tensorrt 识别onnx 里所有的算子，对于tensort 已经发布的自带算子，可以直接解析， 对于tensorrt 没有定义的算子，需要通过继承的方式重写自定义算子类，使得解析时能找到对应的解析方法。

1、插件具体实现

继承 nvinfer1::IPluginV2DynamicExt类
插件的具体实现逻辑一般用cuda 核函数重写，在enqueue()函数中调用核函数。
IPluginV2DynamicExt继承自IPluginV2Ext, IPluginV2Ext又继承自IPluginV2
所以需要实现的overwrite 函数有：

IPluginV2DynamicExt基类：
    构造函数和析构函数
    virtual DimsExprs getOutputDimensions(）：输出数据的尺寸
    virtual bool supportsFormatCombination(）：支持的数据类型，int8，float16，float32等
    virtual void configurePlugin（）： 配置插件格式（这个算子所采用的数据格式和类型）
    virtual size_t getWorkspaceSize()： 需要的额外空间大小
    virtual int enqueue()： 推理具体逻辑
IPluginV2Ext基类：
    virtual nvinfer1::DataType getOutputDataType()
IPluginV2基类：
    virtual AsciiChar const* getPluginType()
    virtual AsciiChar const* getPluginVersion()
    virtual int32_t getNbOutputs()
    virtual size_t getSerializationSize()
    virtual void serialize(void* buffer)

2、插件实例创建

继承nvinfer1::IPluginCreator类，用来调用插件具体实现。
IPluginCreator基类，主要需要实现的虚函数如下：

构造函数和析构函数
virtual AsciiChar const* getPluginName()
virtual AsciiChar const* getPluginVersion()
virtual PluginFieldCollection const* getFieldNames()
virtual IPluginV2* createPlugin()
virtual IPluginV2* deserializePlugin()
virtual void setPluginNamespace()
virtual AsciiChar const* getPluginNamespace()

3、插件注册

采用宏REGISTER_TENSORRT_PLUGIN注册插件

4、编译推理过程

编译阶段
        通过MySELUPluginCreator::createPlugin创建plugin
        期间会调用MySELUPlugin::clone克隆插件
        调用MySELUPlugin::supportsFormatCombination判断该插件所支持的数据格式和类型
        在这里我们告诉引擎，本插件可以支持什么类型的推理
        可以支持多种，例如fp32、fp16、int8等等
        调用MySELUPlugin::getOutputDimensions获取该层的输出维度是多少
        调用MySELUPlugin::enqueue进行性能测试（不是一定会执行）
        如果支持多种，则会在多种里面进行实际测试，选择一个性能最好的配置
        调用MySELUPlugin::configurePlugin配置插件格式
        告诉你目前这个层所采用的数据格式和类型
        调用MySELUPlugin::serialize将该层的参数序列化储存为trtmodel文件

推理阶段
        通过MySELUPluginCreator::deserializePlugin反序列化插件参数进行创建
        期间会调用MySELUPlugin::clone克隆插件
        调用MySELUPlugin::configurePlugin配置当前插件使用的数据类型和格式
        调用MySELUPlugin::enqueue进行推理

#ifndef CUSTOM_MYSELU_PLUGIN_H
#define CUSTOM_MYSELU_PLUGIN_H

#include <NvInferPlugin.h>
#include <string>
#include <vector>

class MySELUPlugin : public nvinfer1::IPluginV2DynamicExt {
public:
	MySELUPlugin(const std::string name, const std::string attr1, float attr3);  // 接受算子名称属性，build engine时构造函数
	MySELUPlugin(const std::string name, const void* data, size_t length);  // 接受算子名称和反序列化的engine data，推理时构造函数
	MySELUPlugin() = delete;

	int getNbOutputs() const noexcept override;
	virtual nvinfer1::DataType getOutputDataType(int32_t index,
		nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept override {
		return inputTypes[0];
	}
	virtual nvinfer1::DimsExprs getOutputDimensions(int32_t outputIndex,
		const nvinfer1::DimsExprs* inputs, int32_t nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept override;

	int initialize() noexcept override;
	void terminate() noexcept override;

	virtual size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
		int32_t nbInputs, const nvinfer1::PluginTensorDesc* outputs,
		int32_t nbOutputs) const noexcept override {
		return 0;
	};

	int enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc,
		const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override;

	size_t getSerializationSize() const noexcept override;
	void serialize(void* buffer)  const noexcept override;

	virtual void configurePlugin(const  nvinfer1::DynamicPluginTensorDesc* in, int32_t nbInputs,
		const  nvinfer1::DynamicPluginTensorDesc* out, int32_t nbOutputs) noexcept override;

	virtual bool supportsFormatCombination(int32_t pos, const nvinfer1::PluginTensorDesc* inOut, int32_t nbInputs,
		int32_t nbOutputs) noexcept override;

	const char* getPluginType() const noexcept override;
	const char* getPluginVersion() const noexcept override;

	void destroy() noexcept override;
	nvinfer1::IPluginV2DynamicExt* clone() const noexcept override;
	void setPluginNamespace(nvinfer1::AsciiChar const* pluginNamespace) noexcept override;
	const char* getPluginNamespace()const noexcept override;

private:
	const std::string mLayerName;
	std::string mattr1;
	float mattr3;
	size_t mInputVolume;
	std::string mNamespace;
};


class MySELUPluginCreator : public nvinfer1::IPluginCreator {
public:
	MySELUPluginCreator();
	const char* getPluginName() const noexcept override;
	const char* getPluginVersion() const noexcept override;
	const nvinfer1::PluginFieldCollection* getFieldNames() noexcept override;

	nvinfer1::IPluginV2* createPlugin(nvinfer1::AsciiChar const* name,
		nvinfer1::PluginFieldCollection const* fc) noexcept override;
	nvinfer1::IPluginV2* deserializePlugin(nvinfer1::AsciiChar const* name,
		void const* serialData, size_t serialLength)noexcept override;
	void setPluginNamespace(nvinfer1::AsciiChar const* pluginNamespace) noexcept override;
	const char* getPluginNamespace() const noexcept override;
private:
	static nvinfer1::PluginFieldCollection mfc;
	static std::vector<nvinfer1::PluginField> mPluginAttributes;
	std::string mNamespace;
};

#endif

#include "myselu_plugin.hpp"
#include <NvInfer.h>
#include <cstring>
#include <vector>
#include <cassert>

void myselu_inference(const float* x, float* output, int n, cudaStream_t stream);
// MySELU plugin的特定常量
namespace {
	const char* MYSELU_PLUGIN_VERSION{ "1" };
	const char* MYSELU_PLUGIN_NAME{ "MYSELU" };  //名称要和onnx中对应的一致
}
// 静态类字段的初始化
nvinfer1::PluginFieldCollection MySELUPluginCreator::mfc{};
std::vector<nvinfer1::PluginField> MySELUPluginCreator::mPluginAttributes;
// 实际注册时，注册的是创建器，交给tensorRT管理
REGISTER_TENSORRT_PLUGIN(MySELUPluginCreator);
// 用于序列化插件的Helper function
template <typename T>
void writeToBuffer(char*& buffer, const T& val) {
	*reinterpret_cast<T*>(buffer) = val;
	buffer += sizeof(T);
}
// 用于反序列化插件的Helper function
template <typename T>
T readFromBuffer(const char*& buffer) {
	T val = *reinterpret_cast<const T*>(buffer);
	buffer += sizeof(T);
	return val;
}
// 定义插件类MYSELUPlugin
MySELUPlugin::MySELUPlugin(const std::string name, const std::string attr1, float attr3)
	:mLayerName(name), mattr1(attr1), mattr3(attr3)
{
	printf("==================== 编译阶段，attr1 = %s, attr3 = %f\n", attr1.c_str(), attr3);
};

MySELUPlugin::MySELUPlugin(const std::string name, const void* data, size_t length) 
	:mLayerName(name)
{
	// Deserialize in the same order as serialization
	const char* d = static_cast<const char*>(data);
	const char* a = d;
	int nstr = readFromBuffer<int>(d);
	mattr1 = std::string(d, d + nstr);
	d += nstr;
	mattr3 = readFromBuffer<float>(d);
	assert(d == (a + length));
	printf("==================== 推理阶段，attr1 = %s, attr3 = %f\n", mattr1.c_str(), mattr3);
};

const char* MySELUPlugin::getPluginType() const noexcept
{
	return MYSELU_PLUGIN_NAME;
}
const char* MySELUPlugin::getPluginVersion() const noexcept
{
	return MYSELU_PLUGIN_VERSION;
}
int MySELUPlugin::getNbOutputs() const noexcept {
	return 1;
}
// 获取该层的输出维度是多少
nvinfer1::DimsExprs MySELUPlugin::getOutputDimensions(int32_t outputIndex,
	const nvinfer1::DimsExprs* inputs, int32_t nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept 
{
	// MySELUping不改变输入尺寸，所以输出尺寸将与输入尺寸相同
	return *inputs;
}
int MySELUPlugin::initialize() noexcept
{
	return 0;
}
int MySELUPlugin::enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc,
	const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
{
	void* output = outputs[0];
	size_t volume = 1;
	for (int i = 0; i < inputDesc->dims.nbDims; ++i) {
		volume *= inputDesc->dims.d[i];
	}
	mInputVolume = volume;
	myselu_inference(static_cast<const float*>(inputs[0]),
		static_cast<float*>(output),
		mInputVolume,
		stream
	);
	return 0;
}
size_t MySELUPlugin::getSerializationSize() const noexcept
{
	return sizeof(int) + mattr1.size() + sizeof(mattr3);
}
// 该层的参数序列化储存为trtmodel文件
void MySELUPlugin::serialize(void* buffer)  const noexcept
{
	char* d = static_cast<char*>(buffer);
	const char* a = d;
	int nstr = mattr1.size();
	writeToBuffer(d, nstr);
	memcpy(d, mattr1.data(), nstr);
	d += nstr;
	writeToBuffer(d, mattr3);
	assert(d == a + getSerializationSize());
}
// 判断该插件所支持的数据格式和类型
bool MySELUPlugin::supportsFormatCombination(int32_t pos, const nvinfer1::PluginTensorDesc* inOut, int32_t nbInputs,
	int32_t nbOutputs) noexcept
{
	auto type = inOut[pos].type;
	auto format = inOut[pos].format;
	// 这个插件只支持普通的浮点数，以及NCHW输入格式
	if (type == nvinfer1::DataType::kFLOAT && format == nvinfer1::PluginFormat::kLINEAR) {
		return true;
	}
	else {
		return false;
	}
}
void MySELUPlugin::terminate() noexcept {}
void MySELUPlugin::destroy() noexcept
{
	// This gets called when the network containing plugin is destroyed
	delete this;
}
// 配置插件格式:目前这个层所采用的数据格式和类型
void MySELUPlugin::configurePlugin(const  nvinfer1::DynamicPluginTensorDesc* in, int32_t nbInputs,
	const  nvinfer1::DynamicPluginTensorDesc* out, int32_t nbOutputs) noexcept
{
	auto type = in->desc.type;
	auto format = in->desc.format;
	assert(nbOutputs == 1);
	assert(type == nvinfer1::DataType::kFLOAT);
	assert(format == nvinfer1::PluginFormat::kLINEAR);
}

// 克隆插件
nvinfer1::IPluginV2DynamicExt* MySELUPlugin::clone() const noexcept
{
	printf("===================克隆插件=================\n");
	auto plugin = new MySELUPlugin(mLayerName, mattr1, mattr3);
	plugin->setPluginNamespace(mNamespace.c_str());
	return plugin;
}
void MySELUPlugin::setPluginNamespace(const char* libNamespace) noexcept
{
	mNamespace = libNamespace;
}
const char* MySELUPlugin::getPluginNamespace() const noexcept
{
	return mNamespace.c_str();
}
// 插件创建器
MySELUPluginCreator::MySELUPluginCreator()
{
	// 描述MySELUPlugin的必要PluginField参数
	mPluginAttributes.emplace_back(nvinfer1::PluginField("attr1", nullptr, nvinfer1::PluginFieldType::kCHAR));
	mPluginAttributes.emplace_back(nvinfer1::PluginField("attr3", nullptr, nvinfer1::PluginFieldType::kFLOAT32));
	// 收集PluginField的参数
	mfc.nbFields = mPluginAttributes.size();
	mfc.fields = mPluginAttributes.data();
}
const char* MySELUPluginCreator::getPluginName() const noexcept
{
	return MYSELU_PLUGIN_NAME;
}
const char* MySELUPluginCreator::getPluginVersion() const noexcept
{
	return MYSELU_PLUGIN_VERSION;
}
const nvinfer1::PluginFieldCollection* MySELUPluginCreator::getFieldNames() noexcept
{
	return &mfc;
}
// 创建plugin
nvinfer1::IPluginV2* MySELUPluginCreator::createPlugin(nvinfer1::AsciiChar const* name,
	nvinfer1::PluginFieldCollection const* fc) noexcept
{
	std::string attr1;
	float attr3;
	const nvinfer1::PluginField* fields = fc->fields;

	// Parse fields from PluginFieldCollection
	for (int i = 0; i < fc->nbFields; ++i) {
		if (strcmp(fields[i].name, "attr1")==0) {
			assert(fields[i].type == nvinfer1::PluginFieldType::kCHAR);
			auto cp = static_cast<const char*>(fields[i].data);
			attr1 = std::string(cp, cp + fields[i].length);
		}
		else if (strcmp(fields[i].name, "attr3") == 0) {
			assert(fields[i].type == nvinfer1::PluginFieldType::kFLOAT32);
			attr3 = *(static_cast<const float*>(fields[i].data));
		}
	}
	return new MySELUPlugin(name, attr1, attr3);
}
// 反序列化插件参数进行创建
nvinfer1::IPluginV2* MySELUPluginCreator::deserializePlugin(nvinfer1::AsciiChar const* name,
	void const* serialData, size_t serialLength)noexcept
{
	// This object will be deleted when the network is destroyed, which will
	// call MySELUPlugin::destroy()
	return new MySELUPlugin(name, serialData, serialLength);
}

void MySELUPluginCreator::setPluginNamespace(const char* libNamespace) noexcept
{
	mNamespace = libNamespace;
}
const char* MySELUPluginCreator::getPluginNamespace() const noexcept
{
	return mNamespace.c_str();
}

核函数myselu_kernel.cu

#include <cuda_runtime.h>
#include <cmath>

static __device__ float sigmoid(float x) {
	return 1 / (1 + expf(-x));
}

static __global__ void myselu_kernel(const float* x, float* output, int n)
{
	int position = threadIdx.x + blockDim.x*blockIdx.x;
	if (position >= n) return;
	output[position] = x[position]*sigmoid(x[position]);
}

void myselu_inference(const float* x, float* output, int n, cudaStream_t stream)
{
	const int nthreads = 512;
	int block_size = n > nthreads ? nthreads : n;
	int grid_size = (n + block_size - 1) / block_size;
	myselu_kernel<<<grid_size, block_size, 0, stream>>>(x, output, n);
}

主函数main.cpp

// tensorRT include
// 编译用的头文件
#include <NvInfer.h>
// onnx解析器的头文件
#include <NvOnnxParser.h>
// 推理用的运行时头文件
#include <NvInferRuntime.h>
// cuda include
#include <cuda_runtime.h>
// system include
#include <stdio.h>
#include <math.h>

#include <iostream>
#include <fstream>
#include <vector>

using namespace std;

inline const char* severity_string(nvinfer1::ILogger::Severity t) {
	switch (t) {
	case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: return "internal_error";
	case nvinfer1::ILogger::Severity::kERROR:   return "error";
	case nvinfer1::ILogger::Severity::kWARNING: return "warning";
	case nvinfer1::ILogger::Severity::kINFO:    return "info";
	case nvinfer1::ILogger::Severity::kVERBOSE: return "verbose";
	default: return "unknow";
	}
}

class TRTLogger : public nvinfer1::ILogger {
public:
	virtual void log(Severity severity, nvinfer1::AsciiChar const* msg) noexcept override {
		if (severity <= Severity::kINFO) {
			if (severity == Severity::kWARNING) {
				printf("\033[33m%s: %s\033[0m\n", severity_string(severity), msg);
			}
			else if (severity <= Severity::kERROR) {
				printf("\033[31m%s: %s\033[0m\n", severity_string(severity), msg);
			}
			else {
				printf("%s: %s\n", severity_string(severity), msg);
			}
		}
	}
} logger;

bool build_model() {
	TRTLogger logger;

	// 这是基本需要的组件
	nvinfer1::IBuilder* builder = nvinfer1::createInferBuilder(logger);
	nvinfer1::IBuilderConfig* config = builder->createBuilderConfig();
	nvinfer1::INetworkDefinition* network = builder->createNetworkV2(1);

	// 通过onnxparser解析器解析的结果会填充到network中，类似addConv的方式添加进去
	nvonnxparser::IParser* parser = nvonnxparser::createParser(*network, logger);
	if (!parser->parseFromFile("myselu.onnx", 1)) {
		network->destroy();
		config->destroy();
		builder->destroy();
		printf("load onnx file failed\n");
		return false;
	}

	int maxBatchSize = 10;
	printf("Workspace Size = %.2f MB\n", (1 << 28) / 1024.0f / 1024.0f);
	config->setMaxWorkspaceSize(1 << 28);

	// 如果模型有多个输入，则必须多个profile
	auto profile = builder->createOptimizationProfile();
	auto input_tensor = network->getInput(0);
	int input_channel = input_tensor->getDimensions().d[1];

	// 配置输入的最小、最优、最大的范围
	profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims4(1, input_channel, 3, 3));
	profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kOPT, nvinfer1::Dims4(1, input_channel, 3, 3));
	profile->setDimensions(input_tensor->getName(), nvinfer1::OptProfileSelector::kMAX, nvinfer1::Dims4(maxBatchSize, input_channel, 5, 5));
	config->addOptimizationProfile(profile);

	nvinfer1::ICudaEngine* engine = builder->buildEngineWithConfig(*network, *config);
	if (engine == nullptr) {
		printf("build engine failed\n");
		network->destroy();
		config->destroy();
		builder->destroy();
		return false;
	}

	// 将模型序列化，并储存为文件
	nvinfer1::IHostMemory* model_data = engine->serialize();
	FILE* f = fopen("myselu.trtmodel", "wb");
	fwrite(model_data->data(), 1, model_data->size(), f);
	fclose(f);

	// 卸载顺序按照构建顺序倒序
	model_data->destroy();
	parser->destroy();
	engine->destroy();
	network->destroy();
	config->destroy();
	builder->destroy();
	printf("Done.\n");
	return true;
}

///////////////////////////////////////////////////////////////////////////////////////////////////////

vector<unsigned char> load_file(const string& file) {
	ifstream in(file, ios::in | ios::binary);
	if (!in.is_open())
		return {};

	in.seekg(0, ios::end);
	size_t length = in.tellg();

	std::vector<uint8_t> data;
	if (length > 0) {
		in.seekg(0, ios::beg);
		data.resize(length);

		in.read((char*)&data[0], length);
	}
	in.close();
	return data;
}

void inference() {

	TRTLogger logger;
	auto engine_data = load_file("myselu.trtmodel");
	nvinfer1::IRuntime* runtime = nvinfer1::createInferRuntime(logger);
	nvinfer1::ICudaEngine* engine = runtime->deserializeCudaEngine(engine_data.data(), engine_data.size());
	if (engine == nullptr) {
		printf("Deserialize cuda engine failed.\n");
		runtime->destroy();
		return;
	}

	nvinfer1::IExecutionContext* execution_context = engine->createExecutionContext();
	cudaStream_t stream = nullptr;
	cudaStreamCreate(&stream);

	float input_data_host[] = {
		// batch 0
		1,   1,   1,
		1,   1,   1,
		1,   1,   1,

		// batch 1
		-1,   1,   1,
		1,   0,   1,
		1,   1,   -1
	};
	float* input_data_device = nullptr;

	// 3x3输入，对应3x3输出
	const int ib = 2;
	const int iw = 3;
	const int ih = 3;
	float output_data_host[ib * iw * ih];
	float* output_data_device = nullptr;
	cudaMalloc(&input_data_device, sizeof(input_data_host));
	cudaMalloc(&output_data_device, sizeof(output_data_host));
	cudaMemcpyAsync(input_data_device, input_data_host, sizeof(input_data_host), cudaMemcpyHostToDevice, stream);

	// 明确当前推理时，使用的数据输入大小
	execution_context->setBindingDimensions(0, nvinfer1::Dims4(ib, 1, ih, iw));
	float* bindings[] = { input_data_device, output_data_device };
	bool success = execution_context->enqueueV2((void**)bindings, stream, nullptr);
	cudaMemcpyAsync(output_data_host, output_data_device, sizeof(output_data_host), cudaMemcpyDeviceToHost, stream);
	cudaStreamSynchronize(stream);

	for (int b = 0; b < ib; ++b) {
		printf("batch %d. output_data_host = \n", b);
		for (int i = 0; i < iw * ih; ++i) {
			printf("%f, ", output_data_host[b * iw * ih + i]);
			if ((i + 1) % iw == 0)
				printf("\n");
		}
	}

	printf("Clean memory\n");
	cudaStreamDestroy(stream);
	cudaFree(input_data_device);
	cudaFree(output_data_device);
	execution_context->destroy();
	engine->destroy();
	runtime->destroy();
}

int main() {
	if (!build_model()) {
		return -1;
	}
	inference();
	return 0;
}