适用场景：你已经有一个 YOLO 检测模型，希望在 NVIDIA Jetson 上实时显示检测结果，同时把模型关注区域以热力图方式融合到原始视频里，看模型到底在关注哪里。
请添加图片描述

请添加图片描述
本教程从零开始讲清楚一件事：为什么普通 PyTorch Hook 热力图会卡，为什么只用普通 TensorRT 检测引擎又拿不到特征图，以及如何把 YOLO 的 Detect 层输入特征 P2/P3/P4/P5 一起导出进 ONNX 和 TensorRT，让一次 TensorRT 推理同时产出：

  • output0：YOLO 检测输出，用来画检测框。
  • p2：高分辨率浅层特征，通常是细节和边缘响应。
  • p3：中等尺度特征。
  • p4：更抽象的中低分辨率特征。
  • p5：低分辨率高语义特征，通常更集中关注目标主体。

最终效果：

  • 左侧：实时可见光检测视频。
  • 右侧：四个尺度的“原视频 + 热力图”融合画面，2x2 展示。
  • 整条链路主要走 TensorRT，不再每帧用 PyTorch CPU 跑 Hook。

1. 为什么一开始会卡

很多热力图教程会这样做：

  1. 用 YOLO("xxx.pt") 加载 PyTorch 模型。
  2. 给 Detect 层注册 hook。
  3. 每帧把图像送进 PyTorch。
  4. 从 hook 里取出 P2/P3/P4/P5 特征图。
  5. 把特征图转成热力图。

这个思路在电脑上做离线图片可视化没问题，但在 Jetson 上实时视频会很卡，尤其是：

  • 每帧 PyTorch 前向计算开销大。
  • 如果 PyTorch CUDA 环境不完整，可能退到 CPU。
  • Hook 取出的特征还要从 GPU/CPU 转 NumPy，再 resize、applyColorMap、融合。
  • Jetson 如果供电不足，还会降频，例如出现 System throttled due to Over-current。

所以真正实时的做法不是“每帧 PyTorch Hook”，而是：

把 Hook 想要的中间特征作为模型正式输出导出到 ONNX，再构建成多输出 TensorRT engine。

这样一次 TensorRT 推理就能同时拿到检测结果和四个尺度特征图。

2. 本项目文件说明

本文档所在目录：

C:\e\workspace\2026世校赛\Quantification\step_jeston_RepC3_0520

核心文件：

RepC3_chip3.pt
原始 PyTorch YOLO 模型。

export_repC3_attention_onnx.py
把 YOLO 模型导出成多输出 ONNX：output0 + p2 + p3 + p4 + p5。

jetson_fullscreen_attention.py
实时显示脚本：TensorRT 多输出推理、检测框绘制、四尺度热力图融合、全屏展示。

jetson_display_detect.py
单纯检测显示脚本。

Jetson 远端工作目录：

/home/nvidia/yolo_deploy/ReqC3_chip3

最终生成的关键文件：

RepC3_chip3_attention.onnx
RepC3_chip3_attention_fp16.trt

3. 总体流程

整个流程分成 4 步：

第 1 步：确认 PyTorch 模型能取到 P2/P3/P4/P5
第 2 步：导出多输出 ONNX
第 3 步：用 trtexec 构建多输出 TensorRT FP16 engine
第 4 步：实时读取摄像头，一次 TensorRT 推理同时显示检测和热力图融合

传统做法：

摄像头帧 -> TensorRT 检测
摄像头帧 -> PyTorch Hook 热力图
两个模型/两条推理链路 -> 慢

优化后：

摄像头帧 -> 一个 TensorRT 多输出 engine -> 检测框 + P2/P3/P4/P5 -> 实时显示

4. 什么是 P2/P3/P4/P5

YOLO 检测头通常不是只看一张特征图，而是看多个尺度。

以输入 1024x1024 为例：

P2: 256x256, stride=4
P3: 128x128, stride=8
P4: 64x64,   stride=16
P5: 32x32,   stride=32

含义：

  • P2 分辨率最高，适合看边缘、纹理、细小瑕疵。
  • P3 兼顾细节和语义。
  • P4 更关注中等目标结构。
  • P5 分辨率最低，但语义最强，常常集中在目标主体区域。

