以下文章来源于：启思@知乎
作者：启思
原文链接：https://zhuanlan.zhihu.com/p/527238167

本文仅用于学术分享，如有侵权，请联系后台作删文处理
导读
模型部署作为人工智能的落地的“最后一步”，也是算法能够转换为生产力的重要环节。本文作者分享了TensorRT 的部署流程，希望能对各位读者有所帮助。
前段时间用 TensorRT 部署了一套模型，速度相比 Python 实现的版本快了 20 多倍，中间踩了许多坑，但是最后发现流程其实相当简单，特此记录一下踩坑过程。
01
TensorRT
这东西就是 NVidia 在自家显卡上做了一个深度学习 inference 加速的框架，只要你把训练好的模型参数和结构告诉他，他就能自动帮你优化（硬件相关），以达到最快速度。
这涉及两个问题：
  1. 应该以什么模型格式把模型喂给 TensorRT？
  2. 如何使用 TensorRT 优化后的模型？
对于第一个问题：现在的深度学习框架非常多，不止常用的 pytorch/tensorflow，而即使是同一种框架还可以使用不同的编程语言实现。让 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。
02
你需要做的
  1. 把模型导出成 ONNX 格式。
  2. 安装 TensorRT 和 CUDA。注意二者和 driver 的版本号对应，我用的是 ZIP 安装，跟着这个把流程走一遍。
    https://docs.nvidia.com/deeplearning/tensorrt/install-guide/index.html#installing-zip
  3. 设置优化参数，使用 TensorRT 把 ONNX 优化成 Engine，得到当前硬件上优化后的模型。
  4. 使用 TensorRT Runtime API 进行 inference。
