# 使用ascend-pytorch适配的profiling接口,推荐先运行10个step,再profiling 1个step。 with torch.autograd.profiler.profile(use_npu=True) as prof: out = model(input_tensor) loss=loss_func(out) loss.backward() optimizer.zero_grad() optimizer.step() # 打印Profiling结果信息 print(prof) # 导出chrome_trace文件到指定路径 output_path = '/home/HwHiAiUser/profile_data.json' prof.export_chrome_trace(output_path)
PyTorch profiling更多内容请参见《性能分析工具使用指南》"高级功能>性能数据采集(AI框架方式)>PyTorch Profiling"章节。