简介面向算法工程师与深度学习部署开发者这套部署实战资源聚焦于使用 TensorRT 推理加速器完成 MobileViT 视觉模型的落地转换。流程设计覆盖模型训练与验证、ONNX 中间表示转换、TensorRT 引擎构建、自定义插件接入、量化校准以及最终推理测试能够帮助读者打通从 PyTorch 代码到 GPU 高性能推理的完整链路。压缩包共包含51个文件核心为21个Python脚本与11个编译后的pyc文件其中涉及模型转换、引擎构建、量化校准等工具另有2个CUDA算子源码及其头文件配合说明文档、演示幻灯片、权重文件和测试数据整体大小约为64.25MB目前已有136人学习或下载。项目结合MobileViT轻量化结构与视觉Transformer的全局建模优势示范了剪枝、层融合、精度校准等优化操作在降低显存占用和计算开销的同时提升推理速度。读者能够从中掌握TensorRT的接口调用、插件开发及int8量化部署方法理解边缘设备上的模型压缩与加速思路适合有一定深度学习基础且希望完成算法工程化落地的中高级开发人员。1. 用 TensorRT 部署 MobileViT这份算法部署项目先解决哪个卡点在算法部署这个方向上MobileViT 是个很典型的“看着轻量转起来处处是坑”的模型。它把 MobileNet 的卷积和 Vision Transformer 的 attention 混在一起想在保持精度的同时压体积可真到了 TensorRT 这一步self-attention 和 LayerNorm 这类算子并不总能被原生解析器干净地消化。这份“使用 TensorRT 部署 MobileViT”的算法部署项目把从 PyTorch 权重导出 ONNX、写自定义 plugin、做 INT8 校准再到精度对比和 benchmark 的完整链路都打包好了。适合三种人要把 MobileViT 落到 NVIDIA GPU 上做实时推理的工程师想搞明白 TensorRT 自定义插件怎么写的算法岗需要一套可复现部署基线去估算并发路数的团队。2. 部署链路拆解从 PyTorch 权重到 TensorRT 引擎每一步在做什么部署一个 MobileViT不是把 .pt 文件丢给 trtexec 就能完事的。MobileViT 内部有 unfold、fold、Multi-head Self-Attention、LayerNorm 这些结构ONNX 导出一旦出问题后面 TensorRT 构建引擎时要么报算子不支持要么构建成功但输出完全不对。所以先按项目目录把链路理清楚再执行每一步比直接跑脚本靠谱得多。2.1 先认目录这份项目把部署切成了哪几个阶段拿到资源包后别急着解压跑命令。先看文件是怎么分组的基本就能猜到作者的部署路径。我把关键文件整理成了下表阶段涉及文件在链路里干什么模型导出convert_to_onnx.pyPyTorch 权重转 ONNX支持动态 batch图修改onnx_add_plugin.py在 ONNX 图里把 attention / LayerNorm 子图替换成自定义节点引擎构建convert_to_trt.pyONNX 转 TensorRT engine支持 FP16 / INT8插件实现attentionPlugin.cu / .h、layerNormPlugin.cu / .hCUDA kernel 与 TensorRT plugin 封装校准工具calibrator.py、imagenet_lmdb_datasets.pyINT8 量化时加载校准集测试验证test_trt.py、test_torch_precision.py、test_trt_precision.py、test_attention_plugin.py冒烟测试、精度对比、插件单测辅助工具gen_test_data.py、benchmark、Makefile、doc造数据、跑性能、一键编译这个项目最有价值的部分不是 convert_to_trt.py 本身而是 attentionPlugin 和 layerNormPlugin 那两对文件。MobileViT 的 transformer 块里LayerNorm 和 attention 在 TensorRT 的 ONNX parser 里经常要一层一层拆解才能支持一旦拆错精度就崩了。项目把它固化成 plugin相当于把最容易翻车的那段预先封装好。weights-file 目录里应该放预训练权重如果你拿到的包里这个目录是空的先准备好 MobileViT 的 PyTorch 权重再往下走。2.2 PyTorch 模型导出 ONNXconvert_to_onnx.py 怎么用先把最基础的导出命令跑通。多数场景下下面这几项就够用# 导出 MobileViT 到 ONNXweights 路径按实际位置改 python convert_to_onnx.py \ --weights weights-file/mobilevit_s.pt \ --output mobilevit_s.onnx \ --opset 13 \ --dynamic-batchopset 选 13 是有讲究的。TensorRT 对 ONNX 算子版本的支持通常滞后opset 开太高反而容易在解析时踩到不认识的算子组合。MobileViT 的 attention 里有些操作在 opset 13 下有明确的算子映射配合项目里的 onnx_add_plugin.py 改图比一味升级 opset 更可控。再看 convert_to_onnx.py 内部最常见的实现逻辑。它本质上就是一次 torch.onnx.exportimport torch from models import MobileViT # 项目里的 MobileViT 定义 model MobileViT(num_classes1000) model.load_state_dict(torch.load(args.weights, map_locationcpu)[model]) model.eval() # dummy 输入的空间尺寸必须和训练/部署时保持一致 dummy torch.randn(1, 3, 256, 256) torch.onnx.export( model, dummy, args.output, opset_versionargs.opset, input_names[input], output_names[output], dynamic_axes{ input: {0: batch}, output: {0: batch}, }, do_constant_foldingTrue, )逻辑说明dummy 输入的 HxW 直接决定了后面 TensorRT 里输入张量的形状。MobileViT 的 attention 逻辑和输入分辨率绑定在一起训练时用 256部署时最好别随手改成其他值。dynamic_axes 这里只给 batch 维开动态是因为在 TensorRT 的 optimization profile 里宽高维度参与动态变化会大幅增加 shape 推断的复杂度容易触发限制。注意导出 ONNX 后先用 onnxruntime 验证一下输出再进 TensorRT。常见做法是拿同一张图分别跑 PyTorch 和 ONNX Runtimetop-1 一致再往下走这一步能挡掉后续 80% 的问题。2.3 ONNX 转 TensorRT 引擎convert_to_trt.py 的参数选择很多人拿到 ONNX 第一反应是跑trtexec --onnxmobilevit_s.onnx --fp16。这个命令做基准测试没问题但做不了两件事一是加载自定义 plugin二是挂自定义 INT8 calibrator。项目里用 convert_to_trt.py 而不是直接甩一条 trtexec正是因为 attentionPlugin 和 layerNormPlugin 需要先构建成 .so 注册进 TensorRT再通过 ONNX 图中的 plugin 节点加载。convert_to_trt.py 的骨架一般是这样的import tensorrt as trt logger trt.Logger(trt.Logger.WARNING) builder trt.Builder(logger) network builder.create_network(1 int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) profile builder.create_optimization_profile() # 输入 1x3x256x256最小 batch 1常见 batch 4最大 batch 8 input_tensor network.add_input(input, trt.float32, (-1, 3, 256, 256)) profile.set_shape(input_tensor.name, (1, 3, 256, 256), (4, 3, 256, 256), (8, 3, 256, 256)) config builder.create_builder_config() config.add_optimization_profile(profile) config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 30) # workspace 1GB config.set_flag(trt.BuilderFlag.FP16) # 开 FP16参数说明profile 的三个 shape 分别是最小、常见、最大 batch。常见值选 4是因为线上推理多数按 batch 4 喂图性能曲线比较平。workspace 给 1GB 是给算子融合留出尝试空间太小的话某些融合方案不会触发。构建完成后用serialize()落盘 engine 文件之后部署只加载 engine不再走 parser。这里有个容易忽略的点TensorRT 的输出不像 ONNX Runtime 那样直接给 Python ndarray它需要预先绑定额外的设备内存。这个细节在第 4 章的 test_trt.py 里会体现。如果你之前只玩过 trtexec建议先老老实实把这一个 Python API 链路跑通再回去用 trtexec 做压力测试。3. 自定义插件与 INT8 校准跑通 MobileViT 的两个关键关卡引擎能不能构建出来很多时候取决于自定义插件写得对不对。MobileViT 里最核心的两个结构是 LayerNorm 和 Multi-head Self-AttentionTensorRT 对它们的原生支持并不完整。项目里那对 .cu 文件就是专门补这块短板的。3.1 为什么 attention 和 LayerNorm 要单独写插件MobileViT 的 transformer 块结构可以简化理解成把 feature map unfold 成 patch在 patch 内做 self-attention再 fold 回去。这个“unfold attention fold”的组合在 ONNX 里会被展开成大量 Reshape、Transpose、MatMul 节点。TensorRT 的 ONNX parser 遇到这种长串子图做得最多的事情是“战术性放弃”——要么解析失败要么解析出来但效率很差。attentionPlugin.cu 的目标是把整段 attention 子图合并成一个自定义 operator。在 device 端直接完成 patch 切分、QKV 计算、softmax、加权求和一次 kernel 调用把整段算完。LayerNorm 也是同样的道理MobileViT 里 LayerNorm 作用在最后一维unfold 之后维度顺序经常是 [B, P, N, C]TensorRT 用 ReduceMean、Sub、Div 的组合能实现功能但中间会多出好几层临时张量性能上不划算。看 plugin 接口的写法// attentionPlugin.h 里暴露的派生类核心在 enqueue 里调 kernel class AttentionPlugin : public trt::IPluginV2DynamicExt { public: const char* getPluginType() const override { return AttentionPlugin; } int32_t enqueue(const trt::PluginTensorDesc* inputDesc, const trt::PluginTensorDesc* outputDesc, const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) override { attention_kernel_launch(..., stream); // attentionPlugin.cu 里实现 return 0; } };逻辑说明enqueue是真正在 GPU 上干活的入口所有输入输出都在设备端不能在这里面做 cudaMemcpy 这类同步操作否则性能直接崩。getPluginType这个方法决定了 engine 反序列化时能不能找到这个插件。项目里的 Makefile 就是用来把这些 .cu 编成 .so 的编完之后要保证加载路径对得上否则推理时会报 unknown plugin。3.2 onnx_add_plugin.py把不相容的算子替换成自定义节点有了 plugin 本体下一步就是把 ONNX 图里的旧子图替换成 plugin 节点。这一步用的是 onnx_add_plugin.py。# 先用 convert_to_onnx.py 得到原始模型再改图 python onnx_add_plugin.py --onnx mobilevit_s.onnx --output mobilevit_s_plugin.onnx这类脚本最常见的实现是基于 onnx_graphsurgeon。匹配 LayerNorm 和 Attention 子图把匹配到的节点删除插入自定义节点再重新导出import onnx import onnx_graphsurgeon as gs graph gs.import_onnx(onnx.load(args.onnx)) # 找到 attention 子图的输入输出节点具体匹配逻辑要看实际模型结构 x graph.tensors[args.input_node] out graph.tensors[args.output_node] # 删除旧子图节点插入自定义节点 plugin_node gs.Node(opAttentionPlugin, inputs[x], outputs[out]) graph.nodes.append(plugin_node) graph.cleanup().toposort() onnx.save(gs.export_onnx(graph), args.output)逻辑说明改图最怕匹配不完整导致旧节点还残留在图上执行时等于跑了两遍。我一般改完图先跑一下onnx.checker.check_model再让 TensorRT 把 network 的层数打印出来对比改图前后的层数变化。如果层数只减了一点点多半是替换没生效plugin 只是挂名实际还在跑旧子图。3.3 calibrator.py 与 ImageNet LMDBINT8 校准集怎么喂FP16 跑通之后想再压吞吐就得看 INT8。MobileViT 这类轻量模型对 INT8 量化比大模型更敏感校准集的质量直接决定精度掉多少。项目的 calibrator.py 继承的是 TensorRT 的 IInt8Calibrator2 接口校准数据从 ImageNet LMDB 里读。# calibrator.py 的骨架lmdb 按序读图速度快且缓存稳定 import tensorrt as trt class MobileViTCalibrator(trt.IInt8Calibrator2): def __init__(self, lmdb_path, batch_size16): super().__init__() self.dataset imagenet_lmdb_datasets.ImageNetLMDB(lmdb_path) self.batch_size batch_size self.buffers [torch.zeros(batch_size, 3, 256, 256, dtypetorch.uint8)] def get_batch_size(self): return self.batch_size def get_batch(self, names): # 读下一批图返回设备端指针列表 return [self.buffers[i].data_ptr() for i in range(len(names))] def read_calibration_cache(self): with open(calib.cache, rb) as f: return f.read() def write_calibration_cache(self, cache): with open(calib.cache, wb) as f: f.write(cache)参数说明get_batch返回的必须是设备端指针所以 buffers 要预先放在 GPU 上。校准集一般从 ImageNet 验证集抽 500 到 1000 张类别分布要均衡不能 80% 都是同一个类。构建 INT8 引擎时把 calibrator 挂到 config 上config.set_flag(trt.BuilderFlag.INT8) config.int8_calibrator MobileViTCalibrator(imagenet_val.lmdb, batch_size16)校准 cache 文件建议留档。有了 calib.cache后续重新构建引擎不用再跑一遍校准而且能保证同一批图像的量化结果可复现。没有 cache 的话每次构建 INT8 引擎都重新抽样精度波动会很大调起来特别难定位问题。4. 精度对比与性能基准三个测试脚本帮你判断部署有没有翻车引擎构建成功不等于部署成功。MobileViT 这类带 attention 的模型最常出现的情况是精度悄悄掉了几个点单看一两张图完全看不出来。所以项目里专门放了 test_torch_precision.py 和 test_trt_precision.py用来对齐两边的输出。我每次转换完模型第一件事不是看速度而是跑精度对比。4.1 test_torch_precision.py 和 test_trt_precision.py同一份数据双端跑执行顺序上先跑 test_torch_precision.py 生成 PyTorch 端的输出基线再用 test_trt_precision.py 加载 engine 输出。两边用同一份预处理后的数据按 cosine similarity 和最大绝对误差来对比# test_trt_precision.py 的对比逻辑简化 import numpy as np trt_outputs run_engine(engine_path, inputs) # TensorRT 端到端推理 torch_outputs torch_ref_outputs # 来自 test_torch_precision.py 的保存结果 for i, (a, b) in enumerate(zip(trt_outputs, torch_outputs)): a np.asarray(a).flatten() b np.asarray(b).flatten() cos np.dot(a, b) / (np.linalg.norm(a) * np.linalg.norm(b) 1e-12) max_diff np.max(np.abs(a - b)) print(f输出[{i}] cosine{cos:.6f} max_diff{max_diff:.6f})判断阈值怎么给我一般这么设FP16 引擎要求 cosine 大于 0.999max_diff 取决于 logits 的数值尺度通常在 0.05 以内。INT8 引擎要求 cosine 大于 0.99同时看 top-1 掉点不超过 1 个点。不要只拿一个 cosine 下结论要同时看 top-5 和每个类别的 logits 排序。cosine 这个指标对整体数值方向敏感但对少数类别的小幅错位不敏感只看它容易漏掉问题。4.2 test_trt.py引擎构建完先做冒烟测试精度对比之前先跑冒烟测试确认引擎能正常推理。命令长这样python test_trt.py --engine mobilevit_s_fp16.engine --input data/test.jpg里面关键的一步是绑定输入输出 buffer。TensorRT 的动态 batch 必须在执行前显式设置形状# test_trt.py 关键步骤 engine deserialize_engine(args.engine) context engine.create_execution_context() context.set_input_shape(input, (1, 3, 256, 256)) # 动态 batch 必须调用 stream torch.cuda.current_stream().cuda_stream d_input torch.empty((1, 3, 256, 256), dtypetorch.float32, devicecuda) d_output torch.empty((1, 1000), dtypetorch.float32, devicecuda) context.execute_async_v2([d_input.data_ptr(), d_output.data_ptr()], stream)这个脚本只要输出 shape 是 (1, 1000)引擎基本可用。它还有一个隐藏作用确认 engine 反序列化时能找到 AttentionPlugin 和 LayerNormPlugin。如果插件 .so 没加载或路径不对这一步会直接抛错而不是等到线上推理才暴露。4.3 用 benchmark 数据估算单卡并发路数T4 上 1080p 25fps 大概几路benchmark 脚本会输出单次推理的平均延迟。拿到延迟之后怎么估算路数这是很多人在选卡、定方案时纠结的点。T4 上跑 1080p 25fps 的视频流每路的帧间隔是 40ms。假设 MobileViT-S 在 FP16、batch1 下单路延迟约 8ms具体以你的 benchmark 实测为准理论并发就是 40 除以 8等于 5 路。但实际评估时每路还有解码、缩放、后处理我给一个保守公式可用路数 ≈ (1000 / fps) / 单帧GPU推理毫秒数 * 0.8场景单帧推理 ms1080p25 路数估算T4 MobileViT-S FP16 batch18示例以实测为准约 4~5T4 YOLO 640 FP16常见参照5~7示例约 5~6这里要提醒一个常见误区瓶颈不一定在 GPU。如果预处理是用 OpenCV 在 CPU 上做的很多路数上不去的案例是 CPU 被打满GPU 利用率反而没拉起来。测路数时要同时盯 CPU 占用和 GPU 利用率别只盯一个。另外MobileViT 的 attention 只支持训练时的分辨率测试时固定输入尺寸不要盲目把分辨率提到 512 或 640那会让 unfold 的 patch 数量变化plugin 就得重新适配。5. 避坑指南MobileViT 转 TensorRT 的常见问题与排查下面这几条都是实际调试时反复遇到的坑按“现象-原因-解决”写排查时直接对号入座。5.1 引擎构建成功但推理结果和 PyTorch 完全对不上现象engine 能加载输出形状也没问题但和 torch 输出对比cosine 只有 0.8 甚至更低分类结果明显不对。原因plugin 没有被真正执行或者 plugin 内部 permute 的顺序和 ONNX 图约定不一致。MobileViT 的 unfold 之后特征维度是 [B, N, C]但有些实现会变成 [B, C, N]softmax 作用的轴一旦错了输出就彻底歪掉。解决先跑 test_attention_plugin.py输入固定张量给 plugin把输出和 torch 的 attention 输出逐元素对齐。如果 plugin 单测不过问题在 .cu 里的 kernel 逻辑如果单测过了但整网不对去看 onnx_add_plugin.py 的输入输出接线是否正确。这个坑排起来很费时间建议第一次就把 plugin 单测写好别等整网跑完再猜。5.2 动态 batch 设置不当构建时报 shape 推断错误现象convert_to_trt.py 在 build 阶段抛 ONNX Parser 的 shape inference error或者 engine 构建成功但运行时set_input_shape报错。原因ONNX 的 dynamic_axes 只标了 batch 维但模型里某个算子把 batch 当作常量参与了计算常见于 unfold 和 reshape 的组合。ONNX Runtime 对动态 shape 容忍度高TensorRT 更严格。解决导出 ONNX 后先用 onnxruntime 的动态 batch 测一遍构建 TensorRT 时把 min/opt/max 三个 batch 写成 (1, 4, 8)宽高保持固定。如果还是推不出来直接去掉动态 batch固定 batch1 或 batch4。多数推理场景根本不需要动态 batch固定形状反而能触发更多融合优化。5.3 INT8 校准后精度暴跌现象FP16 一切正常一开 INT8top-1 掉 10 个点以上。原因校准集太小比如只有 50 张图或者校准集分布和真实场景不一致又或者 calibrator 返回的 batch 数不足。MobileViT 的 attention 对量化误差比纯卷积模型更敏感。解决从验证集抽 500 张左右类别分布均衡每次构建都用同一个 calib.cache保证前后可比如果还是掉对 attention 和 LayerNorm 插件相关层跳过量化在 TensorRT 里用 set_precision 按层做白名单控制。5.4 Ubuntu 上 import tensorrt 失败或版本与 CUDA 不匹配现象运行脚本时直接抛ImportError: libnvinfer.so.8: cannot open shared object file。原因最常见的是用 pip 装了 tensorrt但系统里没有对应版本的 CUDA或者之前用 deb 装过一套、pip 又装了一套两个版本的库互相覆盖。解决先卸载所有 tensorrt 相关包统一用一种安装方式。Ubuntu 上我一般先用dpkg -l | grep TensorRT看系统包再用python -c import tensorrt; print(tensorrt.__version__)验证 pip 版本是不是同一套。同时用trtexec --version确认命令行工具和 python 库的版本一致并核对 CUDA 的小版本跨度。TensorRT 对 CUDA minor 版本有要求差一两个小版本就可能出现库加载失败。5.5 显存占用异常高推理速度比预期慢不少现象FP16 构建成功batch1 延迟比 ONNX Runtime 还慢显存却占用很高。原因层融合没有生效。plugin 节点本身会阻断一部分融合但如果 ONNX 里还残留着一堆未合并的 ElementWise 和 Reduce 节点TensorRT 的优化就无从下手。这种情况常见于跳过 onnx_add_plugin.py 直接拿原图去 build。解决先跑 onnx_add_plugin.py 把子图替换干净再构建引擎看 verbose 日志里的层数和 kernel 数量确认图确实变小了。另外把 workspace 上限调大一点比如 2GB让融合算法有足够空间尝试但不要超过显存容量。6. 把精度验证做成回归测试让每次改图都能自动发现翻车部署项目里最容易烂尾的就是“跑一次没问题改个输入尺寸就翻车”。我习惯把精度验证固化成一个入口让它不仅能手动跑还能在每次改完 plugin 或转换脚本后自动执行。项目里有 Makefile最简单的做法就是把整条验证链串进一个 target# 一键跑完整验证造数据 - 冒烟 - torch 基线 - trt 对比 make test # 等价于 python gen_test_data.py --count 8 --output data/ python test_trt.py --engine mobilevit_s_fp16.engine --input data/ python test_torch_precision.py --data data/ --save torch_outputs.npz python test_trt_precision.py --engine mobilevit_s_fp16.engine --data data/test_torch_precision.py 先把 torch 输出存成 npztest_trt_precision.py 再加载同样的输入去对比。只要测试数据不变任何一端的改动都能被 diff 出来。对比脚本里最好把阈值判断写进去失败时退出非 0这样能直接挂 CI# check_precision.py 片段失败即退出非 0适合接 CI if cos 0.999 or max_diff 0.05: print(fFAIL: cosine{cos:.6f} max_diff{max_diff:.6f}) exit(1)用法上的一个硬性约定测试输入尺寸固定 256x256。MobileViT 的模型结构和插件都是按这个尺寸对齐的改尺寸必须重新导出 ONNX、重新改图、重新校准三条缺一不可否则验证结果没有参考价值。我最早调试这类模型时经常是引擎 build 完就丢上去跑等发现线上精度不对再回来查。后来一次 plugin 里 transpose 顺序写反torch 和 trt 的整体 cosine 能到 0.99 以上top-1 却始终不对最后靠逐层输出对比才发现是 patch 维的顺序问题。从那以后我每次改图都强制走一遍 torch/trt 双端对比阈值写进 Makefile先make test再谈性能。希望帮到你。本文还有配套的精品资源点击获取