热力图做法：

activation = abs(feature).mean(axis=0)

也就是对一个特征图的所有通道求绝对值，再沿通道平均，得到一张二维响应图。

5. 为什么普通 TensorRT 检测模型不够

普通 YOLO TensorRT engine 只输出检测结果，例如：

output0: (1, 8, 87040)

这个输出已经是检测头处理后的结果，里面主要是框和类别分数，不再保留真正的中间特征图。

我们要的热力图来自 Detect 层输入特征：

p2: (1, 32, 256, 256)
p3: (1, 64, 128, 128)
p4: (1, 128, 64, 64)
p5: (1, 256, 32, 32)

所以必须重新导出一个“多输出模型”。

6. 导出多输出 ONNX 的完整代码

文件名：

export_repC3_attention_onnx.py

完整代码如下：

import argparse
import importlib.metadata
import os
import sys

import torch


ONNX_EXTRA_PATH = "/home/nvidia/miniforge3/envs/yolo11m_seg_quant/lib/python3.10/site-packages"


def ensure_torchvision_metadata():
    root = os.path.join(os.path.dirname(os.path.abspath(__file__)), "_fake_torchvision_metadata")
    dist = os.path.join(root, "torchvision-0.21.0.dist-info")
    package = os.path.join(root, "torchvision")
    os.makedirs(dist, exist_ok=True)
    os.makedirs(package, exist_ok=True)
    with open(os.path.join(dist, "METADATA"), "w") as f:
        f.write("Metadata-Version: 2.1\nName: torchvision\nVersion: 0.21.0\n")
    with open(os.path.join(package, "__init__.py"), "w") as f:
        f.write("__version__ = '0.21.0'\n")
    sys.path.insert(0, root)


class DetectFeatureExport(torch.nn.Module):
    def __init__(self, yolo_model):
        super().__init__()
        self.model = yolo_model.model
        self.layers = self.model.model
        self.save = set(int(i) for i in getattr(self.model, "save", []))

    def forward(self, x):
        y = []
        detect_inputs = None
        for layer in self.layers:
            if layer.f != -1:
                if isinstance(layer.f, int):
                    x_in = y[layer.f]
                else:
                    x_in = [x if j == -1 else y[j] for j in layer.f]
            else:
                x_in = x

            if layer is self.layers[-1]:
                detect_inputs = x_in
            x = layer(x_in)
            y.append(x if layer.i in self.save else None)

        if isinstance(x, (list, tuple)):
            detect = x[0]
        else:
            detect = x
        return (detect, detect_inputs[0], detect_inputs[1], detect_inputs[2], detect_inputs[3])


def main():
    parser = argparse.ArgumentParser()
    parser.add_argument("--model", default="RepC3_chip3.pt")
    parser.add_argument("--output", default="RepC3_chip3_attention.onnx")
    parser.add_argument("--imgsz", type=int, default=1024)
    parser.add_argument("--device", default="cpu")
    args = parser.parse_args()

    if os.path.isdir(ONNX_EXTRA_PATH) and ONNX_EXTRA_PATH not in sys.path:
        sys.path.append(ONNX_EXTRA_PATH)

    ensure_torchvision_metadata()
    from ultralytics import YOLO

    yolo = YOLO(args.model)
    yolo.model.eval()
    wrapped = DetectFeatureExport(yolo).to(args.device).eval()
    dummy = torch.zeros(1, 3, args.imgsz, args.imgsz, device=args.device)

    with torch.no_grad():
        outputs = wrapped(dummy)
        for name, output in zip(["output0", "p2", "p3", "p4", "p5"], outputs):
            print(name, tuple(output.shape))

    torch.onnx.export(
        wrapped,
        dummy,
        args.output,
        input_names=["images"],
        output_names=["output0", "p2", "p3", "p4", "p5"],
        opset_version=17,
        do_constant_folding=True,
    )
    print("saved:", args.output)


if __name__ == "__main__":
    main()

7. 这段导出代码的核心原理

重点是这个包装类：

class DetectFeatureExport(torch.nn.Module):

它不是重新写一个模型，而是复用 Ultralytics 已经加载好的模型层。

关键逻辑：

if layer is self.layers[-1]:
    detect_inputs = x_in

