开场（痛点+类比）
你是不是也遇到这种尴尬：模型推理搬上 GPU 以后，GPU 忙得冒烟，CPU 却站在旁边发呆，单次推理还是要 250 多毫秒？别急，这更像一家串串店：总厨（GPU）在火力全开，一旁的配菜师傅（CPU）却在等菜出锅才去洗下一盆。我们要教的，就是让配菜师傅边等边干活，把流水线串起来。
原理速写
- GPU 调度是异步的：PyTorch 在背后维护一条 CUDA 任务队列，
model(dummy)只是把“菜谱”塞进队列，GPU 真正开炒时 CPU 已经获得控制权。也就是说，函数返回不代表 GPU 完活，只代表它“接单”了。 - 同步调用是刹车：一旦你写下
tensor.cpu().numpy()、torch.cuda.synchronize()这类语句，PyTorch 会等到 GPU 完全结束才继续往下走，相当于把正在洗菜的师傅叫停去盯锅。 - 重叠的核心是解耦：如果 CPU 要做的后处理能使用上一批数据就够了，就不用和本批的 GPU 输出死死绑在一起，本质是一种“提前准备、延后消费”的策略。
- 双缓冲最好用：最简单的方式是维护一个长度为 2 的队列，GPU 新产出的张量先入队，CPU 优先处理队头的旧张量，让两个厨师形成稳定流水线。
- 收尾要对齐：流水线上总有最后一个菜没被端出来，收尾时主动
torch.cuda.synchronize()扫尾，确保“锅里没余汤”再统计耗时。
可以把两个版本的时间线对比着看：
串行版: GPU[批0推理]---sync---CPU[批0后处理]---GPU[批1推理]---sync---CPU[批1后处理]
重叠版: GPU[批0推理]---sync---CPU[批0后处理]
\__CPU等待区处理上一批__/
GPU[批1推理]---sync---CPU[批1后处理]
第二条时间线里 CPU 和 GPU 的“空窗期”明显被压缩，这就是重叠能省下来的真时间。
实战步骤
1. 环境准备
python3 -m venv venv
source venv/bin/activate
pip install torch==2.3.0 torchvision==0.18.0
在一台有 CUDA 11.8+ 驱动的 Linux 主机上执行，确保 nvidia-smi 能看到显卡。
2. 写出“串行版”基线
保存为 sequential.py：
import math
import statistics
import time
import torch
import torchvision.models as models
def heavy_cpu_work():
acc = 0.0
for i in range(500_000):
acc += math.sqrt(i) * math.sin(i) * math.cos(i)
return acc
@torch.inference_mode()
def main():
device = "cuda"
model = models.vit_b_16(weights=None).to(device).eval()
dummy = torch.randn(10, 3, 224, 224, device=device)
_ = model(dummy) # 预热
durations = []
for _ in range(50):
start = time.perf_counter()
output = model(dummy)
torch.cuda.synchronize()
heavy_cpu_work()
durations.append(time.perf_counter() - start)
print(f"Sequential avg latency: {statistics.mean(durations):.4f} s")
if __name__ == "__main__":
main()
运行 python sequential.py，你会看到平均延迟大约在 0.26 s 左右。
代码讲解：
@torch.inference_mode()让模型推理时少走计算图的开销，省得构建梯度。dummy = torch.randn(...)构造一份固定尺寸的伪数据，方便你在没有真实输入的情况下测试。torch.cuda.synchronize()是关键：它确保 GPU 的活干完后再让 CPU 去执行heavy_cpu_work()，于是整段流程完全串行。heavy_cpu_work()用大量三角函数替代真实后处理，纯粹为了模拟 CPU 密集操作；black_box没用到是因为 Python 里不需要，但在 Rust 里常见。
3. 改写成“重叠版”
保存为 overlap.py：
import collections
import math
import statistics
import time
import torch
import torchvision.models as models
def heavy_cpu_work(batch_id):
acc = 0.0
for i in range(500_000):
value = batch_id * 31 + i
acc += math.sqrt(value) * math.sin(value) * math.cos(value)
return acc
@torch.inference_mode()
def main():
device = "cuda"
model = models.vit_b_16(weights=None).to(device).eval()
dummy = torch.randn(10, 3, 224, 224, device=device)
_ = model(dummy)
torch.cuda.synchronize()
backlog = collections.deque()
durations = []
for step in range(50):
start = time.perf_counter()
backlog.append((step, model(dummy)))
if len(backlog) >= 2:
old_step, _ = backlog.popleft()
heavy_cpu_work(old_step)
torch.cuda.synchronize()
durations.append(time.perf_counter() - start)
while backlog:
last_step, _ = backlog.popleft()
heavy_cpu_work(last_step)
print(f"Overlap avg latency: {statistics.mean(durations):.4f} s")
if __name__ == "__main__":
main()
运行 python overlap.py，平均延迟一般能掉到 0.17 s 左右 —— 同样的模型，多出三成时间给业务。
代码讲解：
backlog = collections.deque()相当于准备一个托盘，存放 GPU 刚出锅但 CPU 还没处理的菜。- 循环里先让 GPU 接手下一批：
backlog.append((step, model(dummy)))；随后只要队列里超过 1 批，就取出最老的那份交给 CPU 处理。 torch.cuda.synchronize()仍然存在，但位置被挪到了“处理完上一批 CPU 任务之后”，确保 GPU 能尽量长时间不被刹车。- 循环结束后还有一个
while backlog，那是为了兜底处理最后一批 CPU 任务，否则最后一道菜会被忘在台面上。
4. 快速对比
python sequential.py
python overlap.py
把两个结果放在一起，就能量化这波厨房分工带来的收益。
常见坑与对策
- CPU 仍然读取当前批次输出？立刻同步，重叠失效。解决：只处理上一次存好的张量。
- 忘记最后
torch.cuda.synchronize()？尾批推理时间没算进去。解决：收尾时强制同步一次。 - CPU 预处理还在等磁盘？流水线被 I/O 卡住。解决：把读盘放在单独线程或使用数据加载器的预取功能。
- 显存紧张？ backlog 越大越占内存。解决：保持队列长度为 2~3，出队后立刻
del旧张量。 - 不理解算力利用率？可以同时开
watch -n 0.5 nvidia-smi和htop，观察 GPU 带宽保持高位、CPU 核心同样忙碌时，说明重叠已经生效。
总结与下一步
- GPU 推理默认异步，别让 CPU 在门口干等。
- CPU/GPU 分工一旦解耦，平均延迟就能直接降下去。
- PyTorch 上手门槛不高，几行代码就能开启重叠模式。
下一步你可以：
- 把
heavy_cpu_work换成真实的后处理或日志写入，验证收益。 - 接入
torch.cuda.streams.Stream把预处理、推理、后处理拆成多条流水线。 - 用
nvtx标记关键区间，配合nsys看时间线，确保重叠真的发生。
读完有收获？不妨留言聊聊你在推理提速上的小妙招，也欢迎转发给同事一起调教流水线；多一个点赞，我们就多写一篇 CPU/GPU 协作的实战故事。