使用Pydot和graphviz画TensorRT的Engine图

使用Pydot和graphviz画TensorRT的Engine图

TensorRT支持很多层的融合,目的是为了优化访存,减少数据在每层之间传输的消耗。

融之后的模型,一般无法通过Netron查看,毕竟TensorRT是闭源的,如果我们想看到融合后的模型长什么样,只要在build engine开启verbose模式即可。

复制代码
./trtexec --explicitBatch --onnx=debug.onnx --saveEngine=debug.trt  --verbose

[V] [TRT] Engine Layer Information:
Layer(Scale): QuantizeLinear_2_quantize_scale_node, Tactic: 0, input[Float(1,3,-17,-18)] -> 255[Int8(1,3,-17,-18)]
Layer(CaskConvolution): conv1.weight + QuantizeLinear_7_quantize_scale_node + Conv_9 + Relu_11, Tactic: 4438325421691896755, 255[Int8(1,3,-17,-18)] -> 267[Int8(1,64,-40,-44)]
Layer(CudaPooling): MaxPool_12, Tactic: -3, 267[Int8(1,64,-40,-44)] -> Reformatted Output Tensor 0 to MaxPool_12[Int8(1,64,-21,-24)]
Layer(Reformat): Reformatting CopyNode for Output Tensor 0 to MaxPool_12, Tactic: 0, Reformatted Output Tensor 0 to MaxPool_12[Int8(1,64,-21,-24)] -> 270[Int8(1,64,-21,-24)]
Layer(CaskConvolution): layer1.0.conv1.weight + QuantizeLinear_20_quantize_scale_node + Conv_22 + Relu_24, Tactic: 4871133328510103657, 270[Int8(1,64,-21,-24)] -> 284[Int8(1,64,-21,-24)]
Layer(CaskConvolution): layer1.0.conv2.weight + QuantizeLinear_32_quantize_scale_node + Conv_34 + Add_42 + Relu_43, Tactic: 4871133328510103657, 284[Int8(1,64,-21,-24)], 270[Int8(1,64,-21,-24)] -> 305[Int8(1,64,-21,-24)]

pyplot代码

https://github.com/pytorch/pytorch/pull/66431/files

engine_layer_visualize.py

这是jerryzh168大神开源的Facebook内部查看engine的工具,使用pydot和graphviz来画神经网络结构图

复制代码
pip install pydot
conda install python-graphviz

需要注意我们需要输入log_file也就是刚才开启Verbose的构建信息,然后profile_file则是使用TensorRT来profile的信息,最简单的可以通过trtexec这样获取到:

复制代码
./trtexec --loadEngine=debug_int8.trt --dumpProfile --shapes=input:1x3x512x512 --exportProfile=debug_profile

然后通过上述代码生成EngineLayers_0.dot

复制代码
import pydot

graphs = pydot.graph_from_dot_file("EngineLayers_0.dot")
graph = graphs[0]
graph.write_png("trt_engine.png")

TensorRT-engine 模型输入是Float而输出是Int8。这个模型是由TensorRT官方提供的pytorch-quantization[1]工具对Pytorch模型进行量化后导出ONNX,然后再由TensorRT-8转化得到的engine,这个engine的精度是INT8。

相关推荐
风铃喵游32 分钟前
让大模型调用MCP服务变得超级简单
前端·人工智能
旷世奇才李先生36 分钟前
Pillow 安装使用教程
深度学习·microsoft·pillow
booooooty1 小时前
基于Spring AI Alibaba的多智能体RAG应用
java·人工智能·spring·多智能体·rag·spring ai·ai alibaba
PyAIExplorer1 小时前
基于 OpenCV 的图像 ROI 切割实现
人工智能·opencv·计算机视觉
风口猪炒股指标1 小时前
技术分析、超短线打板模式与情绪周期理论,在市场共识的形成、分歧、瓦解过程中缘起性空的理解
人工智能·博弈论·群体博弈·人生哲学·自我引导觉醒
ai_xiaogui2 小时前
一键部署AI工具!用AIStarter快速安装ComfyUI与Stable Diffusion
人工智能·stable diffusion·部署ai工具·ai应用市场教程·sd快速部署·comfyui一键安装
聚客AI3 小时前
Embedding进化论:从Word2Vec到OpenAI三代模型技术跃迁
人工智能·llm·掘金·日新计划
weixin_387545643 小时前
深入解析 AI Gateway:新一代智能流量控制中枢
人工智能·gateway
聽雨2373 小时前
03每日简报20250705
人工智能·社交电子·娱乐·传媒·媒体
二川bro4 小时前
飞算智造JavaAI:智能编程革命——AI重构Java开发新范式
java·人工智能·重构