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