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
python3 train.py >profiler.log
python3 extract_ops_by_steop.py --profiler_file ./profiler.log