前段时间用 TensorRT 部署了一套模型,速度相比 Python 实现的版本快了 20 多倍,中间踩了许多坑,但是最后发现流程其实相当简单,特此记录一下踩坑过程。
顺便推荐一下深蓝学院的CUDA课程CUDA入门与深度神经网络加速 - 深蓝学院 - 专注人工智能与自动驾驶的学习平台 (shenlanxueyuan.com),如果想深入了解CUDA和TensorRT相关知识和运用的朋友可以关注一下
TensorRT
这东西就是 NVidia 在自家显卡上做了一个深度学习 inference 加速的框架,只要你把训练好的模型参数和结构告诉他,他就能自动帮你优化(硬件相关),以达到最快速度。
这涉及两个问题:
- 应该以什么模型格式把模型喂给 TensorRT?
- 如何使用 TensorRT 优化后的模型?
对于第一个问题:现在的深度学习框架非常多,不止常用的 pytorch/tensorflow,而即使是同一种框架还可以使用不同的编程语言实现。让 TensorRT 对每一个框架都直接支持,显然是不可能的。
TensorRT 只需要知道网络的结构和参数即可,它支持三种转换入口:
- TF-TRT,要求是 TensorFlow 模型
- ONNX 模型格式
- 使用 TensorRT API 手动把模型搭起来,然后把参数加载进去
第一种不够灵活,第三种比较麻烦,所以最省事方便的就是第二种方法。本文介绍第二种。
ONNX 就是一个通用的神经网络格式,一个 .onnx
文件内包含了网络的结构和参数。甭管是用什么深度学习框架写的网络,只要把模型导出成 ONNX 格式,就跟原本的代码没有关系了。
转成 ONNX 格式还没有被优化,需要再使用 TensorRT 读取它并优化成 TensorRT Engine。优化参数也在这一步指定。
对于第二个问题:得到的 TensorRT Engine 是硬件相关的,之后跑模型只需要运行这个 Engine 即可。调用 TensorRT Engine 需要使用 TensorRT Runtime API。
所以整个逻辑就是:
- 把你的模型导出成 ONNX 格式。
- 把 ONNX 格式模型输入给 TensorRT,并指定优化参数。
- 使用 TensorRT 优化得到 TensorRT Engine。
- 使用 TensorRT Engine 进行 inference。
你需要做的
- 把模型导出成 ONNX 格式。
- 安装 TensorRT 和 CUDA。注意二者和 driver 的版本号对应,我用的是 ZIP 安装:Installation Guide :: NVIDIA Deep Learning TensorRT Documentation,跟着这个把流程走一遍。
- 设置优化参数,使用 TensorRT 把 ONNX 优化成 Engine,得到当前硬件上优化后的模型。
- 使用 TensorRT Runtime API 进行 inference。
官方文档:Developer Guide :: NVIDIA Deep Learning TensorRT Documentation,写的很详细了。
ONNX 转换
Pytorch 自带导出方法 torch.onnx.export
。
TensorFlow 推荐使用这个 onnx/tensorflow-onnx: Convert TensorFlow, Keras, Tensorflow.js and Tflite models to ONNX (github.com) ,从 checkpoint (.meta 后缀)导出只需要
python -m tf2onnx.convert --checkpoint XXX.meta --inputs INPUT_NAME:0 --outputs OUTPUT_NAME:0 --output XXX.onnx
即可,tf 1.x/2.x 均可以用。
这里一定要注意 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 时会更快一些。