TensorRT是可以在NVIDIA各种GPU硬件平台下运行的一个C++推理框架。我们利用Pytorch、TF或者其他框架训练好的模型,可以转化为TensorRT的格式,然后利用TensorRT推理引擎去运行我们这个模型,从而提升这个模型在英伟达GPU上运行的速度。速度提升的比例是比较可观的。
TensorRT是由C++、CUDA、python三种语言编写成的一个库,其中核心代码为C++和CUDA,Python端作为前端与用户交互。当然,TensorRT也是支持C++前端的,如果我们追求高性能,C++前端调用TensorRT是必不可少的。
下面是网上有人总结的具体TensorRT的加速效果:
训练框架得到的模型通过第二步的一系列转换优化,得到一个序列存下来,然后实际线上运行的时候再通过反序列出TRT的engine模型进行运行。
其中1–4属于上面的第二步,6–8属于第四步,下面是具体的代码过程:
然后具体的代码案例可到其开源代码中的sample文件夹中进行查看。
implicit(隐式) batch
Build 阶段设置:
IBuilder:: createNetwork();
IBuilder:: setMaxBatchSize(maxBatchSize);
(这里面的maxBatchSize根据显卡显存确定)
Infer阶段:
enqueue(batchSize, data, stream, nullptr);
(在infer阶段batchSize不必对应上面的maxBatchSize,可以是具体运行中的batch)
explicit(显式) batch
Build 阶段设置:
IBuilder::createNetworkV2(1U <<
static_cast(NetworkDefinitionCreationFlag::kEXPLICIT_BATCH))
builder->setMaxBatchSize(maxBatchSize);
(以上几步没有变化)
IOptimizationProfile* profile = builder.createOptimizationProfile();
profile->setDimensions(“foo”, OptProfileSelector::kMIN, Dims3(3,100,200); // 运行中使用到的最小batch
profile->setDimensions(“foo”, OptProfileSelector::kOPT, Dims3(3,150,250); // 运行中最多使用的batch
profile->setDimensions(“foo”, OptProfileSelector::kMAX, Dims3(3,200,300); // 运行中使用到的最大batch
config.addOptimizationProfile(profile)
context.setOptimizationProfile(0)
Infer阶段:
context->setBindingDimensions(i, input_dim);
context->allInputDimensionsSpecified();
context->enqueueV2(data, stream, nullptr);
ONNX :https://github.com/NVIDIA/TensorRT/tree/main/parsers
Pytorch:https://github.com/NVIDIA-AI-IOT/torch2trt
对于pytorch还可以通过pytorch->onnx->TRT
TensorFlow:
https://github.com/tensorflow/tensorflow/tree/1cca70b80504474402215d2a4e55bc44621
b691d/tensorflow/compiler/tf2tensorrt
Tencent Forward:https://github.com/Tencent/Forward