简介:面向视频理解算法在真实场景落地时的推理性能瓶颈,这份实战资源完整演示如何用 TensorRT 部署 SlowFast 视频理解模型,以解决其计算开销大、推理速度慢的问题。项目从模型转换入手,先将训练好的模型导出为 ONNX 中间格式,再借助 TensorRT 完成计算图优化、层融合、内核自动调优与精度校准,最终部署到 NVIDIA GPU 上,实现显著加速和低延迟推理,覆盖从模型导出到 TensorRT 部署的完整链路。资源包共 24 个文件,压缩后约 46KB,以 21 个 Python 脚本为主体,涵盖 ONNX 导出、TensorRT 转换、模型优化与推理脚本等关键模块,并配有 YAML 配置文件、README 工程说明与 .gitignore;脚本按转换、优化、推理拆分,另有 models、utils、configs 等目录便于按模块阅读。已有 158 人学习下载,适合具备深度学习基础、希望掌握从模型到产品完整流程的算法工程师和研究者,也可作为视频理解部署项目二次开发参考。
1. 用TensorRT部署SlowFast:视频理解算法从“能跑”到“能上线”
训练好的SlowFast视频理解模型,在PyTorch里跑起来只有十几个FPS,放到线上要同时处理几十路视频流,CPU扛不动,GPU又跑不满。这个标题的核心是把SlowFast从PyTorch搬到TensorRT上,用层融合、FP16/INT8量化把推理延迟压下去。这篇笔记面向做视频结构化、行为识别、内容审核的算法工程师和部署工程师,目标是一步步把pt权重变成可上线的TensorRT engine,并讲清楚参数怎么调、坑在哪里。适合手里已有训练好的SlowFast权重、正准备把它推上线的人,也适合刚接触算法部署、想找一个完整案例练手的人。
2. 部署前先摸清SlowFast的家底:两路分支、侧向连接与TensorRT的契合点
2.1 Slow路径与Fast路径:两路分支在推理时怎么协作
SlowFast不是单一网络,而是双分支结构。Slow路径以低帧率(通常8帧)采样,但空间分辨率保持较高,负责捕捉场景语义和稳定的动作类别;Fast路径以高帧率(通常32帧)采样,空间分辨率降为Slow路径的1/8,负责捕捉快速运动变化,比如挥手、转身、跌倒这类短促动作。两条路径的通道数也不同,Fast路径的通道数只有Slow路径的1/8,设计上就是轻量辅助。
推理时两路输入并行过网络,中间通过侧向连接(lateral connection)把Fast路径的时间特征融合进Slow路径。侧向连接在SlowFast里是一个卷积加转置拼接的操作,PyTorch里做的是torch.cat配合reshape,这一步在后面对接TensorRT时是第一个容易翻车的点。两路输入的shape不一致,意味着TensorRT构建engine时,两条输入张量的profile要分别设置。
import torch from slowfast.models import build_model from slowfast.config.defaults import get_cfg cfg = get_cfg() model = build_model(cfg) ckpt = torch.load("slowfast_8x8_r50.pyth", map_location="cpu") model.load_state_dict(ckpt["model_state_dict"]) model.eval() dummy_slow = torch.randn(1, 3, 8, 224, 224) dummy_fast = torch.randn(1, 3, 32, 224, 224) with torch.no_grad(): out = model([dummy_slow, dummy_fast]) print(out.shape)打印出来的out.shape就是最终分类分数,默认是[1, num_classes]。注意SlowFast的前向输入是一个列表[slow, fast],不是两个独立参数,这在与ONNX导出对接时要用tuple形式传入。模型的输入帧数由配置里的DATA.SAMPLING_RATE和DATA.NUM_FRAMES决定,8x8配置对应的就是slow 8帧、fast 32帧。
2.2 TensorRT对视频理解模型的三个优势:层融合、FP16/INT8量化、多流调度
视频理解模型部署,GPU利用率低是常态。SlowFast这种双分支结构,PyTorch推理时每个算子一次kernel launch,中间还有Python GIL和显存拷贝开销。TensorRT做的第一件事就是层融合,把Conv+BN+ReLU合并成一个kernel,把相邻的elementwise操作拉平,减少kernel launch次数。SlowFast的backbone是ResNet50,堆了大量BasicBlock,BN层特别多,融合收益比纯卷积网络更明显。
第二件事是精度校准。FP16模式直接把权重和激活降到半精度,一般掉点0.1%到0.5%,对动作分类这类任务基本无感。INT8模式需要提供校准数据集,SlowFast的输出是类别概率分布,对量化误差的敏感度中等,校准集选不好会出现个别类别概率偏移,部署时优先用FP16,INT8留到验收阶段再试。
第三件事是多流调度。视频理解场景很少单路推理,通常是多路视频流并发。TensorRT的execute_async_v2配合多个CUDA stream,可以让不同视频帧的推理在GPU上并行交错,把显存利用率拉满。这一点是PyTorch里需要自己写多线程、容易写出显存碎片的问题,TensorRT的显存池机制在engine构建时就把优化计划定好了。
这里插一句版本选择的事:如果你手里的卡是GTX 1070这类Pascal架构老卡,TensorRT版本选不对,后面每一步都膈应。常见组合是CUDA 11.8 + cuDNN 8.9 + TensorRT 8.6 GA,这个组合在Pascal、Volta、Ampere上都有覆盖,踩坑少。TensorRT 10.x换了CUDA 12.x底座,驱动要求高,老卡虽然理论支持,但有些kernel选型会回退到通用实现,性能反而打折。版本匹配问题第三章详细讲。
3. 构建部署环境:TensorRT版本选择、CUDA匹配与依赖清单
3.1 TensorRT 8.x/10.x与CUDA、GPU的兼容矩阵:GTX 1070这类老卡到底能不能用
“TensorRT 10.x是否支持GTX 1070”这个问题被问得很多。GTX 1070的Compute Capability是6.1,属于Pascal架构。TensorRT 10.x最低支持CC 6.0,所以理论上是能用的,但实际部署不建议一上来就上10.x。原因有三:驱动要求高,10.x配套CUDA 12.x,老卡驱动需要升级到545+,生产环境升级驱动是个牵扯面很大的事;kernel覆盖有空洞,Pascal架构缺少很多新指令,10.x的算子库对老架构的优化投入明显不如8.x时期;TensorRT 10.x把工作区参数从--workspace改成了--memPoolSize,网上大量旧教程直接抄会报参数错误。
这块兼容关系用一张表概括,按生产环境稳妥程度排序:
| TensorRT版本 | CUDA版本 | cuDNN版本 | 推荐的GPU架构 | GTX 1070实测状态 |
|---|---|---|---|---|
| TensorRT 8.2 | CUDA 11.2 | cuDNN 8.2 | Pascal/Volta/Turing | 稳定,支持完整 |
| TensorRT 8.6 GA | CUDA 11.8 | cuDNN 8.9 | Pascal/Ampere全系 | 最稳,生产首选 |
| TensorRT 9.x | CUDA 12.0+ | cuDNN 8.9 | Ampere/Ada | Pascal部分算子回退 |
| TensorRT 10.x | CUDA 12.3+ | cuDNN 9.x | Ada/Ampere/Hopper | 可用但不推荐 |
如果你的生产环境已经锁定了TensorRT 10.x,GTX 1070能跑,但建议先构建engine跑一遍--dumpProfile,看有没有kernel选型回退的警告。我从经验出发的建议是:代码和模型结构不变的情况下,先用8.6 GA跑通全流程,再考虑升10.x。配置化部署的项目最怕环境升级带来一长串连锁反应,8.6 GA的算子支持对SlowFast这个模型已经完全够用。
3.2 从零装一套可复现的TensorRT部署环境:安装步骤与校验命令
TensorRT不能只靠pip安装,正确方式是下载对应CUDA版本的安装包,然后手动解压配置。装完之后Python包需要单独处理,常见的翻车点是import tensorrt成功但版本号对不上,或者trtexec二进制找不到,这通常是装了两个版本的TensorRT,LD路径串了。
# 以TensorRT 8.6 GA + CUDA 11.8为例,安装包解压到/opt cd /opt tar -xzvf TensorRT-8.6.1.6.Linux.x86_64-gnu.cuda-11.8.tar.gz # 配置环境变量,写进~/.bashrc export TRT_HOME=/opt/TensorRT-8.6.1.6 export LD_LIBRARY_PATH=$TRT_HOME/lib:$LD_LIBRARY_PATH export PATH=$TRT_HOME/bin:$PATH # Python包安装 pip install $TRT_HOME/python/tensorrt-8.6.1.6-cp38-none-linux_x86_64.whl # 校验:三行命令全部通过才算装好 python -c "import tensorrt as trt; print(trt.__version__)" trtexec --version ldconfig -p | grep nvinferldconfig那一步容易栽跟头,如果系统里之前装过TensorRT,grep nvinfer会搜出多个路径,Python import时加载到哪个库是随机的。解决办法是把$TRT_HOME/lib放到LD_LIBRARY_PATH最前面,并移除旧版本的动态库路径。还有一处要注意:TensorRT的Python wheel只对应特定Python版本,cp38就是Python 3.8,如果生产环境是3.10,要去python/目录找对应的cp310包,或者用pip在线安装但指定--extra-index-url指向NVIDIA的PyPI索引。
环境装完先别急着跑模型,拿一个简单的ONNX文件构建engine验证整条链路通不通。可以用ONNX官方仓库里的resnet50模型,trtexec --onnx=resnet50.onnx --fp16能成功生成engine文件,就说明CUDA、cuDNN、TensorRT三者版本是匹配的。这一步能过滤掉80%的环境类问题,后面再踩坑就基本都是模型本身的了。
3.3 PyTorch与TensorRT的推理环境隔离:一个容易被忽略的实践
部署实践中,PyTorch的版本依赖经常和TensorRT冲突。比如PyTorch 2.x默认自带CUDA 12 runtime,会覆盖系统CUDA,导致TensorRT加载时找不到cuDNN的符号。常见做法是训练和部署环境分开,训练环境留在原机器,部署机上只装Python基础环境加TensorRT,用ONNX作为模型交换格式。这样副作用最小。
隔离部署环境的具体做法是单独建一个conda环境,Python版本固定3.8或3.10,只装numpy、pycuda、tensorrt这几个包。pycuda的安装也常出问题,它需要系统里有gcc和python头文件,pip install pycuda报编译错误时先安装python3-dev。这个环境里不要装torch,导出ONNX时用另一台机器或另一个环境。如果没有条件分开,那至少确保import tensorrt之前不加载torch的CUDA库,否则两个runtime的上下文管理会发生碰撞。
4. 把PyTorch的SlowFast转到TensorRT:pt → ONNX → engine 的完整链路
4.1 第一步:PyTorch模型导出ONNX,动态轴与输入预处理对齐
导出ONNX是整个流程里最容易出hidden bug的一步。SlowFast的模型结构里,常见的坑是torch.cat在时间维度上的拼接、reshape操作、以及lateral connection里的自定义卷积。这些操作在PyTorch里正常,导出ONNX时不一定能映射到标准算子。
import torch from slowfast.models import build_model from slowfast.config.defaults import get_cfg cfg = get_cfg() model = build_model(cfg) ckpt = torch.load("slowfast_8x8_r50.pyth", map_location="cpu") model.load_state_dict(ckpt["model_state_dict"]) model.eval() # 注意:这里是tuple传参,不是list dummy_slow = torch.randn(1, 3, 8, 224, 224) dummy_fast = torch.randn(1, 3, 32, 224, 224) torch.onnx.export( model, (dummy_slow, dummy_fast), "slowfast_r50.onnx", opset_version=13, input_names=["slow_input", "fast_input"], output_names=["prediction"], dynamic_axes={ "slow_input": {0: "batch", 2: "slow_time"}, "fast_input": {0: "batch", 2: "fast_time"}, }, do_constant_folding=True, )导出时input_names和dynamic_axes是关键。SlowFast两个输入的batch维度必须都设动态,时间维度建议也设动态,实际线上推理时视频帧数是变化的,clip长度固定也是靠前处理pad出来的。空间维度224x224建议固定,不要设成动态,视频理解模型的高宽变化对TensorRT优化伤害很大,常见做法是前处理统一resize到224。
导完先用onnxsim精简一次,能去掉一堆冗余的Identity和Cast节点:
python -m onnxsim slowfast_r50.onnx slowfast_r50_sim.onnximport onnx model = onnx.load("slowfast_r50_sim.onnx") onnx.checker.check_model(model) # 打印输入输出确认动态轴生效 print(onnx.helper.printable_graph(model.graph))第一次跑ONNX导出时用opset_version=13最稳,新版PyTorch如果提示opset过低可以升到14或15,但onnxsim和TensorRT对旧版opset的解析更成熟,能不动就不动。导出成功后先拿onnxruntime在CPU上跑一遍验证输出shape和概率分布,再进trtexec,一个问题一个坑排查。
4.2 第二步:用trtexec构建engine,FP16与动态Shape参数怎么设
trtexec是TensorRT自带的可执行工具,推荐先用它跑通构建流程,因为它把profiling和报错信息都打得很清楚。构建命令行里--minShapes、--optShapes、--maxShapes三者必须成对出现,顺序也不能乱。
trtexec \ --onnx=slowfast_r50_sim.onnx \ --saveEngine=slowfast_r50_fp16.engine \ --fp16 \ --minShapes=slow_input:1x3x8x224x224,fast_input:1x3x32x224x224 \ --optShapes=slow_input:8x3x8x224x224,fast_input:8x3x32x224x224 \ --maxShapes=slow_input:16x3x8x224x224,fast_input:16x3x32x224x224 \ --workspace=4096 \ --verbose这里--minShapes的格式是张量名:维度,多个输入用逗号分隔。optShapes的选择很有讲究,它不是设最小值最大值就行,而是你要估算线上最常见的batch和帧数组合。比如线上多是8路视频并发,那就把opt设为8。TensorRT优化器会针对opt shape挑选kernel,实际推理shape偏离opt越远,性能越差。
--workspace=4096是GPU显存上限,不是固定分配。SlowFast模型的中间激活远大于同类分类模型,因为fast路径的32帧输入会产生巨大的时间维激活值,显存给到4GB到6GB比较稳妥。如果构建时OOM,优先调低batch上限,不要调低workspace。--verbose构建费时间但值得,构建失败时没有它只剩下一个笼统错误,开了能看到具体卡在哪个节点。
构建成功的标志是最后出现Engine built in xx seconds。生成的engine文件是二进制序列化产物,不区分源模型,但强制绑定构建时的TensorRT版本和GPU架构。也就是说在A100上构建的engine不能拿到GTX 1070上跑,跨机器部署要在目标机器上重新构建。
4.3 第三步:Python加载engine做推理,输出与原始模型对齐校验
engine文件拿到后,写推理代码主要做三件事:反序列化engine、分配host/device显存、执行推理。
import tensorrt as trt import numpy as np import pycuda.driver as cuda import pycuda.autoinit class SlowFastEngine: def __init__(self, engine_path, num_classes): self.logger = trt.Logger(trt.Logger.WARNING) with open(engine_path, "rb") as f: with trt.Runtime(self.logger) as runtime: self.engine = runtime.deserialize_cuda_engine(f.read()) self.context = self.engine.create_execution_context() self.num_classes = num_classes self.stream = cuda.Stream() self._alloc_buffers() def _alloc_buffers(self): self.inputs = {} self.outputs = {} self.bindings = [] for i in range(self.engine.num_bindings): name = self.engine.get_binding_name(i) dtype = trt.nptype(self.engine.get_binding_dtype(name)) shape = self.engine.get_binding_shape(i) if shape[0] == -1: shape = (1, *shape[1:]) host = cuda.pagelocked_empty(trt.volume(shape), dtype) device = cuda.mem_alloc(host.nbytes) self.bindings.append(device) if self.engine.binding_is_input(name): self.inputs[name] = {"host": host, "device": device} else: self.outputs[name] = {"host": host, "device": device} def infer(self, slow, fast): # slow: [N,3,8,224,224], fast: [N,3,32,224,224] self.context.set_binding_shape(0, slow.shape) self.context.set_binding_shape(1, fast.shape) np.copyto(self.inputs["slow_input"]["host"], slow.ravel()) np.copyto(self.inputs["fast_input"]["host"], fast.ravel()) for name, buf in self.inputs.items(): cuda.memcpy_htod_async(buf["device"], buf["host"], self.stream) self.context.execute_async_v2(self.bindings, self.stream.handle) for name, buf in self.outputs.items(): cuda.memcpy_dtoh_async(buf["host"], buf["device"], self.stream) self.stream.synchronize() return self.outputs["prediction"]["host"].reshape(-1, self.num_classes)逻辑说明:_alloc_buffers阶段,动态shape的binding和静态shape不同,需要先获取默认shape再分配显存。因为SlowFast的batch和time维度是动态的,第一次推理前必须调用set_binding_shape重新指定实际shape,否则execute_async_v2会用构建时opt shape执行,输出shape和预期不一致。
执行链路是host到device拷贝、GPU推理、device回host拷贝三段。这里用execute_async_v2而不是execute_v2,本质区别是异步执行不阻塞CPU,帧与帧之间可以流水线重叠。np.copyto处理的是分页锁内存,显存拷贝走高速通道,避免普通内存的page fault影响吞吐。
完成加载后需要做一次精度对齐,把同一批测试帧分别喂给PyTorch原模型和TensorRT engine,对比softmax后的输出。常见的包装方式是把TensorRT推理包装成model(x)接口,这样验证脚本可以复用。精度对比踩坑放在第五章,这里先留个心眼:对比时输入必须完全一致,包括归一化参数、crop方式、通道顺序,一个像素的偏差都会放大成概率分布差异。
5. 部署SlowFast时的5个高频坑:从NaN输出到显存溢出
5.1 FP16推理输出全NaN,概率分布是乱的
现象:engine构建成功,加载无报错,但推理结果全是NaN或者softmax后概率分布明显错乱。
原因:SlowFast的backbone最后接的全连接层,在FP16下权重范围较大,分类头在类别数多时容易出现激活值溢出。Fast路径的侧向连接在低精度下也会放大误差。
解决:让全连接层保持FP32计算。TensorRT 8.6里可以指定--layerPrecisions,但要精确定位层名比较麻烦。我一般直接改导出脚本,单独控制torch的_modules中final_fc层的导出精度。另一个更快的方案是在trtexec构建时不加--fp16,先确认整条链路在FP32下输出正确,再开FP16,隔离问题。如果确实是最终分类头的问题,就把整个模型降成FP16+FP32混合,做法是--fp16 --layerPrecisions=.*:fp32试出哪个层有问题再精准回退。
5.2 构建engine报Unsupported op,卡在某个ConvTranspose或Resize节点
现象:trtexec跑到中段报错,提示Unsupported operation或Assertion failed,日志指向某个特定的ONNX节点。
原因:SlowFast的lateral connection里用了torch.cat加reshape的组合,导出ONNX后变成了非标准的Concat和Gather节点组合,TensorRT解析器版本对某些opset组合解析不完整。
解决:把opset_version从13升到14或降回12都试一下,通常能覆盖解析器支持的边界。再用onnxsim简化。如果还不行,就用onnx库手动改图,把那个不支持的子图替换成标准的Conv3d + Concat。还有一个偏方是给torch.onnx.export传入operator_export_type=torch.onnx.OperatorExportTypes.ONNX_ATEN_FALLBACK,导出时把ATen算子也一起导出,TensorRT对ATen算子的兼容度反而更高。
5.3 GTX 1070上构建慢,且构建时显存爆掉
现象:同样的ONNX在RTX 3090上几分钟构建完,换到GTX 1070上半小时还没结束,甚至中途报OutOfMemory。
原因:Pascal架构没有Tensor Core,TensorRT构建时无法使用FP16的Tensor Core优化路径,只能走CUDA core通用实现,kernel选择变少,构建策略搜索时间变长。同时动态shape的优化空间爆炸,显存占用被拉高。
解决:构建时关掉动态shape的time维度,固定为8帧和32帧。线上推理不需要每帧长度都变,可以把输入视频前处理成固定长度clip。--minShapes和--maxShapes设置相同时,TensorRT会退化成static shape构建,构建速度快一个数量级。如果必须动态,减小batch的opt值,比如从16降到8,--workspace相应调回2048。
5.4 FP16推理比FP32还慢,延迟不降反升
现象:对比基准测试发现,开FP16后FPS没提升,反而下降20%。
原因:输入的T维度太大。SlowFast的fast路径32帧,batch为8时输入张量是[8, 3, 32, 224, 224],中间激活膨胀得厉害。小batch下FP16 kernel的启动和调度开销占比高,Tensor Core吞吐没吃满,带宽瓶颈突出。
解决:增大batch到16或32,让FP16的吞吐优势显现出来。或者换一个思路,在batch=1的实时推理场景,FP16收益很小,直接保持FP32更稳。这个问题在TensorRT里可以通过--streams=2开启多流并行,让两个不同请求的推理在同一个context里流水线交错,GPU利用率上来后FP16优势才能兑现。
5.5 多路视频流推理时延迟抖动,p99飙升
现象:单路推理延迟稳定30ms,并发8路后平均延迟才40ms,但p99飙到120ms。
原因:动态shape的engine在每路视频帧shape变化时触发context重配置。比如一路视频当前帧是32帧,另一路是16帧,TensorRT的context切换到新shape时要重新绑定内存和调整kernel调度,这个开销比理论值高很多。另外前处理在CPU上做resize和归一化,GPU空闲等待CPU拷帧。
解决:把所有输入视频clip统一pad到固定长度,强行走static shape,完全消除context重配置开销。前处理放到GPU上,用CUDA的torchvision.transforms结合cudastream实现,或者用预处理线程池把CPU的帧解码和GPU推理流水线重叠。部署架构上让每个GPU上跑的推理流数量固定,不要动态增减。
6. 性能调优与验证:把FPS从十几拉到一百的最后一公里
模型在TensorRT上跑通只是第一步,真正上线前还要做三件调优:预处理搬上GPU、推理流水线双stream化、精度回归验证。最常见的前处理瓶颈是视频解码后的resize和归一化在CPU上做,帧从内存拷到显存本身就占掉一次全帧拷贝的带宽。用CUDA的torch.Tensor.to(device, non_blocking=True)把归一化之后的张量直接放到GPU,配合cuda stream让拷贝和计算重叠,单路推理的端到端延迟能再降15%。SlowFast因为有两路输入,slow和fast的预处理可以放到两个不同的cuda stream并行执行,比单stream串行执行少一次等待。
精度回归验证不能只看top-1是否一致。建议准备一个固定测试集,里面包含至少100个clip,覆盖不同光照和运动幅度。跑完之后对比PyTorch原模型和TensorRT的softmax输出,计算KL散度,阈值为0.01。散度超标的clip要单独看,大概率是特定类别在FP16下激活溢出。这个验证要固化成脚本,每次engine重构建后都跑一遍,不要靠肉眼感觉“看起来差不多”。上线后监控指标建议同时看p50和p99延迟,p99延迟比平均延迟更能反映抖动问题。
最后记录一个我之前翻车的教训:第一次部署SlowFast时,我把所有层都开FP16,推理结果在大多数类别上正常,但“跑步”和“快走”这两个相似动作的概率分布颠倒了。当时花了很长时间排查数据预处理,最后才发现是分类头的权重范围太大,FP16下低比特截断导致细微特征丢失。之后我把final_fc层强制保留FP32,这个问题就再没出现过。这个经验固化成了部署规范:凡是最后接全连接分类头的模型,分类头一律FP32,其他层才考虑FP16。希望帮到你。
本文还有配套的精品资源,点击获取