如何让PyTorch预先保留指定GPU内存，解决共享机器内存占用冲突问题？

可以预先锁定22GB GPU内存，避免第一个工作流运行时其他进程抢占资源导致第二个工作流崩溃，以下是主流框架的实现方法：

PyTorch 实现方式 #

在加载任何模型或执行第一个工作流之前，先分配一个占位张量占据目标内存：

import torch

# 计算22GB对应的float32张量元素数（1个float32占4字节）
reserved_elements = 22 * 1024**3 // 4
device = torch.device("cuda")

# 创建空张量锁定内存，保持引用防止被GC回收
dummy_tensor = torch.empty(reserved_elements, dtype=torch.float32, device=device)

第一个工作流运行时，PyTorch会在已分配的内存池中使用所需的8GB，剩余内存被占位张量占用，其他进程无法获取。进入第二个工作流前，释放占位张量即可：

del dummy_tensor
torch.cuda.empty_cache()

TensorFlow 实现方式 #

通过配置固定内存限制，让TensorFlow启动时直接申请22GB内存：

import tensorflow as tf

gpus = tf.config.list_physical_devices('GPU')
if gpus:
    try:
        # 为目标GPU设置22GB内存限制（单位为MB）
        tf.config.set_logical_device_configuration(
            gpus[0],
            [tf.config.LogicalDeviceConfiguration(memory_limit=22 * 1024)]
        )
    except RuntimeError as e:
        # 捕获内存不足或配置失败的异常
        print(f"内存配置失败: {e}")

配置完成后，TensorFlow会一直持有这22GB内存，后续两个工作流都在这个内存池内调度，不会被外部进程抢占。

关键注意点 #

• 确认目标GPU的总物理内存≥22GB，否则预分配操作会直接失败。
• 在共享机器上使用该方法前，建议和管理员或其他用户沟通，避免影响他人的正常使用。
• PyTorch中torch.cuda.empty_cache()仅释放框架内部未使用的缓存，必须显式删除占位张量才能释放锁定的内存供第二个工作流使用。

内容的提问来源于stack exchange，提问作者Nagabhushan S N