最后一层就是 Detect 层。Detect 层的输入通常就是 [P2, P3, P4, P5]。

最终返回：

return (detect, detect_inputs[0], detect_inputs[1], detect_inputs[2], detect_inputs[3])

也就是：

output0: 检测输出
p2: Detect 输入第 1 个尺度
p3: Detect 输入第 2 个尺度
p4: Detect 输入第 3 个尺度
p5: Detect 输入第 4 个尺度

8. 在 Jetson 上导出 ONNX

进入远端目录：

cd /home/nvidia/yolo_deploy/ReqC3_chip3

运行：

export LD_PRELOAD=/home/nvidia/miniforge3/envs/yolo_trt/lib/libstdc++.so.6
/home/nvidia/miniforge3/envs/yolo_trt/bin/python3 ./export_repC3_attention_onnx.py \
  --model RepC3_chip3.pt \
  --output RepC3_chip3_attention.onnx \
  --imgsz 1024 \
  --device cpu

成功时会看到类似：

output0 (1, 8, 87040)
p2 (1, 32, 256, 256)
p3 (1, 64, 128, 128)
p4 (1, 128, 64, 64)
p5 (1, 256, 32, 32)
saved: RepC3_chip3_attention.onnx

9. 用 trtexec 构建 TensorRT FP16 engine

运行：

/usr/src/tensorrt/bin/trtexec \
  --onnx=RepC3_chip3_attention.onnx \
  --saveEngine=RepC3_chip3_attention_fp16.trt \
  --fp16 \
  --memPoolSize=workspace:2048 \
  --verbose \
  > build_attention_trt.log 2>&1

查看结果：

ls -lh RepC3_chip3_attention_fp16.trt
tail -120 build_attention_trt.log

成功时会看到：

Output binding for output0 with dimensions 1x8x87040 is created.
Output binding for p2 with dimensions 1x32x256x256 is created.
Output binding for p3 with dimensions 1x64x128x128 is created.
Output binding for p4 with dimensions 1x128x64x64 is created.
Output binding for p5 with dimensions 1x256x32x32 is created.

性能参考：

Throughput: about 50 qps
GPU Compute Time mean: about 19.5 ms

这说明多输出特征图不会像 PyTorch CPU Hook 那样卡。

10. 实时显示脚本的核心思路

实时脚本：

jetson_fullscreen_attention.py

它做了 5 件事：

1. OpenCV 读取摄像头。
2. TensorRT 加载 RepC3_chip3_attention_fp16.trt。
3. 一次推理得到 output0 + p2 + p3 + p4 + p5。
4. output0 解码检测框。
5. p2-p5 生成热力图，并和原视频融合。

显示布局：

左侧：检测结果视频
右侧：2x2 四尺度热力图融合视频

11. TensorRT 多输出推理代码结构

核心类是：

class TrtRunner(object):

它会自动识别 TensorRT 的输入输出 tensor：

for i in range(self.engine.num_io_tensors):
    name = self.engine.get_tensor_name(i)
    mode = self.engine.get_tensor_mode(name)

对于多输出 engine，输出顺序是：

outputs[0] -> output0
outputs[1] -> p2
outputs[2] -> p3
outputs[3] -> p4
outputs[4] -> p5

实时显示里这样取：

trt_features = [output[0] for output in outputs[1:5]]

12. 检测框解码

检测输出形状：

(1, 8, 87040)

其中前 4 个值一般是框信息，后面是类别分数。本项目里类别为：

DEFAULT_CLASSES = ["ZF-scratch", "scratch", "broken", "pinbreak"]

解码流程：

output0 -> 转置为每行一个候选框 -> 取最大类别分数 -> 阈值过滤 -> NMS -> 画框

NMS 使用：

cv2.dnn.NMSBoxes(...)

13. 热力图如何生成

每个特征图形状类似：

p2: (1, 32, 256, 256)

去掉 batch 后：

(32, 256, 256)

对通道求平均响应：

activation = np.mean(np.abs(feature), axis=0)

再归一化：

activation -= activation.min()
activation /= activation.max()

最后变成 OpenCV 彩色热力图：

color = cv2.applyColorMap(heat, cv2.COLORMAP_JET)

