TensorRT,由NVIDIA精心打造,是一款专为深度学习推理任务设计的高效优化工具。它能够显著降低应用的延迟同时提升处理速度,是大规模数据中心、嵌入式设备以及自动驾驶技术等领域推理加速的理想选择。TensorRT的兼容性广泛,支持包括TensorFlow、Caffe、Mxnet和Pytorch在内的主流深度学习框架,与NVIDIA GPU的结合使用。
TensorRT本质上是一个推理加速器,它允许用户在模型训练完成后,直接将模型文件导入TensorRT进行优化,无需再依赖原始的深度学习框架。
模型训练:需要fp32,避免梯度下溢和舍入误差
梯度下溢:当梯度值非常小(小于 ~6e-8)时,在FP16中会变成0(即下溢)。一旦梯度变为0,对应的参数就无法更新,学习过程会停滞。
舍入误差:在FP16中,由于精度有限,很多微小的更新(比如用一个很小的学习率乘以一个很小的梯度)可能会因为舍入误差而直接变为0。在数百万次的迭代中,这些微小的误差累积起来可能会导致模型无法收敛,或者收敛到一个很差的局部最优点。
模型推理:不涉及到反向传播(梯度计算),只需fp16
在FP32下训练好的模型,其权重和激活值的分布通常在一个合理的范围内。将其转换为FP16进行推理,对于绝大多数任务(如图像分类、目标检测、语音识别等)的最终输出精度影响微乎其微,通常只有小数点后几位的损失,但换来的却是巨大的速度提升和资源节省。
TensorRT 采用多种优化技术来提升深度学习模型的推理性能:
1、层间融合技术:
TensorRT 通过层间融合,将卷积层、偏置层和ReLU激活层合并为单一的CBR结构,实现横向和纵向的层融合。横向融合将这些层合并为单一操作,仅消耗一个CUDA核心,而纵向融合则将具有相同结构但不同权重的层合并成更宽的层,同样只占用一个CUDA核心。这种融合减少了计算图中的层数,降低了CUDA核心的使用量,从而使得模型结构更加紧凑、运行速度更快、效率更高。
2、数据精度优化:
在深度学习模型训练过程中,通常使用32位浮点数(FP32)来保证精度。然而,在推理阶段,由于不需要进行反向传播,可以安全地降低数据精度至FP16或INT8,这不仅减少了内存占用和延迟,还使得模型体积更小,提高了推理速度。
3、Kernel自动调优:
TensorRT 能够自动调整CUDA核心的计算方式,以适应不同的算法、模型结构和GPU平台。这种自动调优确保了模型在特定硬件上以最佳性能运行。
TensorRT 部署流程主要有以下五步:
1.训练模型
2.导出模型为 ONNX 格式
3.选择精度
4.转化成 TensorRT 模型
5.部署模型
Pytorch → ONNX → TensorRT
| 特性/维度 | PyTorch (直接推理) | ONNX Runtime (ONNX Runtime) | TensorRT |
|---|---|---|---|
| 核心定位 | 研究与训练框架,兼顾推理 | 跨平台高性能推理引擎 | NVIDIA平台极致性能推理SDK |
| 工作流程 | torch.jit.trace 或 torch.jit.script 导出模型 -> 用PyTorch C++/Python加载推理 |
PyTorch/TF等 -> 导出ONNX模型 -> ONNX Runtime加载并推理 | PyTorch/TF/ONNX -> 转换和优化 -> 生成TensorRT引擎 -> 执行引擎 |
| 性能 | 较好。利用了PyTorch的优化,但包含一些为训练准备的冗余操作。 | 优秀。专为推理优化,支持图优化、算子融合等,性能通常优于原生PyTorch。 | 极致。在NVIDIA GPU上通常是性能天花板。通过内核融合、量化、自动内核调优等技术达到最优。 |
| 硬件支持 | 主要支持CPU和NVIDIA GPU(通过CUDA)。对其它硬件(如AMD GPU, NPU)支持有限或需要额外适配。 | 非常广泛。通过Execution Providers机制,支持CPU、NVIDIA CUDA/TensorRT、AMD ROCm、Intel OpenVINO、ARM NN、Android NNAPI等。“一次转换,多处运行”。 | 仅限NVIDIA GPU。包括GeForce, Tesla, Jetson等系列。 |
| 易用性与灵活性 | 最高。直接在训练框架内完成,调试方便,动态图模型(如带有控制流的)处理起来更自然。 | 高。ONNX是一个开放的标准,生态庞大。转换过程可能遇到不支持的算子,需要一些调试。 | 较低。优化和构建引擎需要时间,过程像是一个“黑盒”,调试难度较大。量化等操作可能会引入精度损失。 |
| 关键技术 | TorchScript, torch.jit |
ONNX格式、图优化、多种Execution Providers | 层与内核融合、精度校准(INT8量化)、内核自动调优、动态Tensor优化 |
| 典型应用场景 | 1. 研究原型快速验证 2. 对延迟不敏感的服务器端应用 3. 需要利用PyTorch动态图特性的场景 |
1. 需要跨平台部署(如同时部署在服务器CPU和边缘设备) 2. 希望获得比原生框架更好的性能,同时又不想被单一硬件厂商锁定 3. 云原生和微服务环境 |
1. 对延迟和吞吐量要求极致的场景(如自动驾驶、视频直播、高并发推荐系统) 2. 边缘设备(NVIDIA Jetson) 3. 需要INT8量化来大幅提升性能并降低功耗的场景 |