之前发现算法处理完的tensor从gpu转到cpu很耗时,一直找不到解决方法,后面看到有人说测试时间要先同步。
torch.cuda.synchronize()
start = time.time()
result = model(input)
torch.cuda.synchronize()
end = time.time()
才发现耗时的不是这个转换过程
这是因为CUDA kernel函数是异步的,所以不能直接在CUDA函数两端加上time.time()测试时间,这样测出来的只是调用CUDA api的时间,不包括GPU端运行的时间。
我们需要加上线程同步函数,等待kernel中所有线程全部执行完毕再执行CPU端后续指令。上面代码我们将同步指令加在了python端,用的是torch.cuda.synchronize函数。
其中第一次同步是为了防止前面的代码中有未同步还在GPU端运行的指令,第二次同步就是为了等result = model(input)所有线程执行完毕后再统计时间。
关于同步和异步的区别,参考下面两图
同步执行当调用方法执行完成后并返回结果,才能执行后续代码。
异步调用的话可参考ajax,调用方法后不会等到sum方法执行完成,而是直接执行后续代码。sum方法执行完成后主要通过状态通知主线程,或者通过回调处理这次异步方法执行的结果。
版权声明:本文为CSDN博主「Stars-Chan」的原创文章
原文链接:https://blog.csdn.net/weixin_44942126/article/details/117605711
参考链接:
pytorch 正确的测试时间的代码 torch.cuda.synchronize()
PyTorch自定义CUDA算子教程与运行时间分析
同步(Synchronous)和异步(Asynchronous)