异步执行
默认情况下,GPU 操作是异步的。当您调用使用 GPU 的函数时,操作会被排入特定设备的队列,但不一定要等到稍后才会执行。这允许我们并行执行更多计算,包括在 CPU 或其他 GPU 上的操作。
一般来说,异步计算的影响对调用者是不可见的,因为(1)每个设备按照它们排队的顺序执行操作,以及(2)PyTorch 在 CPU 和 GPU 之间或两个 GPU 之间复制数据时会自动执行必要的同步。因此,计算将继续进行,就好像每个操作都是同步执行的一样。
您可以通过设置环境变量 CUDA_LAUNCH_BLOCKING=1 来强制同步计算。当 GPU 发生错误时,这会很方便。 (对于异步执行,这样的错误直到操作实际执行后才会报告,因此堆栈跟踪不会显示它被请求的位置。)
异步计算的结果是没有同步的时间测量是不准确的。要获得精确的测量值,应该在测量前调用 torch.cuda.synchronize(),或者使用 torch.cuda.Event 记录时间,如下所示:
start_event = torch.cuda.Event(enable_timing=True)
end_event = torch.cuda.Event(enable_timing=True)
start_event.record()
# 在这里运行一些东西
end_event.record()
torch.cuda.synchronize() # 等待事件被记录!
elapsed_time_ms = start_event.elapsed_time(end_event)