使用torch.zeros()创建CUDA张量时异常占用主机RAM的问题求助

首先，你遇到的这个现象其实是PyTorch + CUDA运行时的正常行为，并不是“无理由”占用主机内存，下面我来拆解原因并给出针对性解决方案：

为什么torch.zeros(..., device='cuda')会涨主机RAM？ #

你看到的从533.28到661.29的RAM上涨，核心原因是CUDA上下文的初始化开销，而非张量本身占用的主机内存：

• 当你第一次创建CUDA张量时，CUDA运行时会初始化整个进程级别的CUDA上下文（包含驱动、运行时、设备管理等一系列底层资源），这部分会占用几百MB的主机内存，而且是持久化的——只要Python进程不结束，这部分内存就会被保留，gc.collect()也无法回收它，因为它不属于Python垃圾回收的管辖范围。
• 至于你创建的大CUDA张量，虽然数据本身存储在GPU显存中，但PyTorch仅需在主机端维护少量元数据（比如张量形状、设备信息、GPU内存指针等），这部分占用的内存非常小，几乎可以忽略。

你可以做一个简单的验证来确认：

getram()
# 手动初始化CUDA上下文
torch.cuda.init()
getram()  # 此处的RAM涨幅就是CUDA上下文的固定开销
# 再创建一个极小的CUDA张量
tiny_tensor = torch.zeros(1, device='cuda')
getram()  # 此时RAM几乎不会上涨，因为上下文已经初始化完成

针对你的场景的优化方案 #

如果你的主机RAM资源紧张，或者想避免这种突发的RAM上涨，可以试试以下几种方法：

1. 避免预先分配大CUDA张量，改用列表收集后拼接

这种方式不会触发提前的大内存开销，而是逐步处理每个图片，最后在GPU上拼接成完整数据集，主机RAM占用会更平缓：

def load_dataset(dir, filenames):
    dataset_list = []
    getram()
    for i, filename in enumerate(filenames):
        f = read_image(f"{dir}/{filename}")
        if f.shape[0] != 3:
            print(f"Skipping {filename}: wrong channel count")
            continue
        # 直接转成CUDA张量后加入列表，non_blocking=True可异步传输提升速度
        dataset_list.append(f.to(device, non_blocking=True))
    getram()
    # 最后在GPU上拼接成大张量
    dataset = torch.stack(dataset_list, dim=0)
    getram()
    return dataset

2. 用torch.empty()替代torch.zeros()（如果不需要初始值）

如果你不需要把张量初始化为0，torch.empty()会更高效——它只会在GPU上分配内存，不会执行填充0值的操作，减少了不必要的底层开销：

dataset = torch.empty((len(filenames), 3, 256, 256), device=device)

不过这仍会触发CUDA上下文初始化（如果还未初始化的话），第一次创建时还是会有RAM上涨，但后续操作的开销会更小。

3. 接受CUDA上下文的固定开销（RAM充足时）

如果你的主机RAM足够，其实不需要过度担心这部分开销——CUDA上下文是进程级别的，初始化一次后，后续所有CUDA操作都可以复用，不会再重复占用新的RAM。你看到最后两次RAM稳定在678.28，说明循环结束后临时的图片数据已经被正常回收，这是健康的状态。

为什么gc.collect()没用？ #

Python的gc.collect()只能回收Python层面的对象（比如你用read_image加载的主机端图片张量），但CUDA上下文是由CUDA运行时管理的进程级资源，PyTorch的CUDA张量元数据也由PyTorch的C++后端管理，这些都不在Python垃圾回收的范围内，所以gc.collect()对这部分内存完全无效。

备注：内容来源于stack exchange，提问作者Nex