14. 为什么要和原视频融合

纯热力图虽然能看响应强弱，但专家不容易知道热区对应芯片上的哪个位置。

融合显示更直观：

fused = cv2.addWeighted(frame, 1.0 - overlay_alpha, color, overlay_alpha, 0)

例如：

overlay_alpha = 0.45

意思是：

55% 原图 + 45% 热力图

这样既能看到芯片本身，又能看到模型关注区域。

15. 运行实时融合界面

在 Jetson 桌面会话中，显示号通常是 :1：

cd /home/nvidia/yolo_deploy/ReqC3_chip3

export DISPLAY=:1
export XAUTHORITY=/run/user/1000/gdm/Xauthority
export LD_PRELOAD=/home/nvidia/miniforge3/envs/yolo_trt/lib/libstdc++.so.6

/home/nvidia/miniforge3/envs/yolo_trt/bin/python3 ./jetson_fullscreen_attention.py \
  --engine /home/nvidia/yolo_deploy/ReqC3_chip3/RepC3_chip3_attention_fp16.trt \
  --feature-source trt \
  --screen-width 1920 \
  --screen-height 1080 \
  --left-ratio 0.50 \
  --heat-alpha 0.45

后台启动：

nohup /home/nvidia/miniforge3/envs/yolo_trt/bin/python3 ./jetson_fullscreen_attention.py \
  --engine /home/nvidia/yolo_deploy/ReqC3_chip3/RepC3_chip3_attention_fp16.trt \
  --feature-source trt \
  --screen-width 1920 \
  --screen-height 1080 \
  --left-ratio 0.50 \
  --heat-alpha 0.45 \
  > fullscreen_attention.log 2>&1 &
echo $! > fullscreen_attention.pid

停止：

kill $(cat /home/nvidia/yolo_deploy/ReqC3_chip3/fullscreen_attention.pid)

16. 快捷键

程序窗口内支持：

q
x
Esc

任意一个都可以退出。

17. 参数解释

--engine

指定 TensorRT engine。实时热力图必须使用多输出 engine：

RepC3_chip3_attention_fp16.trt

--feature-source trt

强制使用 TensorRT 输出的 P2-P5 特征。如果 engine 没有这些输出，程序会直接报错，避免悄悄退回慢速 PyTorch。

--heat-alpha 0.45

热力图融合比例。越大热力图越明显，越小原图越明显。

推荐范围：

0.35 - 0.55

--left-ratio 0.50

左侧检测画面占屏幕宽度比例。

--screen-width 1920 --screen-height 1080

全屏画布大小。如果显示器不是 1080p，可以改成实际分辨率。

18. 验证多输出 engine

运行：

cd /home/nvidia/yolo_deploy/ReqC3_chip3
export LD_PRELOAD=/home/nvidia/miniforge3/envs/yolo_trt/lib/libstdc++.so.6

/home/nvidia/miniforge3/envs/yolo_trt/bin/python3 ./jetson_fullscreen_attention.py \
  --engine /home/nvidia/yolo_deploy/ReqC3_chip3/RepC3_chip3_attention_fp16.trt \
  --feature-source trt \
  --check-feature

成功输出：

P2 TRT output shape: (1, 32, 256, 256)
P3 TRT output shape: (1, 64, 128, 128)
P4 TRT output shape: (1, 128, 64, 64)
P5 TRT output shape: (1, 256, 32, 32)

19. 常见问题

19.1 右侧热力图是空的

原因通常是用了旧 engine：

ReqC3_chip3_fp16.trt

它只有：

output0

没有：

p2/p3/p4/p5

解决：

--engine /home/nvidia/yolo_deploy/ReqC3_chip3/RepC3_chip3_attention_fp16.trt
--feature-source trt

19.2 画面卡

先确认是不是跑了 PyTorch：

ps -ef | grep jetson_fullscreen_attention.py

如果命令里有：

--feature-source pytorch

或者 engine 不是多输出 engine，就可能退回慢路径。

正确实时路径应该是：

--feature-source trt
RepC3_chip3_attention_fp16.trt

19.3 出现过流降频提示

Jetson 弹出：

System throttled due to Over-current

说明供电不足或功耗模式不稳定。建议：

