如何使用PyTorch彻底释放CUDA占用的GPU显存？

安全阿岳

2026-8-30

问题背景

基于spark + face-alignment实现高相似度人脸生成，核心代码逻辑如下：

fa = face_alignment.FaceAlignment(face_alignment.LandmarksType._2D, flip_input=False) # 基于PyTorch依赖调用GPU
imageVector.append( convertImagefa(image, fa))
del fa
gc.collect()
torch.cuda.empty_cache() # 尝试清理CUDA显存
return imageVector

问题现象

单机部署环境下运行4个可访问GPU的线程，设计调度策略为：每4个请求中仅允许1个请求调用GPU，该策略原本设计可匹配当前显存容量。
执行CUDA显存清理操作后显存无法完全释放：可观测到负载在线程间转移、部分显存被释放，但CUDA始终残留624MiB显存无法释放。
通过nvidia-smi命令查询到的显存占用情况如下：

nvidia-smi                                                                                                                                                              
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A     17132  C   .../face-the-same/bin/python      624MiB |
|    0   N/A  N/A     17260  C   .../face-the-same/bin/python     1028MiB |
|    0   N/A  N/A     17263  C   .../face-the-same/bin/python      624MiB |
|    0   N/A  N/A     17264  C   .../face-the-same/bin/python      624MiB |

根因分析

固定残留的624MiB显存是PyTorch为进程初始化CUDA运行环境时预留的基础上下文显存，这部分显存不受torch.cuda.empty_cache()控制。empty_cache()仅能释放PyTorch缓存池中未被使用的张量显存，只要进程不退出，这部分基础上下文显存就会一直被占用。
现有调度策略仅限制了同时执行GPU计算的请求数量，但只要对应Python进程曾经触发过CUDA调用，就会独立初始化一份CUDA上下文，每个进程都会占用一份固定的基础预留显存，和观测到的多进程各占624MiB的现象完全吻合。

解决方案

已验证落地方案：通过分布式锁将GPU计算任务固定绑定到单个executor/进程ID上执行，从根源上避免多进程各自初始化CUDA上下文，仅单个进程保留CUDA基础上下文，其余进程不触发任何CUDA调用，就不会产生额外的预留显存开销。
其他可选优化方向：
- 多进程部署场景下，可单独常驻1个专属GPU工作进程，所有GPU计算任务通过进程间队列提交给该进程处理，禁止其他业务进程直接调用CUDA相关接口。
- 不要尝试在不终止进程的前提下强制销毁CUDA上下文，当前PyTorch没有提供稳定的公开API支持该操作，强制卸载CUDA动态库极易触发段错误、隐式显存泄漏等不可预期问题。

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