如何为特定PyTorch实例强制设置GPU内存占用上限？

当然有办法实现！在PyTorch生态里，有几种实用的方式可以给单个进程限制GPU内存的使用上限，刚好匹配你12GB GPU只分配4GB的需求，我给你拆解一下常用的方案：

方法一：使用PyTorch内置的进程级显存限制API #

这是最直接的代码层面控制方式，用torch.cuda.set_per_process_memory_fraction()函数，它可以指定当前进程能使用的GPU内存占总显存的比例。

比如你的GPU总显存是12GB，要限制到4GB，那比例就是4/12≈0.333，代码示例如下：

import torch

# 必须在任何CUDA操作（比如模型移到GPU、创建CUDA张量）之前调用
torch.cuda.init()
# device参数指定GPU编号，单GPU场景填0即可
torch.cuda.set_per_process_memory_fraction(0.333, device=0)

# 接下来正常编写你的PyTorch逻辑
model = torch.nn.Sequential(
    torch.nn.Linear(1024, 2048),
    torch.nn.ReLU(),
    torch.nn.Linear(2048, 10)
).cuda()

注意：这个函数会严格限制进程的显存使用，一旦超过比例就会触发显存不足的错误，适合需要精确控制的场景。

方法二：通过环境变量全局配置 #

如果你不想修改代码，或者需要快速测试不同的显存限制，可以通过设置PYTORCH_CUDA_ALLOC_CONF环境变量来实现，直接指定最大可用显存（单位是MB）。

方式1：启动脚本时设置 #

在终端里运行你的Python脚本时，直接加上环境变量：

PYTORCH_CUDA_ALLOC_CONF=memory_limit:4096 python your_training_script.py

这里的4096就是4GB（因为1GB=1024MB）。

方式2：在代码开头设置 #

如果要在代码里动态设置，记得要在导入torch之前配置环境变量：

import os
# 设置最大可用显存为4096MB（4GB）
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'memory_limit:4096'
import torch

# 之后正常使用PyTorch
x = torch.randn(1000, 1000).cuda()

方法三：手动监控与释放显存（非强制限制） #

如果你需要更精细的显存管理，但不一定要强制上限，可以通过torch.cuda.memory_allocated()实时监控当前已用显存，配合torch.cuda.empty_cache()释放闲置显存。不过这种方式是手动管理，不是强制限制，适合需要灵活调整的场景：

import torch

# 监控当前已用显存
print(f"已分配显存: {torch.cuda.memory_allocated(0)/1024**2:.2f} MB")

# 释放闲置显存
torch.cuda.empty_cache()

以上几种方法里，前两种是真正的强制限制，推荐根据你的使用场景选择——代码里控制用方法一，快速测试用方法二的环境变量配置。

内容的提问来源于stack exchange，提问作者Giorgos Sfikas