Python如何在应用退出前自动释放当前程序所占用的全部内存
机器学习训练内存不释放问题解决方案
首先明确Linux和Windows表现差异的核心原因:Linux默认使用glibc的内存分配器,会将Python/机器学习框架释放的空闲内存作为缓存保留,不会主动归还给操作系统,因此即使变量已经被回收,进程的常驻内存占用仍然会居高不下,Windows的内存分配策略更激进,会及时归还空闲内存,因此问题表现更轻。
以下是不需要手动定位删除变量的内存自动释放方案:
方案1:替换内存分配器(无代码侵入,优先推荐)
无需修改任何业务代码,仅在启动进程时替换为更激进回收内存的分配器即可,对PyTorch/TensorFlow等框架的内存碎片问题缓解效果显著:
- 先安装jemalloc:
Debian/Ubuntu系统执行:sudo apt install libjemalloc2
CentOS/RHEL系统执行:sudo dnf install jemalloc - 启动训练脚本时预加载jemalloc:
LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libjemalloc.so.2 python your_train_script.py - 实测该方案可降低30%以上的无效内存占用,且不会影响训练性能。
方案2:调用框架全局内存清理接口
不用手动指定变量,直接调用各框架提供的全局缓存释放接口即可清理绝大多数闲置内存:
- PyTorch用户:任务执行完成后依次执行
torch.cuda.empty_cache()(清理GPU缓存)、torch.cuda.ipc_collect()(清理IPC残留内存)、import gc; gc.collect()(清理CPU层面无引用的对象) - TensorFlow/Keras用户:执行
tf.keras.backend.clear_session(),会清空全局计算图和所有缓存张量 - Dask/Pandas等数据处理库用户:执行
client.close()(Dask)或主动触发全局gc即可,不需要手动删除DataFrame变量
方案3:子进程隔离执行任务(100%释放内存)
如果你的脚本需要执行多个独立的训练/测试任务,将每个任务放到单独的子进程中运行,任务结束后子进程退出,操作系统会自动回收该进程占用的全部内存,完全不需要手动处理内存:
import multiprocessing def run_single_task(config): # 所有模型加载、训练、测试逻辑全部放在该函数内 from your_module import YourModel, train, test model = YourModel(**config) train(model) test(model) if __name__ == "__main__": all_task_configs = [...] for task_config in all_task_configs: p = multiprocessing.Process(target=run_single_task, args=(task_config,)) p.start() p.join() # 此处上一个任务占用的所有内存已经被系统完全回收,不会影响后续任务
方案4:内存上限兜底
如果担心内存溢出影响其他服务,可以通过ulimit限制进程最大内存占用,超过阈值会被系统自动终止,避免占满整机内存:ulimit -v 62914560 && python your_script.py
其中62914560对应60GB内存的KB值,可根据自身硬件配置调整。
内容的提问来源于stack exchange,提问作者Sina
相关产品推荐
相关产品推荐

