PyTorch 测量代码段的运行时间

Contents

timeit

  • timeit. 测量代码开始时刻和结束时刻,然后求差

  • pytorch 的代码经常会运行在 GPU 上,而在 GPU 上的运行都是异步的,意味着采用一般的 timeit 操作不能准确地得到运行时总和,因此我们一般需要用 pytorch 内置的计时工具和同步工具 (单位:ms)
python 复制代码
start = torch.cuda.Event(enable_timing=True)
end = torch.cuda.Event(enable_timing=True)

start.record()
z = x + y
end.record()

# Waits for everything to finish running
torch.cuda.synchronize()

print(start.elapsed_time(end))

profile

  • profile. pytorch 自带或者第三方的代码耗时工具

  • timeit 的方法测试一些小代码还勉强适用,但是在大规模的测试中显然会变得很麻烦,当然,你可以通过添加修饰器的方式去简化一行行重复人工添加这些时间测量代码的枯燥,但是这也并不是最好的解决方案
  • 幸运的是,pytorch 自带了 profile 用于计算模型每个部分耗时 ,其既可以计算 cpu 耗时,也可以计算 gpu 耗时
python 复制代码
x = torch.randn((1, 1), requires_grad=True)
with torch.autograd.profiler.profile(enabled=True) as prof:
	for _ in range(100):  # any normal python code, really!
    	y = x ** 2
print(prof.key_averages().table(sort_by="self_cpu_time_total"))

References

相关推荐
槑槑紫2 小时前
深度学习pytorch整体流程
人工智能·pytorch·深度学习
IT古董5 小时前
【第三章:神经网络原理详解与Pytorch入门】02.深度学习框架PyTorch入门-(5)PyTorch 实战——使用 RNN 进行人名分类
pytorch·深度学习·神经网络
机器学习之心7 小时前
小波增强型KAN网络 + SHAP可解释性分析(Pytorch实现)
人工智能·pytorch·python·kan网络
Green1Leaves13 小时前
pytorch学习-11卷积神经网络(高级篇)
pytorch·学习·cnn
灵智工坊LingzhiAI15 小时前
人体坐姿检测系统项目教程(YOLO11+PyTorch+可视化)
人工智能·pytorch·python
William.csj1 天前
Pytorch/CUDA——flash-attn 库编译的 gcc 版本问题
pytorch·cuda
Green1Leaves2 天前
pytorch学习-9.多分类问题
人工智能·pytorch·学习
摸爬滚打李上进2 天前
重生学AI第十六集:线性层nn.Linear
人工智能·pytorch·python·神经网络·机器学习
HuashuiMu花水木2 天前
PyTorch笔记1----------Tensor(张量):基本概念、创建、属性、算数运算
人工智能·pytorch·笔记
喝过期的拉菲3 天前
如何使用 Pytorch Lightning 启用早停机制
pytorch·lightning·早停机制