PyTorch 性能分析实战：像手术刀一样精准控制 Nsys Timeline（附自定义颜色教程）

在进行深度学习模型性能优化时，NVIDIA Nsight Systems (nsys) 是我们最得力的助手。但在默认情况下，nsys 往往会抓取整个运行过程，导致生成的文件巨大且难以分析。

本文将介绍如何在 PyTorch 代码中"手动插桩"，实现仅抓取特定 Step 的数据，并为 Timeline 上的不同阶段打上自定义颜色的标签，让性能分析变得清晰直观。

一、 核心接口对照表

如果你熟悉 Paddle 的 core.nvprof_* 接口，迁移到 PyTorch 非常简单。PyTorch 的相关功能主要位于 torch.cuda 模块下：

功能	Paddle 接口	PyTorch 接口 (新)
开始抓取	core.nvprof_start()	torch.cuda.cudart().cudaProfilerStart()
停止抓取	core.nvprof_stop()	torch.cuda.cudart().cudaProfilerStop()
标记开始 (Push)	core.nvprof_nvtx_push(msg)	torch.cuda.nvtx.range_push(msg)
标记结束 (Pop)	core.nvprof_nvtx_pop()	torch.cuda.nvtx.range_pop()

---

二、 基础实战：精准抓取指定 Step

为了避免 Trace 文件过大，我们通常采用 "预热 -> 开启抓取 -> 运行 N 步 -> 关闭抓取" 的策略。

以下是一个最简 Demo，模拟了前 90 步预热，仅抓取第 90 到 95 步的逻辑：

import torch

# 检查环境
if not torch.cuda.is_available():
    raise RuntimeError("需要 GPU 环境运行此 Demo")

device = torch.device("cuda:0")
x = torch.randn(2000, 2000, device=device)
y = torch.randn(2000, 2000, device=device)

print(">>> 开始训练循环...")

for ii in range(100):
    
    # 1. Start: 在第 90 步开启 Profiler
    # 配合 nsys 命令行的 -c cudaProfilerApi 参数生效
    if ii == 90:
        print(">>> 开启 Profiler (Start)")
        torch.cuda.cudart().cudaProfilerStart()

    # 2. Push: 标记当前 Step 的开始
    # 这会在 Timeline 上生成一个名为 "Step_XX" 的区间块
    torch.cuda.nvtx.range_push(f"Step_{ii}_Calculation")

    # ...... 实际模型计算逻辑 ......
    z = torch.matmul(x, y)
    torch.cuda.synchronize() # 仅做演示，保证有负载

    # 3. Pop: 标记当前 Step 结束 (必须与 Push 成对)
    torch.cuda.nvtx.range_pop()

    # 4. Stop: 在第 95 步结束抓取
    if ii == 95:
        print("<<< 关闭 Profiler (Stop)")
        torch.cuda.cudart().cudaProfilerStop()
        # 此时 Trace 文件已生成，无需继续记录
        break

print("循环结束")

---

三、 进阶实战：自定义 Timeline 颜色

为了在 Nsys 的 GUI 界面中一眼区分出 前向传播 (Forward) 和 反向传播 (Backward)，或者是区分不同的 Layer，我们可以给 NVTX 标记上色。

这里推荐使用 NVIDIA 官方提供的 Python 库 nvtx，它比 PyTorch 原生接口提供了更丰富的可视化选项。

1. 安装依赖

pip install nvtx

2. 带颜色的代码实现

nvtx 提供了非常优雅的上下文管理器 (with 语句) 和装饰器，让代码更整洁。

import torch
import nvtx  # 引入 nvidia 官方包

# 定义颜色常量 (ARGB 格式: 0xAARRGGBB)
COLOR_GREEN = 0xFF00FF00
COLOR_RED   = 0xFFFF0000

device = torch.device("cuda:0")
x = torch.randn(2000, 2000, device=device)
y = torch.randn(2000, 2000, device=device)

print(">>> 开始带颜色的 Profiling...")

# 开启抓取
torch.cuda.cudart().cudaProfilerStart()

for ii in range(5):
    
    # === 方式 A: 上下文管理器 (推荐) ===
    # 这段逻辑在 Timeline 上将显示为【绿色】长条
    with nvtx.annotate(message=f"Step_{ii}_Forward", color=COLOR_GREEN):
        z = torch.matmul(x, y)
        torch.cuda.synchronize()

    # === 方式 B: 手动 Push/Pop ===
    # 这段逻辑在 Timeline 上将显示为【红色】长条
    nvtx.push_range(message=f"Step_{ii}_Backward", color=COLOR_RED)
    
    z = z * 2 # 模拟反向传播计算
    torch.cuda.synchronize()
    
    nvtx.pop_range() # 记得关闭区间

# 停止抓取
torch.cuda.cudart().cudaProfilerStop()
print("结束。请使用 Nsys 查看红绿相间的 Timeline。")

---

四、 运行 Nsys 的正确姿势

写好了代码，还需要配合正确的启动命令。最关键的参数是 --capture-range=cudaProfilerApi，它告诉 nsys 不要从头抓到尾，而是听从代码中 Start/Stop 接口的指挥。

在终端执行以下命令：

nsys profile \
    --capture-range=cudaProfilerApi \
    --trace=cuda,nvtx,osrt \
    --force-overwrite=true \
    --output=my_colored_timeline \
    python your_script.py

参数解析：

• -c cudaProfilerApi (或 --capture-range): 核心参数，激活代码中的 Start/Stop 开关。
• -t cuda,nvtx (或 --trace): 显式声明我们需要抓取 CUDA 内核流和 NVTX 标记流。
• -o: 输出文件名。

五、 效果展示

运行完成后，使用 PC 端的 NVIDIA Nsight Systems GUI 打开生成的 .nsys-rep 文件。你将看到：

1. 时间轴极其干净：只包含你关心的那几步（Step 90-95）。
2. NVTX 层次分明：在 Timeline 顶部，你能清晰地看到绿色的 Forward 块和红色的 Backward 块交替出现，不再是一堆杂乱无章的内核调用。

希望这篇指南能帮助你更高效地分析模型性能！