PyTorch首次推理GPU显存占用与耗时远高于后续推理的原因咨询

这是PyTorch+CUDA技术栈的正常冷启动行为，和代码逻辑、统计方式无关，核心成因分为三类：

• CUDA内核JIT编译开销
  PyTorch默认对CUDA算子采用懒加载策略：第一次调用某个算子时，才会触发对应CUDA内核的即时编译、GPU硬件指令适配流程，编译过程本身会消耗数秒时间，同时产生大量临时显存占用。编译完成的内核会常驻CUDA上下文缓存，后续同形状输入调用算子时直接读取缓存，无需重复编译。本次测试首次推理4.7s的耗时中，大部分来自这部分编译成本。
• CUDA上下文初始化与显存预分配
  第一次向GPU提交计算任务时，CUDA驱动会完成全量运行环境初始化：涵盖设备能力探测、驱动与用户态runtime握手、显存管理结构初始化、cuBLAS/cuDNN等加速库加载流程。这一步会预先申请一块较大的显存作为后续计算的缓存池，首次测试测得的10.4GB显存中，大部分是这部分预分配的管理显存，并非模型推理实际使用的权重、激活显存。初始化完成的上下文会在进程生命周期内常驻，后续推理直接复用，不会重复触发该流程。

测试代码中调用的torch.cuda.empty_cache()仅会释放PyTorch缓存池中未被占用的显存块归还给CUDA驱动，不会销毁已初始化的CUDA上下文，也不会清除已编译完成的内核缓存，因此后续迭代不会重复产生初始化开销。

• cuDNN算法自动调优开销
  如果开启了torch.backends.cudnn.benchmark配置，第一次遇到固定尺寸的输入时，cuDNN会枚举所有适配当前GPU型号、输入形状的算子实现，通过实际运行测速选出性能最优的算法，该过程同样会产生额外耗时与临时显存占用。选中的最优算法会被缓存，后续同形状输入直接复用即可。

实测数据匹配性说明 #

给出的测试结果完全符合冷启动特征：

• 显存维度：首次推理的显存统计值包含了初始化、编译、算法调优产生的临时显存与预分配显存池，首次推理结束后临时显存被释放，显存池回落到满足常规推理需求的1.2GB左右，后续轮次无额外临时开销，因此显存占用保持稳定。
• 耗时维度：首次推理的耗时包含了所有一次性冷启动成本，后续轮次仅包含纯计算、CPU-GPU数据传输的固定开销，因此耗时稳定在0.85-0.9s区间，波动幅度极小。

如果需要获取准确的推理性能、显存指标，可在正式统计前执行1-3次预热（warmup）推理，待所有冷启动流程完成后再开始记录数据，即可得到稳定的真实推理表现。

内容的提问来源于stack exchange，提问作者Leo