dump op方法

  1. 使用profile接口对原始代码训练脚本的Loss计算和优化过程进行改造,打印算子信息。代码样例如下:

    with torch.autograd.profiler.profile() as prof:
        out = model(input_tensor)
        loss = out.sum()
        loss.backward()
    
    print(prof.key_averages().table(sort_by="self_cpu_time_total"))

    用户也可参考profiler.py文件修改使能代码,该文件默认位于/{python_install_path}/site-packages/torch/autograd/profiler.py

  2. 运行训练脚本并将日志重定向。

    python3 train.py >profiler.log

  3. 使用算子提取脚本解析日志生成all_ops.txt,脚本样例参考算子提取脚本样例

    python3 extract_ops_by_steop.py --profiler_file ./profiler.log