深度学习中TensorFlow/PyTorch CPU与GPU运行内存占用差异问题问询

• CUDA运行时固有基础开销：CUDA初始化时会默认加载全量算子内核、驱动上下文、PTX编译缓存，单进程独立创建CUDA上下文时这部分开销不共享。CUDA 11.x版本为兼容多架构GPU，预编译内核体积更大，空载基础开销可达2GB GPU内存、5GB CPU内存，和业务代码的实际张量占用无关。
• 框架默认内存策略：TensorFlow v1 Session默认预占全部可用GPU内存，TF_CUDNN_USE_AUTOTUNE=0仅关闭卷积算法搜索缓存，不会改变预占逻辑；PyTorch默认关闭内核懒加载，第一次CUDA调用时就会加载全量算子内核到内存，进一步放大空载开销。
• 环境版本不兼容：你使用的370版本Nvidia驱动不支持CUDA 11.x（CUDA 11.x最低要求驱动版本≥450.80.02），强制运行会触发兼容层转译、PTX实时编译逻辑，额外产生30%以上的内存开销。
• 内存碎片化：多进程并发分配释放显存时，CUDA分配器的碎片化问题会导致实际内存占用远高于业务所需的张量大小。

1. 修复环境兼容问题
   升级Nvidia驱动到450.80.02及以上的稳定版本，匹配当前使用的CUDA 11.x版本，消除兼容层带来的额外开销。
2. 开启CUDA内核懒加载
   设置环境变量CUDA_MODULE_LOADING=LAZY（CUDA 11.7及以上）或CUDA_DISABLE_PRELOAD_KERNELS=1（CUDA 11.0-11.6），仅加载业务实际用到的算子内核，可将空载基础内存开销降低70%以上。
3. 调整框架内存配置

• TensorFlow 2兼容v1接口场景，创建Session前添加如下配置，关闭显存预占、限制单进程显存上限：

import tensorflow as tf
config = tf.compat.v1.ConfigProto()
# 开启按需分配显存
config.gpu_options.allow_growth = True
# 单进程最多占用10%显存，可按需调整
config.gpu_options.per_process_gpu_memory_fraction = 0.1
sess = tf.compat.v1.Session(config=config)

• PyTorch场景，设置环境变量PYTORCH_NO_CUDA_MEMORY_CACHING=1关闭不必要的显存缓存，或调用torch.cuda.set_per_process_memory_fraction(0.1)限制单进程显存上限。

4. 共享CUDA上下文
   同GPU上运行的多进程可开启CUDA上下文共享，10个进程共用同一份基础运行开销，避免重复占用内存：TensorFlow可使用分布式集群配置，PyTorch可使用torch.multiprocessing并设置启动方式为spawn开启上下文共享。
5. 降低CUDA版本
   如果仍需更高部署密度，可切换到CUDA 10.2版本，其空载基础开销比CUDA 11.x低40%左右，可进一步减少内存占用。

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