TensorRT部署流程

TensorRT 只需要知道网络的结构和参数即可，它支持三种转换入口：

1. TF-TRT，要求是 TensorFlow 模型
2. ONNX 模型格式
3. 使用 TensorRT API 手动把模型搭起来，然后把参数加载进去

第一种不够灵活，第三种比较麻烦，所以最省事方便的就是第二种方法。本文介绍第二种。

ONNX 就是一个通用的神经网络格式，一个 .onnx 文件内包含了网络的结构和参数。甭管是用什么深度学习框架写的网络，只要把模型导出成 ONNX 格式，就跟原本的代码没有关系了。

转成 ONNX 格式还没有被优化，需要再使用 TensorRT 读取它并优化成 TensorRT Engine。优化参数也在这一步指定。

对于第二个问题：得到的 TensorRT Engine 是硬件相关的，之后跑模型只需要运行这个 Engine 即可。调用 TensorRT Engine 需要使用 TensorRT Runtime API。

所以整个逻辑就是：

1. 把你的模型导出成 ONNX 格式。
2. 把 ONNX 格式模型输入给 TensorRT，并指定优化参数。
3. 使用 TensorRT 优化得到 TensorRT Engine。
4. 使用 TensorRT Engine 进行 inference。

你需要做的

1. 把模型导出成 ONNX 格式。
2. 安装 TensorRT 和 CUDA。注意二者和 driver 的版本号对应，我用的是 ZIP 安装：Installation Guide :: NVIDIA Deep Learning TensorRT Documentation，跟着这个把流程走一遍。
3. 设置优化参数，使用 TensorRT 把 ONNX 优化成 Engine，得到当前硬件上优化后的模型。
4. 使用 TensorRT Runtime API 进行 inference。

官方文档：Developer Guide :: NVIDIA Deep Learning TensorRT Documentation，写的很详细了。

ONNX 转换

Pytorch 自带导出方法 torch.onnx.export。

一定要注意 INPUT_NAME 和 OUTPUT_NAME 有没有写对，它决定了网络的入口和出口，转换 ONNX 错误/导出 Engine 错误很有可能是这个没指定对。

不确定的话可以用 onnxruntime 看一看对不对

import onnxruntime as rt
sess = rt.InferenceSession(onnx_model_path)
inputs_name = sess.get_inputs()[0].name
outputs_name = sess.get_outputs()[0].name
outputs = sess.run([outputs_name], {inputs_name: np.zeros(shape=(batch_size, xxx), dtype=np.float32)})

还可以在这个网站可视化导出的 ONNX：Netron

使用 trtexec.exe 测试

参考：TensorRT: TensorRT Command-Line Wrapper: trtexec (ccoderun.ca)

TensorRT 安装流程走完之后就能在 TensorRT-x-x-x-x/bin/ 文件夹下看到 trtexec.exe。

trtexec 是 TensorRT sample 里的一个例子，把 TensorRT 许多方法包装成了一个可执行文件。它可以把模型优化成 TensorRT Engine ，并且填入随机数跑 inference 进行速度测试。

这个命令：

./trtexec --onnx=model.onnx

把 onnx 模型优化成 Engine ，然后多次 inference 后统计并报时。

报时会报很多个， Enqueue Time 是 GPU 任务排队的时间，H2D Latency 是把网络输入数据从主存 host 拷贝进显存 device 的时间，D2H Latency 是把显存上的网络输出拷贝回主存的时间，只有 GPU Compute Time 是真正的网络 inference 时间。

当然可以把 Engine 文件导出，使用 --saveEngine 参数

./trtexec --onnx=model.onnx --saveEngine=xxx.trt

一般来说模型的第一维大小是任意的（batch size 维度），而 TensorRT 不能把任意 batch size 都加速。可以指定一个输入范围，并且重点优化其中一个 batch size。例如网络输入格式是 [-1, 3, 244, 244]：

./trtexec --onnx=model.onnx --minShapes=input:1x3x244x244 --optShapes=input:16x3x244x244 --maxShapes=input:32x3x244x244 --shapes=input:5x3x244x244

这个 input 对应于 INPUT_NAME:0。

还可以降低精度优化速度。一般来说大家写的模型内都是 float32 的运算，TensorRT 会默认开启 TF32 数据格式，它是截短版本的 FP32，只有 19 bit，保持了 fp16 的精度和 fp32 的指数范围。

另外，TensorRT 可以额外指定精度，把模型内的计算转换成 float16 或者 int8 的类型，可以只开一个也可以两个都开，trtexec 会倾向于速度最快的方式（有些网络模块不支持 int8）

./trtexec --onnx=model.onnx --saveEngine=xxx.trt --int8 --fp16

trtexec 还提供了 --best 参数，这相当于 --int8 --fp16 同时开。