1. 使用足够功率的电源。
2. 避免同时构建 TensorRT engine 和跑实时显示。
3. 构建 engine 时先停止显示进程。

19.4 OpenCV 报 Qt 字体警告

类似：

QFontDatabase: Cannot find font directory ...

一般不影响显示，可以忽略。

20. 性能优化建议

优先级从高到低：

1. 必须使用多输出 TensorRT engine。
2. 不要实时跑 PyTorch Hook。
3. 尽量一次推理同时拿检测和特征。
4. 热力图只做必要的 NumPy 处理。
5. 供电必须稳定，避免 Jetson 降频。

如果还要继续优化，可以考虑：

1. 降低输入尺寸，例如 1024 -> 640。
2. 降低显示分辨率。
3. 减少右侧热力图刷新频率。
4. 把热力图通道均值计算放到 CUDA/TensorRT 插件或模型输出里。
5. 锁定 Jetson 性能模式和风扇。

21. 最终复现命令清单

21.1 上传模型和脚本

scp RepC3_chip3.pt nvidia@192.168.0.166:/home/nvidia/yolo_deploy/ReqC3_chip3/
scp export_repC3_attention_onnx.py nvidia@192.168.0.166:/home/nvidia/yolo_deploy/ReqC3_chip3/
scp jetson_fullscreen_attention.py nvidia@192.168.0.166:/home/nvidia/yolo_deploy/ReqC3_chip3/

21.2 导出 ONNX

cd /home/nvidia/yolo_deploy/ReqC3_chip3
export LD_PRELOAD=/home/nvidia/miniforge3/envs/yolo_trt/lib/libstdc++.so.6

/home/nvidia/miniforge3/envs/yolo_trt/bin/python3 ./export_repC3_attention_onnx.py \
  --model RepC3_chip3.pt \
  --output RepC3_chip3_attention.onnx \
  --imgsz 1024 \
  --device cpu

21.3 构建 TensorRT

/usr/src/tensorrt/bin/trtexec \
  --onnx=RepC3_chip3_attention.onnx \
  --saveEngine=RepC3_chip3_attention_fp16.trt \
  --fp16 \
  --memPoolSize=workspace:2048 \
  --verbose \
  > build_attention_trt.log 2>&1

21.4 验证输出

/home/nvidia/miniforge3/envs/yolo_trt/bin/python3 ./jetson_fullscreen_attention.py \
  --engine /home/nvidia/yolo_deploy/ReqC3_chip3/RepC3_chip3_attention_fp16.trt \
  --feature-source trt \
  --check-feature

21.5 启动实时界面

export DISPLAY=:1
export XAUTHORITY=/run/user/1000/gdm/Xauthority
export LD_PRELOAD=/home/nvidia/miniforge3/envs/yolo_trt/lib/libstdc++.so.6

/home/nvidia/miniforge3/envs/yolo_trt/bin/python3 ./jetson_fullscreen_attention.py \
  --engine /home/nvidia/yolo_deploy/ReqC3_chip3/RepC3_chip3_attention_fp16.trt \
  --feature-source trt \
  --screen-width 1920 \
  --screen-height 1080 \
  --left-ratio 0.50 \
  --heat-alpha 0.45
# ubuntu # 嵌入式硬件
Logo
AtomGit开源社区

AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念，把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起，为开发者提供从开发、训练到部署的一站式体验。

更多推荐

cover

速度与精度的结合：Faster R-CNN模型的性能剖析

avatar AtomGit开源社区
cover

昇腾NPU小模型推理性能调优实战：从1.5s到0.7s的优化之路

avatar AtomGit开源社区

09 GPT-2 论文精读：语言模型如何走向 Zero-shot？

Article:[一篇新闻文章]TL;GPT-2 会根据前面的文章和TL;这个提示继续生成摘要。这里的关键是TL;。它在网络文本中经常表示“太长不看版摘要”。如果模型在预训练中学过这种模式，它就可能在看到TL;后生成摘要内容。GPT-2 的摘要能力比较初级。当输入文章并使用TL;作为提示时，模型能生成一些摘要式文本，但定量指标仍然比较弱，只是刚刚接近一些传统神经基线。这说明 GPT-2 已经开始学

avatar AtomGit开源社区