Ubuntu20.04下PyTorch转换YOLOv7为.ts时CUDA内存不足求助

我尝试用以下代码将YOLOv7模型通过PyTorch转换为.ts文件：

gc.collect()
torch.cuda.empty_cache()
traced_model = torch.jit.trace(model, torch.empty([10, 3, 768, 1280]).half().to(self.device))
trt_script_module = torch_tensorrt.compile(
            traced_model,
            inputs = [
            torch_tensorrt.Input(
                min_shape=[10, 3, 768, 1280],
                opt_shape=[10, 3, 768, 1280],
                max_shape=[10, 3, 768, 1280],
                dtype=torch.half,
            )
            ], enabled_precisions={torch.half}
        )
torch.jit.save(trt_script_module, "yolo_trt_script.ts")

执行torch.jit.trace时出现运行时错误：

RuntimeError: CUDA out of memory. Tried to allocate 38.00 MiB (GPU 0; 7.79 GiB total capacity; 6.52 GiB already allocated; 34.69 MiB free; 6.61 GiB reserved in total by PyTorch) If reserved memory is >> allocated memory try setting max_split_size_mb to avoid fragmentation. See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF

环境详情 #

• 操作系统：Ubuntu 20.04 LTS
• GPU：NVIDIA GeForce RTX 2070（8GB）
• 驱动版本：495.29.05
• Torch版本：1.11.0+cu115
• Torch-Tensorrt版本：1.1.0
• CUDA：11.5

疑问 #

1. 同一台电脑在Windows10环境下能生成.onnx文件，为什么Ubuntu环境下不行（我也试过生成.onnx）？
2. 有没有办法优化转换时的GPU负载？

已尝试的解决方法 #

• 关闭其他进程
• 升级Torch版本
• 清空CUDA缓存

解答 #

关于Windows和Ubuntu下生成ONNX的差异 #

Windows和Ubuntu的CUDA内存管理机制存在核心差异：

• Windows的WDDM驱动支持GPU内存自动分页到系统内存，当GPU内存不足时会动态交换，变相扩大可用空间；而Ubuntu默认使用的NVIDIA驱动内存策略更保守，不会自动做跨内存域的交换，相同模型下更易触发OOM。
• Ubuntu环境中后台进程的GPU消耗更高，比如桌面环境的GPU加速服务、系统级图形渲染进程，会占用额外显存；Windows下这类后台消耗通常更低。
• 不同系统下PyTorch的内存分配逻辑细节有区别，Ubuntu版本的PyTorch可能预留了更多显存用于内部管理，导致实际可用空间被压缩。

优化GPU负载的方法 #

1. 降低trace时的batch size
   当前使用的batch size为10，尝试改为1或2，转换完成后不影响后续推理时使用大batch size：

model.eval()
with torch.no_grad():
    traced_model = torch.jit.trace(model, torch.empty([1, 3, 768, 1280]).half().to(self.device))
# 同步修改torch_tensorrt.Input的shape参数为[1, 3, 768, 1280]

2. 优化CUDA内存分配策略
   设置环境变量减少内存碎片，在代码开头添加：

import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'max_split_size_mb:128'

或者在终端执行转换前先运行：

export PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:128

3. 启用模型eval模式+无梯度上下文
   trace前切换模型到eval模式，关闭训练相关层（BatchNorm、Dropout等）的动态计算，同时用torch.no_grad()禁用梯度追踪，大幅减少内存占用：

model.eval()
with torch.no_grad():
    traced_model = torch.jit.trace(model, torch.empty([1, 3, 768, 1280]).half().to(self.device))

4. 分步转换：先导出ONNX再转TS
   避免PyTorch直接编译TensorRT时的内存叠加，拆分步骤：

# 第一步：导出ONNX
model.eval()
with torch.no_grad():
    torch.onnx.export(model, torch.empty([1, 3, 768, 1280]).half().to(self.device), "yolov7.onnx", opset_version=13)

# 第二步：用TensorRT解析ONNX并转TS
import tensorrt as trt
import torch_tensorrt

builder = trt.Builder(trt.Logger(trt.Logger.WARNING))
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
parser = trt.OnnxParser(network, builder.logger)

with open("yolov7.onnx", "rb") as f:
    parser.parse(f.read())

config = builder.create_builder_config()
config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, 1 << 30)  # 分配1GB工作空间
config.set_flag(trt.BuilderFlag.FP16)

engine = builder.build_engine(network, config)
trt_script_module = torch_tensorrt.ts.export_to_script_module(engine)
torch.jit.save(trt_script_module, "yolo_trt_script.ts")

5. 减少后台GPU消耗
   在Ubuntu下可以切换到命令行模式执行转换，或者关闭桌面环境的GPU特效（如Gnome的窗口动画、 compositor），释放后台占用的显存。

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