一般来说，只开 fp16 可以把速度提一倍并且几乎不损失精度；但是开 --int8 会大大损失精度，速度会比 fp16 快，但不一定能快一倍。

int8 优化涉及模型量化，需要校准（calibrate）提升精度。TensorRT 有两种量化方法：训练后量化和训练中量化。二者的校准方法不同，精度也不同，后者更高一些。具体参考 Developer Guide :: NVIDIA Deep Learning TensorRT Documentation。

trtexec 采用的是训练后量化，写起来更方便一些。不过看看源码就能发现，因为 trtexec 只为了测试速度，所以校准就象征性做了一下，真想自己部署 int8 模型还得自己写校准。

使用 TensorRT C++ API

trtexec 只能看模型最快能跑多快，它是不管精度的，如果真想实际部署上又快又好的模型还是要自己调 TensorRT 的 API。

可以用 C++ API、Python API、TF-TRT Runtime，因为 TF-TRT 有局限性，C++ API 的速度比 Python API 快，所以我选择 C++ API。三者区别可以参考：TensorRT/5. Understanding TensorRT Runtimes.ipynb at main · NVIDIA/TensorRT (github.com)

参考 TensorRT 的 sample 自己写并不难。把 ONNX 转换成 TensorRT Engine 的代码是：

class Logger : public nvinfer1::ILogger {
	void log(Severity severity, const char* msg) noexcept override {
		// suppress info-level messages
		if (severity <= Severity::kWARNING)
			std::cout << msg << std::endl;
	}
} logger;

auto builder = std::unique_ptr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(logger));
// network definition
uint32_t explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); // indicate this network is explicit batch
auto network = std::unique_ptr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(explicitBatch));
// parser to parse the ONNX model
auto parser = std::unique_ptr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, logger));
// import ONNX model
parser->parseFromFile(onnx_filename.c_str(), static_cast<int32_t>(nvinfer1::ILogger::Severity::kWARNING));
// build engine
auto config = std::unique_ptr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig()); // optimization config
auto serializedModel = std::unique_ptr<nvinfer1::IHostMemory>(builder->buildSerializedNetwork(*network, *config));
// deserializing
auto runtime = std::unique_ptr<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(logger));
// load engine
engine = std::shared_ptr<nvinfer1::ICudaEngine>(runtime->deserializeCudaEngine(serializedModel->data(), serializedModel->size()));

大概逻辑就是，先声明一个 explicit batch 的网络（ONNX 是这样的），然后用 parser 从 ONNX 文件里读出来模型，根据 config 指定的优化参数进行 serialized，然后在 runtime 时把它 deserialize 得到的就是 Engine 了。

如果要加上 fp16 或者 int8 优化，需要在 serialized 之前，修改 config

auto profileOptimization = builder->createOptimizationProfile();
// We do not need to check the return of setDimension and setCalibrationProfile here as all dims are explicitly set
profileOptimization->setDimensions(network->getInput(0)->getName(), nvinfer1::OptProfileSelector::kMIN, nvinfer1::Dims2{ batch_size, xxx });
profileOptimization->setDimensions(network->getInput(0)->getName(), nvinfer1::OptProfileSelector::kOPT, nvinfer1::Dims2{ batch_size, xxx });
profileOptimization->setDimensions(network->getInput(0)->getName(), nvinfer1::OptProfileSelector::kMAX, nvinfer1::Dims2{ batch_size, xxx });
config->addOptimizationProfile(profileOptimization);
config->setCalibrationProfile(profileOptimization);

if (is_fp16) {
    if (builder->platformHasFastFp16()) {
        config->setFlag(nvinfer1::BuilderFlag::kFP16);
    }
    else {
        std::cout << "This platform does not support fp16" << std::endl;
    }
}

if (is_int8) {
    if (builder->platformHasFastInt8()) {
        config->setFlag(nvinfer1::BuilderFlag::kINT8);
        int batch_count = 4096; // calibrate size
        config->setInt8Calibrator(new NetworkInt8Calibrator(batch_count, { batch_size * xxx }, input_range, *network, std::cerr));
    }
    else {
        std::cout << "This platform does not support int8" << std::endl;
    }
}

这里我把输入的范围规定成唯一了。int8 这里的校准是我仿照 trtexec 写的，具体看下一节。

如果直接从文件中读取 Engine 就更简单了

std::ifstream engineFile(engine_filename, std::ios::binary);
engineFile.seekg(0, std::ifstream::end);
int64_t fsize = engineFile.tellg();
engineFile.seekg(0, std::ifstream::beg);
std::vector<uint8_t> engineBlob(fsize);
engineFile.read(reinterpret_cast<char*>(engineBlob.data()), fsize);
// deserializing a plan
auto runtime = std::unique_ptr<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(logger));
engine = std::shared_ptr<nvinfer1::ICudaEngine>(runtime->deserializeCudaEngine(engineBlob.data(), fsize));