官方文档，写的很详细了。https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/#overview
03
ONNX 转换
Pytorch 自带导出方法 torch.onnx.export。
TensorFlow 推荐使用这个 https://github.com/onnx/tensorflow-onnx，从 checkpoint （.meta 后缀）导出只需要
python-mtf2onnx.convert--checkpointXXX.meta--inputsINPUT_NAME:0--outputsOUTPUT_NAME:0--outputXXX.onnx
即可，tf 1.x/2.x 均可以用。
这里一定要注意 INPUT_NAME 和 OUTPUT_NAME 有没有写对，它决定了网络的入口和出口，转换 ONNX 错误/导出 Engine 错误很有可能是这个没指定对。
不确定的话可以用 onnxruntime 看一看对不对
import onnxruntime as rtsess = rt.InferenceSession(onnx_model_path)inputs_name = sess.get_inputs()[0].nameoutputs_name = sess.get_outputs()[0].nameoutputs = sess.run([outputs_name], {inputs_name: np.zeros(shape=(batch_size, xxx), dtype=np.float32)})
还可以在这个网站可视化导出的 ONNX：https://netron.app/
04
使用 trtexec.exe 测试
参考：
https://www.ccoderun.ca/programming/doxygen/tensorrt/md_TensorRT_samples_opensource_trtexec_README.html
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 有两种量化方法：训练后量化和训练中量化。二者的校准方法不同，精度也不同，后者更高一些。具体参考 
https://docs.nvidia.com/deeplearning/tensorrt/developerguide/index.html#working-with-int8
trtexec 采用的是训练后量化，写起来更方便一些。不过看看源码就能发现，因为 trtexec 只为了测试速度，所以校准就象征性做了一下，真想自己部署 int8 模型还得自己写校准。
05
使用 TensorRT C++ API
trtexec 只能看模型最快能跑多快，它是不管精度的，如果真想实际部署上又快又好的模型还是要自己调 TensorRT 的 API。
可以用 C++ API、Python API、TF-TRT Runtime，因为 TF-TRT 有局限性，C++ API 的速度比 Python API 快，所以我选择 C++ API。三者区别可以参考：
https://github.com/NVIDIA/TensorRT/blob/main/quickstart/IntroNotebooks/5.%20Understanding%20TensorRT%20Runtimes.ipynb
参考 TensorRT 的 sample 自己写并不难。把 ONNX 转换成 TensorRT Engine 的代码是：
classLogger :public nvinfer1::ILogger {voidlog(Severity severity, constchar* msg)noexcept override {// suppress info-level messagesif (severity <= Severity::kWARNING)std::cout << msg << std::endl; }} logger;auto builder = std::unique_ptr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(logger));// network definitionuint32_t explicitBatch = 1U << static_cast<uint32_t>(nvinfer1::NetworkDefinitionCreationFlag::kEXPLICIT_BATCH); // indicate this network is explicit batchauto network = std::unique_ptr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(explicitBatch));// parser to parse the ONNX modelauto parser = std::unique_ptr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, logger));// import ONNX modelparser->parseFromFile(onnx_filename.c_str(), static_cast<int32_t>(nvinfer1::ILogger::Severity::kWARNING));// build engineauto config = std::unique_ptr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig()); // optimization configauto serializedModel = std::unique_ptr<nvinfer1::IHostMemory>(builder->buildSerializedNetwork(*network, *config));// deserializingauto runtime = std::unique_ptr<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(logger));// load engineengine = 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 setprofileOptimization->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 planauto 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 hostvoid* 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()); // inferencecudaMemcpy(output.data(), mOutput_device_buffer.data(), mOutput_device_buffer_size * sizeof(float)), cudaMemcpyDeviceToHost)
这是同步版的，当然还有异步版的 enqueueV2。
06
使用 TensorRT C++ API 实现 int8 校准
这里用的还是训练后校准。逻辑是：搞一些真实输入数据（不需要输出），告诉 TensorRT，它会根据真实输入数据的分布来调整量化的缩放幅度，以最大程度保证精度合适。
理论上校准数据越多，精度越高，但实际上不需要太多数据，TensorRT 官方说 500 张图像就足以校准 ImageNet 分类网络。
我的校准部分是
int batch_count = 4096; // calibrate sizeconfig->setInt8Calibrator(new NetworkInt8Calibrator(batch_count, { batch_size * xxx }, input_range, *network, std::cerr));
就是搞了 4096 个 batch 的数据（这个数看着设就行），后面那个类是自己实现的，负责告诉校准器每个 batch 的数据是什么
classNetworkInt8Calibrator :public nvinfer1::IInt8EntropyCalibrator2 {public: NetworkInt8Calibrator(int batches, conststd::vector<int64_t>& elemCount, conststd::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); } }boolgetBatch(void* bindings[], constchar* names[], int nbBindings)noexcept{if (mCurrentBatch >= mBatches) {returnfalse; }for (int i = 0; i < nbBindings; ++i) { bindings[i] = mInputDeviceBuffers[names[i]][mCurrentBatch]; } ++mCurrentBatch;returntrue; }intgetBatchSize()constnoexcept{return mBatches; }constvoid* readCalibrationCache(size_t& length)noexcept{ returnnullptr; }voidwriteCalibrationCache(constvoid*, 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 ，这个得根据需要选择，不同类型的网络不一样，详见 
https://docs.nvidia.com/deeplearning/tensorrt/developerguide/index.html#enable_int8_c
这里面有一个 read 和 write 函数我没实现，它们负责从文件中读和写 calibration cache 。如果用这个 cache 那么省去了生成数据和网络跑数据的时间，生成 Engine 时会更快一些。

注意：大白梳理对接AI行业的一些中高端岗位，年薪在50W~120W之间，图像算法、搜索推荐等热门岗位，欢迎感兴趣的小伙伴联系大白，提供全流程交流跟踪，各岗位详情如下：
《AI未来星球》陪伴你在AI行业成长的社群，各项福利重磅开放：
（1）198元《31节课入门人工智能》视频课程；
（2）大白花费近万元购买的各类数据集；
（3）每月自习活动，每月17日星球会员日，各类奖品送不停；
（4）加入《AI未来星球》内部微信群；
还有各类直播时分享的文件、研究报告，一起扫码加入吧！
人工智能行业，研究方向很多，大大小小有
几十个方向
。
为了便于大家学习交流，大白创建了一些不同方向的行业交流群。
每个领域，都有各方向的行业实战高手，和大家一起沟通交流。
目前主要开设：Opencv项目方面、目标检测方面、模型部署方面，后期根据不同领域高手的加入，建立新的方向群！
大家可以根据自己的兴趣爱好，加入对应的微信群，一起交流学习！
© THE END 
大家一起加油！

继续阅读
阅读原文