如何降低PyTorch中CUDA初始化的内存占用（Jetson TX2场景）

针对你的Jetson TX2（GPU/CPU共享内存）、PyTorch v1.10.0、JetPack 4.6.1、CUDA 10.2环境，以下是几个可降低CUDA初始化内存占用的方案，以牺牲部分性能换取GPU可用空间：

1. 禁用CUDA内核预加载 #

PyTorch默认预加载大量优化内核是初始化内存占用高的核心原因，可通过环境变量强制关闭预加载，改为动态按需加载：

export PYTORCH_JIT=0
export CUDA_MODULE_LOADING=LAZY

• PYTORCH_JIT=0：禁用JIT编译内核的预加载，削减初始化阶段内存占用
• CUDA_MODULE_LOADING=LAZY：让CUDA仅在实际执行对应操作时才加载内核，而非一次性全部加载

2. 限制CUDA上下文内存占比 #

Jetson的CUDA上下文默认占用较多共享内存，可通过环境变量或代码限制内存分配上限：

• 环境变量设置：

export CUDA_MEMORY_LIMIT=1073741824  # 限制为1GB，可根据剩余内存调整数值

• PyTorch代码内设置：

import torch
# 限制当前进程使用50%的GPU共享内存，可按需调整比例
torch.cuda.set_per_process_memory_fraction(0.5, device=0)

3. 优化模型与训练流程 #

针对你要训练的小型模型（Conv2D、池化、Dense），可进一步压缩内存使用：

• 用torch.no_grad()包裹无梯度需求的代码段（如推理），避免存储梯度信息
• 启用FP16混合精度训练，在精度损失可控的前提下减少内存占用：

from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()

for epoch in range(epochs):
    for inputs, labels in dataloader:
        optimizer.zero_grad()
        with autocast():
            outputs = model(inputs)
            loss = criterion(outputs, labels)
        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

• 调小批量大小（batch_size），使用最小可行批次降低单步内存负载

4. 清理系统共享内存缓存 #

Jetson的共享内存可能被系统缓存占用，可手动清理临时释放空间：

sudo sync && sudo sysctl -w vm.drop_caches=3

该操作仅清理系统页缓存、目录项等，不影响正在运行的进程。

5. 编译PyTorch时裁剪冗余内核（进阶方案） #

若上述方法仍无法满足需求，可重新编译PyTorch，仅保留所需算子内核：

• 下载PyTorch v1.10.0源码
• 编译时指定仅保留Jetson TX2对应的CUDA架构（sm_62），并关闭不必要的优化：

export TORCH_CUDA_ARCH_LIST="6.2"
export USE_CUDNN=0  # 若可接受卷积性能损失，关闭CuDNN
python setup.py install

此方法能最大程度削减内核数量，但操作复杂度较高，需熟悉编译流程。

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