把 Engine 保存到文件：

std::ofstream engineFile(engine_filename, std::ios::binary);
auto serializedEngine{ engine->serialize() };
engineFile.write(static_cast<char*>(serializedEngine->data()), serializedEngine->size());

使用时，需要告诉网络输入输出在显存上的 cuda 指针。

std::vector<float> input, output; // data on host
void* mInput_device_buffer;
void* mOutput_device_buffer;
// cudaMalloc(&xx_device_buffer, size)
cudaMemcpy(mInput_device_buffer, input.data(), input.size() * sizeof(float), cudaMemcpyHostToDevice)
std::vector<void*> bindings = { mInput_device_buffer, mOutput_device_buffer };
context->executeV2(bindings.data()); // inference
cudaMemcpy(output.data(), mOutput_device_buffer.data(), mOutput_device_buffer_size * sizeof(float)), cudaMemcpyDeviceToHost)

这是同步版的，当然还有异步版的 enqueueV2。

使用 TensorRT C++ API 实现 int8 校准

这里用的还是训练后校准。逻辑是：搞一些真实输入数据（不需要输出），告诉 TensorRT，它会根据真实输入数据的分布来调整量化的缩放幅度，以最大程度保证精度合适。

理论上校准数据越多，精度越高，但实际上不需要太多数据，TensorRT 官方说 500 张图像就足以校准 ImageNet 分类网络。

我的校准部分是

int batch_count = 4096; // calibrate size
config->setInt8Calibrator(new NetworkInt8Calibrator(batch_count, { batch_size * xxx }, input_range, *network, std::cerr));

就是搞了 4096 个 batch 的数据（这个数看着设就行），后面那个类是自己实现的，负责告诉校准器每个 batch 的数据是什么

class NetworkInt8Calibrator : public nvinfer1::IInt8EntropyCalibrator2 {
public:
	NetworkInt8Calibrator(int batches, const std::vector<int64_t>& elemCount, const std::vector<std::pair<float, float>>& data_range,
		const nvinfer1::INetworkDefinition& network, std::ostream& err) : mBatches(batches)
		, mCurrentBatch(0)
		, mErr(err)
		, data_range(data_range)
	{

		std::default_random_engine generator;
		std::uniform_real_distribution<float> distribution(0.0F, 1.0F);
		auto gen = [&generator, &distribution]() { return distribution(generator); };

		for (int i = 0; i < network.getNbInputs(); i++)
		{
			auto* input = network.getInput(i);
			std::vector<float> rnd_data(elemCount[i]);
			//std::generate_n(rnd_data.begin(), elemCount[i], gen);
			std::vector<void*> data(mBatches);
			for (int c = 0; c < mBatches; c++) {
				// use `gen` generate `elemCount[i]` data to `rnd_data`
				cudaCheck(cudaMalloc(&data[c], elemCount[i] * sizeof(float)), mErr);
				cudaCheck(cudaMemcpy(data[c], rnd_data.data(), elemCount[i] * sizeof(float), cudaMemcpyHostToDevice), mErr);
			}
			mInputDeviceBuffers.insert(std::make_pair(input->getName(), data));
		}
	}


	~NetworkInt8Calibrator()
	{
		for (auto& elem : mInputDeviceBuffers)
		{
			for(int i = 0;i < mBatches;i ++)
				cudaCheck(cudaFree(elem.second[i]), mErr);
		}
	}

	bool getBatch(void* bindings[], const char* names[], int nbBindings) noexcept {
		if (mCurrentBatch >= mBatches)
		{
			return false;
		}

		for (int i = 0; i < nbBindings; ++i)
		{
			bindings[i] = mInputDeviceBuffers[names[i]][mCurrentBatch];
		}

		++mCurrentBatch;

		return true;
	}

	int getBatchSize() const noexcept {
		return mBatches;
	}

	const void* readCalibrationCache(size_t& length) noexcept { return nullptr; }

	void writeCalibrationCache(const void*, size_t) noexcept {}

private:
	int mBatches{};
	int mCurrentBatch{};
	std::vector<std::pair<float, float>> data_range;
	std::map<std::string, std::vector<void*>> mInputDeviceBuffers;
	std::vector<char> mCalibrationCache;
	std::ostream& mErr;
};

继承的类是 IInt8EntropyCalibrator2 ，这个得根据需要选择，不同类型的网络不一样，详见 Developer Guide :: NVIDIA Deep Learning TensorRT Documentation。

这里面有一个 read 和 write 函数我没实现，它们负责从文件中读和写 calibration cache 。如果用这个 cache 那么省去了生成数据和网络跑数据的时间，生成 Engine 时会更快一些。

来自于作者：https://zhuanlan.zhihu.com/p/527238167

感谢作者。