昇腾社区首页
中文
注册
开发者
下载
支持

PyTorch训练场景性能分析快速入门

PyTorch训练场景下，推荐通过Ascend PyTorch Profiler接口采集并解析性能数据，用户可以根据结果自行分析和识别性能瓶颈。
Ascend PyTorch Profiler接口进行采集任务时，进程与Device之间的关系如下：
  • 多进程多Device场景：支持每个Device下分别设置一个采集进程。
  • 单进程多Device场景：支持。须配套PyTorch 2.1.0post14、2.5.1post2、2.6.0及之后的版本。
  • 多进程单Device场景：需要保证多进程之间的采集动作是串行的，即各个采集动作不在同一时间开始，且各个采集动作须包含完整的启动和停止。

前提条件

  • 请确保安装Ascend-cann-toolkit包。
  • 准备好基于PyTorch 2.1.0或更高版本开发的训练模型以及配套的数据集，并按照《PyTorch 训练模型迁移调优指南》中的“模型迁移”完成PyTorch原始模型向昇腾AI处理器的迁移。

采集并解析性能数据

  1. 使用Ascend PyTorch Profiler接口开启PyTorch训练时的性能数据采集。

    在训练脚本（如train_*.py文件）内添加如下示例代码进行性能数据采集参数配置，之后启动训练。如下列示例代码。

    • 下列示例中的接口详细介绍请参见Ascend PyTorch Profiler接口说明。
    • PyTorch场景性能数据采集详细介绍请参见Ascend PyTorch Profiler。
    • 性能数据会占据一定的磁盘空间，可能存在磁盘写满导致服务器不可用的风险。性能数据所需空间跟模型的参数、采集开关配置、采集的迭代数量有较大关系，须用户自行保证落盘目录下的可用磁盘空间。
    • 示例一：使用with语句调用torch_npu.profiler.profile接口，自动创建Profiler，采集with范围内代码段的性能数据。
       1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
      import torch
import torch_npu

...

experimental_config = torch_npu.profiler._ExperimentalConfig(
    export_type=[
        torch_npu.profiler.ExportType.Text
        ],
    profiler_level=torch_npu.profiler.ProfilerLevel.Level0,
    msprof_tx=False,
    aic_metrics=torch_npu.profiler.AiCMetrics.AiCoreNone,
    l2_cache=False,
    op_attr=False,
    data_simplification=False,
    record_op_args=False,
    gc_detect_threshold=None
)

with torch_npu.profiler.profile(
    activities=[
        torch_npu.profiler.ProfilerActivity.CPU,
        torch_npu.profiler.ProfilerActivity.NPU
        ],
    schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active=1, repeat=1, skip_first=1),    # 与prof.step()配套使用
    on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./result"),
    record_shapes=False,
    profile_memory=False,
    with_stack=False,
    with_modules=False,
    with_flops=False,
    experimental_config=experimental_config) as prof:
    for step in range(steps):
        train_one_step(step, steps, train_loader, model, optimizer, criterion)
        prof.step()    # 与schedule配套使用
    • 示例二：创建torch_npu.profiler.profile对象，通过start和stop接口控制采集性能数据，用户可自定义采集启动的位置。
       1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
      import torch
import torch_npu
...

experimental_config = torch_npu.profiler._ExperimentalConfig(
    export_type=[
        torch_npu.profiler.ExportType.Text
        ],
    profiler_level=torch_npu.profiler.ProfilerLevel.Level0,
    msprof_tx=False,
    aic_metrics=torch_npu.profiler.AiCMetrics.AiCoreNone,
    l2_cache=False,
    op_attr=False,
    data_simplification=False,
    record_op_args=False,
    gc_detect_threshold=None
)

prof = torch_npu.profiler.profile(
    activities=[
        torch_npu.profiler.ProfilerActivity.CPU,
        torch_npu.profiler.ProfilerActivity.NPU
        ],
    schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active=1, repeat=1, skip_first=1),    # 与prof.step()配套使用
    on_trace_ready=torch_npu.profiler.tensorboard_trace_handler("./result"),
    record_shapes=False,
    profile_memory=False,
    with_stack=False,
    with_modules=False,
    with_flops=False,
    experimental_config=experimental_config)
prof.start()
for step in range(steps):
    train_one_step()
    prof.step()    # 与schedule配套使用
prof.stop()

    以上两个示例主要使用tensorboard_trace_handler导出性能数据，也可以使用以下prof.export_chrome_trace方式导出单个性能文件“chrome_trace_{pid}.json”。其中tensorboard_trace_handler导出的性能数据包含了prof.export_chrome_trace导出的性能数据，所以根据实际需求选择一种方式即可，两种方式无冲突，也可以同时开启。

    1
2
3
4
5
6
7
8
9
    import torch
import torch_npu

...

with torch_npu.profiler.profile() as prof:
    for step in range(steps):
        train_one_step(step, steps, train_loader, model, optimizer, criterion)
prof.export_chrome_trace("./chrome_trace_14.json")
  2. 查看采集到的性能数据结果文件。
    训练结束后，在torch_npu.profiler.tensorboard_trace_handler接口指定的目录下生成Ascend PyTorch Profiler接口的采集结果目录，如下示例。
    以下数据文件用户无需打开查看，可使用《MindStudio Insight工具用户指南》工具进行性能数据的查看和分析，如需了解详细字段解释请参见MindSpore&PyTorch框架性能数据文件参考。
     1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
    └── msprof_1784298_20250620085947065_ascend_pt
    ├── ASCEND_PROFILER_OUTPUT
    │   ├── kernel_details.csv
    │   ├── operator_details.csv
    │   ├── step_trace_time.csv
    │   └── trace_view.json
    ├── FRAMEWORK
...
    ├── PROF_000001_20250620085947066_FLRBJLNFMBIDRPMB
    │   ├── device_1
    │   │   ├── data
...
    │   ├── host
    │   │   ├── data
...
    │   ├── mindstudio_profiler_log
...
    │   └── mindstudio_profiler_output
    │       ├── api_statistic_20250620085954.csv
    │       ├── msprof_20250620085953.json
    │       ├── op_summary_20250620085954.csv
    │       ├── README.txt
    │       └── task_time_20250620085954.csv
    ├── profiler_info.json
    └── profiler_metadata.json