pytorch 模型GPU推理时间探讨3——正确计算模型推理时间

前言

  上文说到，在统计pytorch模型的推理时间时发现每次的前几次推理耗时都非常多，而且在后面多次的推理中，其时间也呈现出很大的变化，后来经过调研，得知模型在GPU上推理时，需要对GPU进行一个warm up阶段，使得显卡达到工作状态。对于后续多次推理时间的变化，差异较大，这一问题，是因为使用的计时工具是python标准库time里的perf_counter方法，该方法由CPU来执行，也就是说计时是在CPU上进行的，而实际计算是在GPU上进行的，这就涉及到同步异步计算的问题了，下面介绍同步异步计算以及怎么正确的计算模型推理时间。
  还没看过前两篇的同学可以去看一看。下面是链接
  Pytorch GPU模型推理时间探讨
  Pytorch GPU模型推理时间探讨2——显卡warm up

同步计算

  上图表示同步计算，（手工画图，字见谅），CPU将要计算的数据以及模型通过总线发送给GPU，等发送完后通知GPU开始计算，这时开始perf_counter计数，显卡开始计算，此时CPU陷入阻塞，啥事也不做，等待显卡完成计算，当显卡完成计算后，将结果和结束消息发送给CPU，CPU停止计数。这一过程就是同步计算，CPU和GPU之间进行通信需要通过总线，而通信就有开销，所以这会导致CPU的计时极大的不准确性，同步的方式还带来了CPU的资源浪费，在一段时间里啥事也不做。

异步计算

  知道了同步计算，异步计算就好理解了，上图中，CPU将数据和模型传到GPU后，通知GPU后就开始一个计数的线程，然后就去做其他事，等GPU完成计算后，给CPU发一个消息，CPU再停止计数。这样的异步计算，虽然没有CPU资源的浪费，但还是有通信上的损失，所以计时还是不准确。
  所以，我们需要在GPU上，让GPU自己对自己进行计时。

正确计算GPU推理时间

  如上图所示，正确的计时方法应该是由GPU自身发起，而不是在CPU上执行time的方法。在torch中，有torch.cuda.Event这个类可以用来计时。具体用法例子如下：

model = Model()
input = th.randn(batch_size, in_c, in_h, in_w, dtype = th.float32)

starter, ender = th.cuda.Event(enable_timing=True), th.cuda.Event(enable_timing=True)

start = perf_counter()
model.to(device)
input = input.to(device)
to_gpu = (perf_counter() - start) * 1000

# GPU warm-up
starter.record()
for _ in range(10):
    _ = model(input)
ender.record()
th.cuda.synchronize()
warm_up_time = starter.elapsed_time(ender)
print("GPU warm up time: ", warm_up_time)

timings = []
with th.no_grad():
    for i in range(100):
        starter.record()
        res = model(input)
        ender.record()
        # wait for GPU sync
        th.cuda.synchronize()
        curr_timing = starter.elapsed_time(ender)
        timings.append(round(curr_timing, 3))

# do whatever you want

如果对你有帮助，不妨点个赞再走？