MXNet 的 TensorRT 集成指南从 contrib.tensorrt API 到子图融合加速【免费下载链接】mxnetLightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more项目地址: https://gitcode.com/gh_mirrors/mxnet1/mxnet本指南围绕 MXNet 仓库中mxnet.contrib.tensorrt模块对应 API 文档 contrib.tensorrt展开完整讲解该模块提供的 FP16 精度开关与 TensorRT 子图参数初始化接口并结合源码、教程与测试用例说明如何启用 TensorRT 图编译、如何使用符号式 API 完成加速推理以及 MXNet 内部子图扫描—算子融合—参数裁剪的实现原理。读完本文你将掌握set_use_fp16、get_use_fp16、init_tensorrt_params三个接口的用法并能独立搭建一条 MXNet TensorRT 的 GPU 推理流程。模块概览contrib.tensorrt 提供什么在 MXNet 的 Python 侧TensorRT 集成的入口是 python/mxnet/contrib/tensorrt.py该模块通过 python/mxnet/contrib/init.py 中的from . import tensorrt被导入为mxnet.contrib.tensorrt。API 文档页contrib.tensorrt通过 Sphinx 的automodule指令自动收集该模块的 docstring 与成员因此它的正文就是模块中的三个公开函数函数作用set_use_fp16(status)设置环境变量MXNET_TENSORRT_USE_FP16启用或禁用 TensorRT 的 FP16 精度get_use_fp16()读取MXNET_TENSORRT_USE_FP16返回 TensorRT 当前是否运行在 FP16 模式init_tensorrt_params(sym, arg_params, aux_params)把子图节点的权重写入 TensorRT 节点属性并返回剩余的非 TensorRT 参数在 C 侧真正的算子与图优化逻辑位于 src/operator/subgraph/tensorrt/ 目录包含tensorrt-inl.h、tensorrt.cc、tensorrt.cu、nnvm_to_onnx-inl.h、onnx_to_tensorrt.cc等文件。也就是说Python 层的三个函数只是开关与辅助工具真正把 MXNet 计算图翻译成 TensorRT 引擎的是一条完整的 C 子图融合链路下文会逐层展开。FP16 精度开关set_use_fp16 / get_use_fp16TensorRT 支持混合精度推理mxnet.contrib.tensorrt通过一个环境变量来控制在图编译时是否使用 FP16。import mxnet as mx # 启用 FP16整个 TensorRT 节点将强制以 FP16 执行 mx.contrib.tensorrt.set_use_fp16(True) # 查询当前模式返回 True 表示 FP16False 表示 FP32 is_fp16 mx.contrib.tensorrt.get_use_fp16() print(is_fp16)从源码看二者的实现非常直接set_use_fp16(status)执行os.environ[MXNET_TENSORRT_USE_FP16] str(int(status))见 python/mxnet/contrib/tensorrt.pyget_use_fp16()执行bool(int(os.environ.get(MXNET_TENSORRT_USE_FP16, 1)) 1)见 python/mxnet/contrib/tensorrt.py。两个细节值得注意默认值是开启get_use_fp16()在环境变量不存在时回退到1即默认按 FP16 处理。如果某个模型在 FP16 下精度不达标需要显式调用set_use_fp16(False)回退到 FP32。全节点强制语义docstring 明确指出FP16 模式会强制整个 TRT 节点以 FP16 执行The mode FP16 force the whole TRT node to be executed in FP16而不是按层自动选择精度。这意味着启用后图中所有被融合进 TensorRT 子图的算子都会统一使用 FP16 计算精度特性与混合精度自动调度AMP不同使用时需结合实际任务评估数值误差。由于该环境变量在 C 侧图编译时被读取set_use_fp16必须在构建/绑定 executor 之前调用才生效。子图参数初始化init_tensorrt_paramsinit_tensorrt_params(sym, arg_params, aux_params)用于把原本属于 TensorRT 子图内部的权重从 MXNet 参数表中搬移到子图节点的属性里这样 TensorRT 引擎构建时能直接拿到权重同时 MXNet 侧的参数表可以释放这些内存。其工作流程见 python/mxnet/contrib/tensorrt.py浅拷贝arg_params与aux_params避免污染调用方传入的原始字典遍历sym.get_internals()中的所有符号若某个节点带有subgraph_params_names属性该属性由 C 侧CreateSubgraphNode写入见 src/operator/subgraph/tensorrt/tensorrt-inl.h内容是以;分隔的子图参数名列表则逐个参数名检查参数在arg_params中 → 以subgraph_param_name为键移入tensorrt_params参数在aux_params中 → 同样以subgraph_param_name为键移入把tensorrt_params中的每个 NDArray 的handle.value即底层的 DLTensor 指针以字符串形式写入节点属性s._set_attr(**new_attrs)更新subgraph_params_names为剩余非 TensorRT参数名返回移除了子图参数后的arg_params, aux_params。典型用法如下与 tests/python/tensorrt/test_tensorrt_lenet5.py 中的模式一致import mxnet as mx sym, arg_params, aux_params mx.model.load_checkpoint(model, 0) # 获取 TensorRT 后端符号 trt_sym sym.get_backend_symbol(TensorRT) # 将子图权重写入 TensorRT 节点属性返回剩余参数 arg_params, aux_params mx.contrib.tensorrt.init_tensorrt_params( trt_sym, arg_params, aux_params) # 用剩余参数绑定 executor executor trt_sym.simple_bind(ctxmx.gpu(0), data(1, 3, 224, 224), grad_reqnull, force_rebindTrue) executor.copy_params_from(arg_params, aux_params)从源码结构看init_tensorrt_params与 C 侧的参数去重机制相配合教程文档指出MXNet 在图的初始化阶段会尝试移除仅在 TensorRT 段使用的重复权重以降低内存占用见 docs/python_docs/python/tutorials/performance/backend/tensorrt/tensorrt.md这个函数正是该机制在 Python 侧的执行入口。启用 TensorRT 的两种路径仓库中存在两种使用 TensorRT 的方式分别对应不同的MXNET_USE_TENSORRT环境变量用法。路径一tensorrt_bind实验性 API在 MXNet 1.3.0 时代集成以实验形态发布需要显式设置环境变量并调用mx.contrib.tensorrt.tensorrt_bind该 API 封装在 libmxnet 的 C 接口中Python 侧通过contrib.tensorrt暴露import os import mxnet as mx os.environ[MXNET_USE_TENSORRT] 1 # 开启 TensorRT 图编译 # 合并参数并搬到 GPU arg_params.update(aux_params) all_params dict([(k, v.as_in_context(mx.gpu(0))) for k, v in arg_params.items()]) # 用 tensorrt_bind 替代 simple_bind executor mx.contrib.tensorrt.tensorrt_bind( sym, ctxmx.gpu(0), all_paramsall_params, data(1, 3, 224, 224), grad_reqnull, force_rebindTrue)tensorrt_bind刻意模拟了simple_bind的签名区别在于参数以单个合并字典传入以配合上文所述的 TensorRT 权重清理流程。教程文档明确提到随着子图 API 的成熟社区目标是逐步弃用tensorrt_bind让用户透明地使用 TensorRT见 docs/python_docs/python/tutorials/performance/backend/tensorrt/tensorrt.md。由于simple_bind与tensorrt_bind高度相似迁移成本很低。路径二get_backend_symbol 子图 API推荐当前仓库的主路径是基于子图subgraph机制先通过sym.get_backend_symbol(TensorRT)得到优化后的符号再用init_tensorrt_params处理参数最后走常规的simple_bind。这正是 tests/python/tensorrt/test_tensorrt_lenet5.py 所验证的流程run_inference在use_tensorrtTrue时先取后端符号、初始化子图参数再绑定执行并把 MXNet 与 MXNet-TensorRT 的推理准确率差值控制在阈值3e-2以内见 tests/python/tensorrt/test_tensorrt_lenet5.py。注意Gluon 用户必须先hybridize()并把网络导出为符号export再用mx.model.load_checkpoint加载才能走 TensorRT 路径实验阶段的集成仅支持符号式 API见 docs/python_docs/python/tutorials/performance/backend/tensorrt/tensorrt.md。幕后原理算子扫描、子图融合与权重裁剪哪些算子可以进入 TensorRT 子图C 侧的选择器TensorrtSelector见 src/operator/subgraph/tensorrt/tensorrt-inl.h定义了算子兼容性判定isTRTCompatible可归纳为三类无条件支持unconditionalTRTops_copy、clip、elemwise_add/sub/mul、Flatten、Pad、relu、rsqrt、SoftmaxOutput带权重算子withWeightsOpsBatchNorm、Convolution、Deconvolution、FullyConnected其权重会以变量输入接入子图带条件支持的算子Pooling仅支持valid卷积约定或全局池化平均池化要求显式count_include_padfalse不支持 NHWC/NDHWC 布局Convolution/Deconvolution仅支持 NCHW/NCW/NCDHW 布局遇到 NHWC/NDHWC 或未知布局会打印警告并返回不支持Concat仅当拼接维dim ! 0时支持Dropout仅当mode kTraining且axes为空时支持推理期的 dropout 语义Activation仅relu/tanh/sigmoid三种激活类型BatchNorm要求axis 1即 NC(D)(H)W 布局。Select、SelectInput、SelectOutput方法进一步规定子图边界上的输入输出节点也必须兼容且带权重的算子只把非自身权重的变量纳入子图输入。Filter方法则要求候选子图至少包含两个非变量算子否则放弃融合——即单个算子不值得动用 TensorRT。融合与执行流程整个流程可概括为教程文档 docs/python_docs/python/tutorials/performance/backend/tensorrt/tensorrt.md 与源码相互印证MXNet 构建计算图后扫描其中所有算子找出连续且全部兼容的可融合区域每个区域被抽取出来替换为一个名为TensorRTid的特殊节点CreateSubgraphNode对应算子_TensorRT原区域被保存为该节点的子符号subgraphs[0]所有内部参数名写入subgraph_params_names属性执行到TensorRT节点时MXNet 调用 TensorRT 库TRTEngineParam见 src/operator/subgraph/tensorrt/tensorrt-inl.h持有ICudaEngine、IExecutionContext、IParser与TRT_Logger负责管理引擎的 binding 顺序与输入输出缓冲TensorRT 用自己优化过的内核常把多个算子融合进单个 CUDA kernel运行子图MXNet 只负责传入输入、取回输出权重去重仅存在于 TensorRT 段的参数从 MXNet 参数表中移除并释放内存这正是 Python 侧init_tensorrt_params配合完成的工作。模型转换链路NNVM → ONNX → TensorRT子图符号不能直接喂给 TensorRT中间需要一次模型格式转换。仓库的 src/operator/subgraph/tensorrt/nnvm_to_onnx-inl.h 与 src/operator/subgraph/tensorrt/onnx_to_tensorrt.cc 分别实现NNVM → ONNX把 MXNet 子符号转换成 ONNX 图ONNX → TensorRT借助onnx-tensorrt的NvOnnxParser头文件见 3rdparty/onnx-tensorrt/解析 ONNX 并构建 TensorRT 引擎。端到端示例ResNet-18 推理加速仓库教程docs/python_docs/python/tutorials/performance/backend/tensorrt/tensorrt.md给出了完整的 ResNet-18 示例流程import mxnet as mx from mxnet.gluon.model_zoo import vision import time import os batch_shape (1, 3, 224, 224) # 1. 从 Gluon Model Zoo 加载预训练 ResNet-18 并 hybridize resnet18 vision.resnet18_v2(pretrainedTrue) resnet18.hybridize() resnet18.forward(mx.nd.zeros(batch_shape)) resnet18.export(resnet18_v2) # 2. 以符号方式加载 sym, arg_params, aux_params mx.model.load_checkpoint(resnet18_v2, 0) # 3. MXNet 基线显式关闭 TensorRT os.environ[MXNET_USE_TENSORRT] 0 executor sym.simple_bind(ctxmx.gpu(0), databatch_shape, grad_reqnull, force_rebindTrue) executor.copy_params_from(arg_params, aux_params) # ... 预热 10 次后循环 forward 计时 ... # 4. TensorRT 路径开启环境变量改用 tensorrt_bind os.environ[MXNET_USE_TENSORRT] 1 arg_params.update(aux_params) all_params dict([(k, v.as_in_context(mx.gpu(0))) for k, v in arg_params.items()]) executor mx.contrib.tensorrt.tensorrt_bind(sym, ctxmx.gpu(0), all_paramsall_params, databatch_shape, grad_reqnull, force_rebindTrue) # ... 同样预热后计时 ...在教程记录的测试机器上Titan V GPUMXNet 基线耗时约 33.73 秒启用 TensorRT 后约 18.99 秒教程给出的结论是约 1.8 倍加速见 docs/python_docs/python/tutorials/performance/backend/tensorrt/tensorrt.md。加速主要来自算子融合ResNet 的整个计算图对 TensorRT 完全兼容因此优化后的图就是一个单一的TensorRT节点。需要说明的是该数据来自官方教程在特定硬件Titan V、CUDA 9.x 时代上的基准实际加速比会随模型、驱动、TensorRT 版本与输入规模变化。另外教程还提供了 Wavenet 优化前后的计算图可视化wavenet_unoptimized.svg/wavenet_optimized.svg位于 docs/python_docs/python/tutorials/performance/backend/tensorrt/ 目录可以直观看到多个子图被提取并替换为TensorRT节点的过程。安装与前置条件教程docs/python_docs/python/tutorials/performance/backend/tensorrt/tensorrt.md给出的实验性集成安装方式如下系统要求Ubuntu 16.04已更新显卡驱动安装 CUDA 9.0 或 9.2需要 Pascal 或更新的 NVIDIA GPUTensorRT 库按 NVIDIA 官方安装指南单独下载并安装 TensorRT 运行时库安装 MXNet TensorRT 构建PyPI 上有专门版本# CUDA 9.0 pip install mxnet-tensorrt-cu90 # CUDA 9.2 pip install mxnet-tensorrt-cu92使用官方 Docker 镜像其他操作系统或希望免去手工装环境nvidia-docker run -ti mxnet/tensorrt python测试代码 tests/python/tensorrt/test_tensorrt_lenet5.py 中的check_tensorrt_installation通过find_library(nvinfer)检查nvinfer共享库是否存在可作为运行 TensorRT 相关测试的前置自检手段。需要强调的是这些安装步骤针对的是仓库文档写作时MXNet 1.3.0 实验特性阶段的环境。当前仓库源码中的集成是实验性功能且实测代码要求 GPU 环境测试均绑定mx.gpu(0)。在实际使用时请以你所使用的 MXNet 发行版与 TensorRT 版本的兼容矩阵为准并在有 GPU 的机器上验证。配套测试与验证仓库在 tests/python/tensorrt/ 目录提供了完整测试集可作为 API 用法的权威参考test_tensorrt_lenet5.pyLeNet-5 在 MNIST 上的 MXNet 与 MXNet-TensorRT 推理准确率对比要求两者绝对差值小于3e-2test_resnet18.py、test_cvnets.py面向 ResNet-18 与 CV 网络的推理验证test_ops.py针对单个算子兼容性的细粒度测试。运行这些测试前先执行check_tensorrt_installation()确认环境再准备对应模型文件LeNet-5 测试通过LENET_MODEL_DIR环境变量指定模型目录默认/tmp。小结mxnet.contrib.tensorrt是 MXNet 接入 NVIDIA TensorRT 推理加速的 Python 入口本指南已覆盖其全部公开 APIset_use_fp16/get_use_fp16控制MXNET_TENSORRT_USE_FP16环境变量管理 FP16/FP32 精度模式init_tensorrt_params配合子图 API把 TensorRT 子图权重写入节点属性并释放 MXNet 参数表内存结合MXNET_USE_TENSORRT环境变量、tensorrt_bind旧实验 API与get_backend_symbol(TensorRT)现行子图 API两条推理路径配合 C 侧的TensorrtSelector算子兼容性判定、_TensorRT子图节点与 ONNX 转换链路构成了完整的图扫描—融合—引擎化加速闭环。如果你正在为 MXNet 模型做 GPU 推理加速可以从 tests/python/tensorrt/ 的测试用例出发先跑通 LeNet-5 的准确率对比再把同样的get_backend_symbolinit_tensorrt_params流程套用到自己的符号模型上。【免费下载链接】mxnetLightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more项目地址: https://gitcode.com/gh_mirrors/mxnet1/mxnet创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 MXNet 的 TensorRT 集成指南：从 contrib.tensorrt API 到子图融合加速
MXNet 的 TensorRT 集成指南从 contrib.tensorrt API 到子图融合加速【免费下载链接】mxnetLightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more项目地址: https://gitcode.com/gh_mirrors/mxnet1/mxnet本指南围绕 MXNet 仓库中mxnet.contrib.tensorrt模块对应 API 文档 contrib.tensorrt展开完整讲解该模块提供的 FP16 精度开关与 TensorRT 子图参数初始化接口并结合源码、教程与测试用例说明如何启用 TensorRT 图编译、如何使用符号式 API 完成加速推理以及 MXNet 内部子图扫描—算子融合—参数裁剪的实现原理。读完本文你将掌握set_use_fp16、get_use_fp16、init_tensorrt_params三个接口的用法并能独立搭建一条 MXNet TensorRT 的 GPU 推理流程。模块概览contrib.tensorrt 提供什么在 MXNet 的 Python 侧TensorRT 集成的入口是 python/mxnet/contrib/tensorrt.py该模块通过 python/mxnet/contrib/init.py 中的from . import tensorrt被导入为mxnet.contrib.tensorrt。API 文档页contrib.tensorrt通过 Sphinx 的automodule指令自动收集该模块的 docstring 与成员因此它的正文就是模块中的三个公开函数函数作用set_use_fp16(status)设置环境变量MXNET_TENSORRT_USE_FP16启用或禁用 TensorRT 的 FP16 精度get_use_fp16()读取MXNET_TENSORRT_USE_FP16返回 TensorRT 当前是否运行在 FP16 模式init_tensorrt_params(sym, arg_params, aux_params)把子图节点的权重写入 TensorRT 节点属性并返回剩余的非 TensorRT 参数在 C 侧真正的算子与图优化逻辑位于 src/operator/subgraph/tensorrt/ 目录包含tensorrt-inl.h、tensorrt.cc、tensorrt.cu、nnvm_to_onnx-inl.h、onnx_to_tensorrt.cc等文件。也就是说Python 层的三个函数只是开关与辅助工具真正把 MXNet 计算图翻译成 TensorRT 引擎的是一条完整的 C 子图融合链路下文会逐层展开。FP16 精度开关set_use_fp16 / get_use_fp16TensorRT 支持混合精度推理mxnet.contrib.tensorrt通过一个环境变量来控制在图编译时是否使用 FP16。import mxnet as mx # 启用 FP16整个 TensorRT 节点将强制以 FP16 执行 mx.contrib.tensorrt.set_use_fp16(True) # 查询当前模式返回 True 表示 FP16False 表示 FP32 is_fp16 mx.contrib.tensorrt.get_use_fp16() print(is_fp16)从源码看二者的实现非常直接set_use_fp16(status)执行os.environ[MXNET_TENSORRT_USE_FP16] str(int(status))见 python/mxnet/contrib/tensorrt.pyget_use_fp16()执行bool(int(os.environ.get(MXNET_TENSORRT_USE_FP16, 1)) 1)见 python/mxnet/contrib/tensorrt.py。两个细节值得注意默认值是开启get_use_fp16()在环境变量不存在时回退到1即默认按 FP16 处理。如果某个模型在 FP16 下精度不达标需要显式调用set_use_fp16(False)回退到 FP32。全节点强制语义docstring 明确指出FP16 模式会强制整个 TRT 节点以 FP16 执行The mode FP16 force the whole TRT node to be executed in FP16而不是按层自动选择精度。这意味着启用后图中所有被融合进 TensorRT 子图的算子都会统一使用 FP16 计算精度特性与混合精度自动调度AMP不同使用时需结合实际任务评估数值误差。由于该环境变量在 C 侧图编译时被读取set_use_fp16必须在构建/绑定 executor 之前调用才生效。子图参数初始化init_tensorrt_paramsinit_tensorrt_params(sym, arg_params, aux_params)用于把原本属于 TensorRT 子图内部的权重从 MXNet 参数表中搬移到子图节点的属性里这样 TensorRT 引擎构建时能直接拿到权重同时 MXNet 侧的参数表可以释放这些内存。其工作流程见 python/mxnet/contrib/tensorrt.py浅拷贝arg_params与aux_params避免污染调用方传入的原始字典遍历sym.get_internals()中的所有符号若某个节点带有subgraph_params_names属性该属性由 C 侧CreateSubgraphNode写入见 src/operator/subgraph/tensorrt/tensorrt-inl.h内容是以;分隔的子图参数名列表则逐个参数名检查参数在arg_params中 → 以subgraph_param_name为键移入tensorrt_params参数在aux_params中 → 同样以subgraph_param_name为键移入把tensorrt_params中的每个 NDArray 的handle.value即底层的 DLTensor 指针以字符串形式写入节点属性s._set_attr(**new_attrs)更新subgraph_params_names为剩余非 TensorRT参数名返回移除了子图参数后的arg_params, aux_params。典型用法如下与 tests/python/tensorrt/test_tensorrt_lenet5.py 中的模式一致import mxnet as mx sym, arg_params, aux_params mx.model.load_checkpoint(model, 0) # 获取 TensorRT 后端符号 trt_sym sym.get_backend_symbol(TensorRT) # 将子图权重写入 TensorRT 节点属性返回剩余参数 arg_params, aux_params mx.contrib.tensorrt.init_tensorrt_params( trt_sym, arg_params, aux_params) # 用剩余参数绑定 executor executor trt_sym.simple_bind(ctxmx.gpu(0), data(1, 3, 224, 224), grad_reqnull, force_rebindTrue) executor.copy_params_from(arg_params, aux_params)从源码结构看init_tensorrt_params与 C 侧的参数去重机制相配合教程文档指出MXNet 在图的初始化阶段会尝试移除仅在 TensorRT 段使用的重复权重以降低内存占用见 docs/python_docs/python/tutorials/performance/backend/tensorrt/tensorrt.md这个函数正是该机制在 Python 侧的执行入口。启用 TensorRT 的两种路径仓库中存在两种使用 TensorRT 的方式分别对应不同的MXNET_USE_TENSORRT环境变量用法。路径一tensorrt_bind实验性 API在 MXNet 1.3.0 时代集成以实验形态发布需要显式设置环境变量并调用mx.contrib.tensorrt.tensorrt_bind该 API 封装在 libmxnet 的 C 接口中Python 侧通过contrib.tensorrt暴露import os import mxnet as mx os.environ[MXNET_USE_TENSORRT] 1 # 开启 TensorRT 图编译 # 合并参数并搬到 GPU arg_params.update(aux_params) all_params dict([(k, v.as_in_context(mx.gpu(0))) for k, v in arg_params.items()]) # 用 tensorrt_bind 替代 simple_bind executor mx.contrib.tensorrt.tensorrt_bind( sym, ctxmx.gpu(0), all_paramsall_params, data(1, 3, 224, 224), grad_reqnull, force_rebindTrue)tensorrt_bind刻意模拟了simple_bind的签名区别在于参数以单个合并字典传入以配合上文所述的 TensorRT 权重清理流程。教程文档明确提到随着子图 API 的成熟社区目标是逐步弃用tensorrt_bind让用户透明地使用 TensorRT见 docs/python_docs/python/tutorials/performance/backend/tensorrt/tensorrt.md。由于simple_bind与tensorrt_bind高度相似迁移成本很低。路径二get_backend_symbol 子图 API推荐当前仓库的主路径是基于子图subgraph机制先通过sym.get_backend_symbol(TensorRT)得到优化后的符号再用init_tensorrt_params处理参数最后走常规的simple_bind。这正是 tests/python/tensorrt/test_tensorrt_lenet5.py 所验证的流程run_inference在use_tensorrtTrue时先取后端符号、初始化子图参数再绑定执行并把 MXNet 与 MXNet-TensorRT 的推理准确率差值控制在阈值3e-2以内见 tests/python/tensorrt/test_tensorrt_lenet5.py。注意Gluon 用户必须先hybridize()并把网络导出为符号export再用mx.model.load_checkpoint加载才能走 TensorRT 路径实验阶段的集成仅支持符号式 API见 docs/python_docs/python/tutorials/performance/backend/tensorrt/tensorrt.md。幕后原理算子扫描、子图融合与权重裁剪哪些算子可以进入 TensorRT 子图C 侧的选择器TensorrtSelector见 src/operator/subgraph/tensorrt/tensorrt-inl.h定义了算子兼容性判定isTRTCompatible可归纳为三类无条件支持unconditionalTRTops_copy、clip、elemwise_add/sub/mul、Flatten、Pad、relu、rsqrt、SoftmaxOutput带权重算子withWeightsOpsBatchNorm、Convolution、Deconvolution、FullyConnected其权重会以变量输入接入子图带条件支持的算子Pooling仅支持valid卷积约定或全局池化平均池化要求显式count_include_padfalse不支持 NHWC/NDHWC 布局Convolution/Deconvolution仅支持 NCHW/NCW/NCDHW 布局遇到 NHWC/NDHWC 或未知布局会打印警告并返回不支持Concat仅当拼接维dim ! 0时支持Dropout仅当mode kTraining且axes为空时支持推理期的 dropout 语义Activation仅relu/tanh/sigmoid三种激活类型BatchNorm要求axis 1即 NC(D)(H)W 布局。Select、SelectInput、SelectOutput方法进一步规定子图边界上的输入输出节点也必须兼容且带权重的算子只把非自身权重的变量纳入子图输入。Filter方法则要求候选子图至少包含两个非变量算子否则放弃融合——即单个算子不值得动用 TensorRT。融合与执行流程整个流程可概括为教程文档 docs/python_docs/python/tutorials/performance/backend/tensorrt/tensorrt.md 与源码相互印证MXNet 构建计算图后扫描其中所有算子找出连续且全部兼容的可融合区域每个区域被抽取出来替换为一个名为TensorRTid的特殊节点CreateSubgraphNode对应算子_TensorRT原区域被保存为该节点的子符号subgraphs[0]所有内部参数名写入subgraph_params_names属性执行到TensorRT节点时MXNet 调用 TensorRT 库TRTEngineParam见 src/operator/subgraph/tensorrt/tensorrt-inl.h持有ICudaEngine、IExecutionContext、IParser与TRT_Logger负责管理引擎的 binding 顺序与输入输出缓冲TensorRT 用自己优化过的内核常把多个算子融合进单个 CUDA kernel运行子图MXNet 只负责传入输入、取回输出权重去重仅存在于 TensorRT 段的参数从 MXNet 参数表中移除并释放内存这正是 Python 侧init_tensorrt_params配合完成的工作。模型转换链路NNVM → ONNX → TensorRT子图符号不能直接喂给 TensorRT中间需要一次模型格式转换。仓库的 src/operator/subgraph/tensorrt/nnvm_to_onnx-inl.h 与 src/operator/subgraph/tensorrt/onnx_to_tensorrt.cc 分别实现NNVM → ONNX把 MXNet 子符号转换成 ONNX 图ONNX → TensorRT借助onnx-tensorrt的NvOnnxParser头文件见 3rdparty/onnx-tensorrt/解析 ONNX 并构建 TensorRT 引擎。端到端示例ResNet-18 推理加速仓库教程docs/python_docs/python/tutorials/performance/backend/tensorrt/tensorrt.md给出了完整的 ResNet-18 示例流程import mxnet as mx from mxnet.gluon.model_zoo import vision import time import os batch_shape (1, 3, 224, 224) # 1. 从 Gluon Model Zoo 加载预训练 ResNet-18 并 hybridize resnet18 vision.resnet18_v2(pretrainedTrue) resnet18.hybridize() resnet18.forward(mx.nd.zeros(batch_shape)) resnet18.export(resnet18_v2) # 2. 以符号方式加载 sym, arg_params, aux_params mx.model.load_checkpoint(resnet18_v2, 0) # 3. MXNet 基线显式关闭 TensorRT os.environ[MXNET_USE_TENSORRT] 0 executor sym.simple_bind(ctxmx.gpu(0), databatch_shape, grad_reqnull, force_rebindTrue) executor.copy_params_from(arg_params, aux_params) # ... 预热 10 次后循环 forward 计时 ... # 4. TensorRT 路径开启环境变量改用 tensorrt_bind os.environ[MXNET_USE_TENSORRT] 1 arg_params.update(aux_params) all_params dict([(k, v.as_in_context(mx.gpu(0))) for k, v in arg_params.items()]) executor mx.contrib.tensorrt.tensorrt_bind(sym, ctxmx.gpu(0), all_paramsall_params, databatch_shape, grad_reqnull, force_rebindTrue) # ... 同样预热后计时 ...在教程记录的测试机器上Titan V GPUMXNet 基线耗时约 33.73 秒启用 TensorRT 后约 18.99 秒教程给出的结论是约 1.8 倍加速见 docs/python_docs/python/tutorials/performance/backend/tensorrt/tensorrt.md。加速主要来自算子融合ResNet 的整个计算图对 TensorRT 完全兼容因此优化后的图就是一个单一的TensorRT节点。需要说明的是该数据来自官方教程在特定硬件Titan V、CUDA 9.x 时代上的基准实际加速比会随模型、驱动、TensorRT 版本与输入规模变化。另外教程还提供了 Wavenet 优化前后的计算图可视化wavenet_unoptimized.svg/wavenet_optimized.svg位于 docs/python_docs/python/tutorials/performance/backend/tensorrt/ 目录可以直观看到多个子图被提取并替换为TensorRT节点的过程。安装与前置条件教程docs/python_docs/python/tutorials/performance/backend/tensorrt/tensorrt.md给出的实验性集成安装方式如下系统要求Ubuntu 16.04已更新显卡驱动安装 CUDA 9.0 或 9.2需要 Pascal 或更新的 NVIDIA GPUTensorRT 库按 NVIDIA 官方安装指南单独下载并安装 TensorRT 运行时库安装 MXNet TensorRT 构建PyPI 上有专门版本# CUDA 9.0 pip install mxnet-tensorrt-cu90 # CUDA 9.2 pip install mxnet-tensorrt-cu92使用官方 Docker 镜像其他操作系统或希望免去手工装环境nvidia-docker run -ti mxnet/tensorrt python测试代码 tests/python/tensorrt/test_tensorrt_lenet5.py 中的check_tensorrt_installation通过find_library(nvinfer)检查nvinfer共享库是否存在可作为运行 TensorRT 相关测试的前置自检手段。需要强调的是这些安装步骤针对的是仓库文档写作时MXNet 1.3.0 实验特性阶段的环境。当前仓库源码中的集成是实验性功能且实测代码要求 GPU 环境测试均绑定mx.gpu(0)。在实际使用时请以你所使用的 MXNet 发行版与 TensorRT 版本的兼容矩阵为准并在有 GPU 的机器上验证。配套测试与验证仓库在 tests/python/tensorrt/ 目录提供了完整测试集可作为 API 用法的权威参考test_tensorrt_lenet5.pyLeNet-5 在 MNIST 上的 MXNet 与 MXNet-TensorRT 推理准确率对比要求两者绝对差值小于3e-2test_resnet18.py、test_cvnets.py面向 ResNet-18 与 CV 网络的推理验证test_ops.py针对单个算子兼容性的细粒度测试。运行这些测试前先执行check_tensorrt_installation()确认环境再准备对应模型文件LeNet-5 测试通过LENET_MODEL_DIR环境变量指定模型目录默认/tmp。小结mxnet.contrib.tensorrt是 MXNet 接入 NVIDIA TensorRT 推理加速的 Python 入口本指南已覆盖其全部公开 APIset_use_fp16/get_use_fp16控制MXNET_TENSORRT_USE_FP16环境变量管理 FP16/FP32 精度模式init_tensorrt_params配合子图 API把 TensorRT 子图权重写入节点属性并释放 MXNet 参数表内存结合MXNET_USE_TENSORRT环境变量、tensorrt_bind旧实验 API与get_backend_symbol(TensorRT)现行子图 API两条推理路径配合 C 侧的TensorrtSelector算子兼容性判定、_TensorRT子图节点与 ONNX 转换链路构成了完整的图扫描—融合—引擎化加速闭环。如果你正在为 MXNet 模型做 GPU 推理加速可以从 tests/python/tensorrt/ 的测试用例出发先跑通 LeNet-5 的准确率对比再把同样的get_backend_symbolinit_tensorrt_params流程套用到自己的符号模型上。【免费下载链接】mxnetLightweight, Portable, Flexible Distributed/Mobile Deep Learning with Dynamic, Mutation-aware Dataflow Dep Scheduler; for Python, R, Julia, Scala, Go, Javascript and more项目地址: https://gitcode.com/gh_mirrors/mxnet1/mxnet创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考