如何在Python中解决TensorFlow与PyTorch显存冲突-Python教程-PHP中文网

tensorflow默认占满gpu显存导致pytorch报cuda内存不足，根本原因是显存被锁死而非总量不足；应启用memory_growth、切换前调用empty_cache、严格进程隔离，禁止单进程混用两框架。

如何在python中解决tensorflow与pytorch显存冲突

TensorFlow默认占满GPU显存，PyTorch无法分配

这是最常见冲突：TensorFlow 1.x/2.x 默认启用“内存增长”关闭状态，启动即申请全部可见GPU显存；PyTorch后续调用 torch.cuda.memory_allocated() 时发现无可用显存，报错 RuntimeError: CUDA out of memory。根本不是显存总量不够，而是被 TensorFlow “锁死”了。

实操建议：

python 查询技能

查询客流数据，输出JSON格式，可直接导入Bitable等可视化工具

下载

TensorFlow 2.4+：在导入 tensorflow 后、创建任何模型前，立即设置内存增长（推荐）：

import tensorflow as tf
gpus = tf.config.list_physical_devices('GPU')
if gpus:
    for gpu in gpus:
        tf.config.experimental.set_memory_growth(gpu, True)

TensorFlow 2.0–2.3：用 tf.config.set_memory_growth()（注意无 experimental. 前缀），行为一致
避免使用 tf.config.set_memory_limit() 硬限制——它仍会预分配指定大小，且 PyTorch 无法感知该限制，容易误判剩余显存

PyTorch与TensorFlow共用同一GPU时显存碎片化

即使开了 set_memory_growth，两者交替运行（如 TF 训练完立刻跑 PyTorch 推理）仍可能失败。原因：TensorFlow 释放显存不及时，PyTorch 的 CUDA 上下文看到的是“碎片化”空间，无法满足连续大块分配请求。

实操建议：

强制清空 PyTorch 缓存：在切换前加 torch.cuda.empty_cache()（注意它不释放 TF 占用的内存，只清理 PyTorch 自己的缓存）
隔离 GPU 设备：用环境变量限定各自可见卡，例如启动时分别设置：
```
CUDA_VISIBLE_DEVICES=0 python tf_script.py
```
和
```
CUDA_VISIBLE_DEVICES=1 python pt_script.py
```
若必须共卡，优先让 TensorFlow 先运行、完成、退出进程，再启动 PyTorch——进程级隔离比库级更可靠

混合框架调试时 `nvtop` 显示显存占用不一致

你看到 nvtop 里显存用了 8GB，但 tf.config.experimental.get_memory_info('GPU:0') 返回 {'current': 2e9, 'peak': 5e9}，而 torch.cuda.memory_reserved() 只有 1GB——这不是 bug，是三套内存管理机制独立计数的结果。

关键点：

TensorFlow 的 get_memory_info 统计其内部 BFCAllocator 分配器的用量（含预留未用部分）
PyTorch 的 memory_reserved() 是 CUDA caching allocator 已保留但未分配给 tensor 的显存
nvtop 显示的是驱动层实际占用（包括 CUDA context、TF/PT allocator 底层页、甚至某些驱动缓存），数值最大也最“真实”
不要试图对齐三者数值——以 nvtop 为准判断物理是否溢出，以各框架自身 API 为准做内部调优

多进程混合调用（如 Flask API 同时加载 TF/PT 模型）

Web 服务中一个进程内同时 import tensorflow 和 import torch 极易崩溃，典型错误是 cuInit failed: CUDA_ERROR_NO_DEVICE 或段错误。主因是两框架对 CUDA context 初始化顺序和共享方式不兼容。

实操建议：

绝对不要在同一个 Python 进程中混用 TF 和 PT 的 GPU 操作——哪怕只是 tf.constant([1]) + torch.tensor([1]).cuda()
改用进程隔离：用 multiprocessing.Process 或 subprocess.Popen 启动子进程，每个子进程只加载一种框架，并通过 pickle / pipe / Redis 通信
若用 FastAPI/Flask，配合 uvicorn --workers 2 并确保每个 worker 只绑定一种框架（例如 worker 0 加载 TF，worker 1 加载 PT），避免单 worker 内混用

TensorFlow 和 PyTorch 的 GPU 内存管理从设计上就互不感知，所谓“解决冲突”，本质是规避而非调和。最容易被忽略的是：即使代码逻辑上没同时调用，只要两个库都被 import 过，CUDA context 就可能在后台悄悄初始化并互相干扰——所以 import 顺序、进程边界、环境变量隔离，比任何运行时参数都关键。