如何将PyTorch Profiler数据导出为CSV文件

PyTorch Profiler的key_averages()返回EventList对象，要把它导出成CSV，直接用Python内置的csv模块就能实现，步骤如下：

核心思路 #

1. 从EventList的每个Event对象中提取需要的分析字段（比如名称、耗时、调用次数等）
2. 用csv模块将字段和对应数据写入CSV文件

示例代码 #

import csv
import torch.profiler

# 假设你已经完成Profiler运行，得到profile对象
# profile = torch.profiler.profile(activities=[...], record_shapes=True)
# ... 执行被分析的代码 ...

# 获取key averages对应的EventList
events = profile.key_averages()

# 定义要导出的字段（可根据table()显示的列自行调整）
export_fields = [
    "operation_name",
    "call_count",
    "total_cpu_time",
    "self_cpu_time",
    "total_cuda_time",
    "self_cuda_time",
    "cpu_memory_usage",
    "cuda_memory_usage"
]

# 写入CSV文件
with open("profiler_output.csv", "w", newline="", encoding="utf-8") as csv_file:
    writer = csv.DictWriter(csv_file, fieldnames=export_fields)
    writer.writeheader()  # 写入表头
    
    for event in events:
        # 构造每行数据，兼容无CUDA的环境
        row_data = {
            "operation_name": event.key,
            "call_count": event.count,
            "total_cpu_time": event.cpu_time_total,
            "self_cpu_time": event.self_cpu_time_total,
            "total_cuda_time": event.cuda_time_total if hasattr(event, "cuda_time_total") else 0,
            "self_cuda_time": event.self_cuda_time_total if hasattr(event, "self_cuda_time_total") else 0,
            "cpu_memory_usage": event.cpu_memory_usage if hasattr(event, "cpu_memory_usage") else 0,
            "cuda_memory_usage": event.cuda_memory_usage if hasattr(event, "cuda_memory_usage") else 0
        }
        writer.writerow(row_data)

补充说明 #

• 可以通过dir(event)查看单个Event对象的所有可用属性，按需增减导出字段
• 如果需要平均耗时，可自行计算（比如total_cpu_time / call_count）后加入导出字段
• 无CUDA环境下，可直接删除CUDA相关字段，避免判断逻辑

内容的提问来源于stack exchange，提问作者aabdel