PyTorch 中用于 主机（CPU）与设备（GPU）同步 的函数 torch.cuda.synchronize()

PyTorch 中用于 主机CPU与设备GPU同步 的函数 torch.cuda.synchronize()flyfish完整代码在文末GPU 是 异步执行 的CPU 发送指令给 GPU比如矩阵乘法、卷积CPU 不会等 GPU 算完直接跑去执行下一行代码GPU 在后台默默计算计算完了再通知 CPUtorch.cuda.synchronize() 到底做什么强制 CPU 停下来等待 GPU 把之前所有的计算任务全部执行完毕再继续运行后面的代码。无同步CPU 发完指令就溜GPU 后台干活加同步CPU 原地等待直到 GPU 干完活torch.cuda.synchronize(deviceNone)它会阻塞当前 CPU 线程直到指定 CUDA 设备上所有 streams 中的所有 kernels计算任务全部完成为止。device参数可选默认为当前设备torch.cuda.current_device()。调用底层 CUDA 的cudaDeviceSynchronize()或类似机制强制等待 GPU 把之前下发的所有工作做完。PyTorch以及几乎所有现代 CUDA 编程默认采用异步执行模型当在 Python 代码里写tensor model(input)、torch.mm(a, b)等 GPU 操作时CPU 只负责下发指令把 kernel 丢到 CUDA stream 里函数几乎立刻返回。GPU 在后台真正执行计算。CPU 和 GPU 是并行的CPU 可以继续往下跑 Python 代码而 GPU 还在算。避免 CPU 频繁等待 GPU。如果直接用time.time()测量时间会只测到下发指令的时间而不是 GPU 真正计算的时间导致计时严重偏小。什么场景使用torch.cuda.synchronize()准确测量 GPU 执行时间Benchmark / Profilingtorch.cuda.synchronize()# 先清空之前的残留任务可选但推荐starttime.time()# 模型/操作outputmodel(input)torch.cuda.synchronize()# 关键等 GPU 全部算完endtime.time()print(end-start)更推荐的现代写法使用 CUDA Event精度更高避免多余阻塞start_eventtorch.cuda.Event(enable_timingTrue)end_eventtorch.cuda.Event(enable_timingTrue)start_event.record()# ... 执行操作 ...end_event.record()torch.cuda.synchronize()# 必须等 event 被记录print(start_event.elapsed_time(end_event))# 单位毫秒调试时需要 GPU 计算真正完成后再看结果比如想print(tensor)或把结果转到 CPU.cpu()、.item()有时会隐式同步但不总是可靠。确认某个 kernel 是否真的执行完排查异步 bug。使用多个 CUDA Stream 时默认 stream 里操作通常不需要手动 syncPyTorch 会自动处理依赖。但如果用了自定义 streamtorch.cuda.Stream不同 stream 之间可能需要显式同步。某些需要严格顺序的场景比如在训练循环中每一步都想确保前一步完全结束一般不推荐会严重降低性能。什么场景不使用torch.cuda.synchronize()正常训练/推理时时几乎永远不要在循环里每一步都加torch.cuda.synchronize()它会强制 CPU 等待 GPU破坏异步并行大幅降低整体吞吐量。性能测试时也只在测量开始前和结束时各加一次不要每 iteration 都加除非故意想测包含同步开销的时间。使用 synchronize() 测量importtorchimporttime# 检查是否有 GPUifnottorch.cuda.is_available():print(CUDA 不可用请检查环境)exit()devicetorch.device(cuda:0)print(f使用设备:{device})# 创建一些较大的数据用于测试xtorch.randn(4096,4096,devicedevice)ytorch.randn(4096,4096,devicedevice)# 使用 synchronize() 测量 print(\n 使用 torch.cuda.synchronize() 测量 )# 预热非常重要第一次运行会有额外开销for_inrange(10):_torch.mm(x,y)torch.cuda.synchronize()# 开始正式计时torch.cuda.synchronize()# 确保之前所有操作完成start_timetime.time()# 要测量的操作这里用矩阵乘法为例foriinrange(100):ztorch.mm(x,y)# GPU 操作torch.cuda.synchronize()# 关键等待所有 GPU 操作完成end_timetime.time()elapsed(end_time-start_time)*1000# 转换为毫秒print(f使用 synchronize() 测量 100 次矩阵乘法耗时:{elapsed:.2f}ms)print(f平均每次:{elapsed/100:.2f}ms)使用 CUDA Event 测量importtorch# 检查 GPUifnottorch.cuda.is_available():print(CUDA 不可用)exit()devicetorch.device(cuda:0)print(f使用设备:{device})xtorch.randn(4096,4096,devicedevice)ytorch.randn(4096,4096,devicedevice)# 使用 CUDA Event 测量推荐 print(\n 使用 CUDA Event 测量推荐 )# 预热for_inrange(10):_torch.mm(x,y)torch.cuda.synchronize()# 创建 Eventstart_eventtorch.cuda.Event(enable_timingTrue)end_eventtorch.cuda.Event(enable_timingTrue)# 开始计时start_event.record()# 记录开始点# 要测量的操作foriinrange(100):ztorch.mm(x,y)end_event.record()# 记录结束点# 必须调用 synchronize() 才能读取时间torch.cuda.synchronize()# 等待两个 Event 都被记录完成elapsed_msstart_event.elapsed_time(end_event)# 返回毫秒print(f使用 CUDA Event 测量 100 次矩阵乘法耗时:{elapsed_ms:.2f}ms)print(f平均每次:{elapsed_ms/100:.2f}ms)