<think>我们正在构建一个用于运行大型ONNX模型的完整程序示例，同时考虑到了大模型（超过2GB）的处理。我们将使用ONNX Runtime，并利用其外部数据功能来避免protobuf的2GB限制。 关键点： 1. 模型准备：假设我们有一个大模型（超过2GB），并且我们已经将其转换为带有外部数据的格式（即权重存储在外部文件中）。 2. 推理程序：我们将编写一个Python脚本，使用ONNX Runtime加载和运行这个模型。 3. 内存优化：我们将使用内存映射等技巧来减少内存占用。 步骤： 1. 加载模型：使用ONNX Runtime的`InferenceSession`，并配置会话选项以使用设备分配器进行初始化的内存映射。 2. 提供输入：准备输入数据（例如numpy数组）。 3. 运行推理：执行模型推理并获取输出。 注意：我们假设模型已经以外部数据格式保存（即一个小的.onnx文件和一个或多个存储权重的外部文件，例如.bin文件）。如果用户已有大模型，可以使用之前介绍的`convert_model_to_external_data`方法转换。 以下是一个完整的示例代码，包括如何准备输入、运行推理以及处理大模型： </think> # 完整超大模型 ONNX Runtime 运行程序示例 下面是一个完整的ONNX Runtime运行程序，专门设计用于处理超过2GB的大型模型。我通过外部数据、内存映射与量化等多种技术来实现高效的大模型推理。 ```python import onnxruntime as ort import numpy as np import psutil # 用于内存监控 import os import time from onnxruntime.quantization import quantize_dynamic, QuantType def log_memory_usage(stage): """记录内存使用情况""" process = psutil.Process(os.getpid()) mem_info = process.memory_info() print(f"[{stage}] 内存使用: {mem_info.rss / (1024 ** 3):.2f} GB") class LargeModelONNXRuntime: def __init__(self, model_path, provider="CUDAExecutionProvider", enable_directml=False): """ 初始化大模型推理引擎 参数: model_path: ONNX模型文件路径(.onnx或带外部数据的.onnx) provider: 默认执行提供程序(CUDAExecutionProvider/CPUExecutionProvider) enable_directml: 是否启用DirectML(Windows GPU加速) """ self.model_path = model_path self.provider = provider # 配置会话选项 self.sess_options = ort.SessionOptions() # 启用优化 self.sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL self.sess_options.enable_cpu_mem_arena = False # 减少内存占用 # 关键配置: 允许使用设备内存映射 self.sess_options.add_session_config_entry( "session.use_device_allocator_for_initializers", "1" ) # 配置执行提供程序 providers = [] if enable_directml: providers = ["DmlExecutionProvider"] elif provider == "CUDAExecutionProvider" and ort.get_device() == "GPU": # 对于大模型，配置CUDA优化参数 providers = [("CUDAExecutionProvider", { "device_id": 0, "arena_extend_strategy": "kSameAsRequested", "cuda_mem_limit": int(80 * 1024 ** 3), # 限制分配80GB显存 "do_copy_in_default_stream": True, })] else: providers = [provider] # 默认为CPU # 创建推理会话 print(f"正在加载模型: {model_path}...") start_time = time.time() self.session = ort.InferenceSession( model_path, sess_options=self.sess_options, providers=providers ) load_time = time.time() - start_time print(f"✓ 模型加载完成! 耗时: {load_time:.2f}秒") # 记录输入输出信息 self.input_details = {input.name: input for input in self.session.get_inputs()} self.output_names = [output.name for output in self.session.get_outputs()] # 测试内存使用情况 log_memory_usage("模型初始化后") @staticmethod def quantize_model(input_model, output_model): """ 量化模型以减小大小 参数: input_model: 输入模型路径 output_model: 输出模型路径 """ print(f"开始量化模型: {input_model} -> {output_model}") quantize_dynamic( input_model, output_model, weight_type=QuantType.QInt8, optimize_model=True, use_external_data_format=True, # 必须使用外部数据格式 ) print("✓ 模型量化完成!") @staticmethod def check_model_shape(model_path): """检查模型输入输出形状""" try: session = ort.InferenceSession(model_path) for i, input in enumerate(session.get_inputs()): print(f"输入 {i+1}: {input.name}, 形状: {input.shape}, 类型: {input.type}") for o, output in enumerate(session.get_outputs()): print(f"输出 {o+1}: {output.name}, 形状: {output.shape}, 类型: {output.type}") return True except Exception as e: print(f"模型检查失败: {e}") return False def predict(self, input_data, warmup=1, repeats=1): """ 执行模型推理 参数: input_data: 字典格式的输入数据 {输入名: 数据} warmup: 预热次数 repeats: 正式推理次数 返回: 输出字典和推理时间统计 """ # 类型转换和格式验证 validated_inputs = {} for name, data in input_data.items(): # 自动转换为float32并验证形状 dt = self.input_details[name].type if 'int' in dt: if not isinstance(data, np.integer): data = np.array(data).astype(np.int64) else: if not isinstance(data, np.float32): data = np.array(data).astype(np.float32) input_shape = self.input_details[name].shape if len(input_shape) > 0 and input_shape != data.shape: print(f"⚠️ 输入形状不匹配: {name}需要{input_shape}但提供了{data.shape}") validated_inputs[name] = data # 预热运行 print(f"正在预热推理({warmup}次)...") for _ in range(warmup): outputs = self.session.run( self.output_names, validated_inputs ) # 正式推理并计时 print(f"进行推理({repeats}次)...") start_time = time.perf_counter() for _ in range(repeats): outputs = self.session.run( self.output_names, validated_inputs ) inference_time = (time.perf_counter() - start_time) / repeats log_memory_usage("推理后") # 转换为输出字典 output_dict = {name: outputs[i] for i, name in enumerate(self.output_names)} # 返回结果和性能数据 return { "outputs": output_dict, "performance": { "batch_size": validated_inputs[list(validated_inputs.keys())[0]].shape[0], "inference_time": inference_time, "throughput": 1 / inference_time if inference_time > 0 else float('inf') } } if __name__ == "__main__": # ========== 使用示例 ========== model_path = "path/to/your/large_model.onnx" # 替换为你的模型路径 # 检查模型(执行前推荐) print("正在检查模型输入输出...") LargeModelONNXRuntime.check_model_shape(model_path) # 创建推理引擎实例 # provider = "CPUExecutionProvider" # 可选CPU engine = LargeModelONNXRuntime( model_path, provider="CUDAExecutionProvider" # 使用GPU加速 ) # 创建模拟输入数据 - 根据您的模型调整 batch_size = 1 input_data = { "input": np.random.randn(batch_size, 1, 5).astype(np.float32) # 5维输入 } # 执行推理 result = engine.predict(input_data, warmup=3, repeats=5) # 打印结果和性能 print("\n推理结果:") for name, value in result["outputs"].items(): print(f"输出 '{name}': 形状={value.shape}, 类型={value.dtype}") print(f"前5个值: {value.flatten()[:5]}") perf = result["performance"] print(f"\n性能统计:") print(f"单次推理时间: {perf['inference_time']:.4f}秒") print(f"处理速度: {perf['throughput']:.2f}批次/秒") # 模型量化示例（单独执行） # LargeModelONNXRuntime.quantize_model( # "original_large_model.onnx", # "quantized_large_model.onnx" # ) ``` ## 程序功能详解 ### 核心组件 1. **推理引擎类** `LargeModelONNXRuntime` - 自动配置最佳会话选项 - 支持CPU、CUDA和DirectML多种执行提供者 - 内置外存分配器减少内存占用[^1] 2. **内存优化功能** ```python # 启用零拷贝内存映射 self.sess_options.add_session_config_entry( "session.use_device_allocator_for_initializers", "1" ) ``` - 当模型初始值设定项超过1GB时，自动使用内存映射文件 3. **量化压缩功能** ```python @staticmethod def quantize_model(input_model, output_model): quantize_dynamic( input_model, output_model, weight_type=QuantType.QInt8, use_external_data_format=True ) ``` - 采用INT8动态量化减小模型体积[^1] ### 模型加载处理流程 ```mermaid graph TD A[初始化配置] -->|设置CUDA显存限制| B[创建推理会话] B --> C[解析输入输出元数据] C -->|内存映射| D{大于1GB的权重?} D -->|是| E[延迟加载磁盘数据] D -->|否| F[加载到显存] E --> G[准备模型执行图] F --> G G --> H[推理执行器就绪] ``` ### 内存监控原理 程序使用内置`log_memory_usage()`监控内存： ```python def log_memory_usage(stage): process = psutil.Process(os.getpid()) mem_info = process.memory_info() print(f"[{stage}] 内存使用: {mem_info.rss / (1024 ** 3):.2f} GB") ``` ## 大模型优化策略 | 技术 | 效果 | 适用场景 | |------|------|----------| | **外部数据格式** | 突破2GB限制 | 所有超大型模型 | | **内存映射** | 减少内存占用50+% | GPU内存不足时 | | **INT8量化** | 模型压缩3-4倍 | FP32/Float16模型 | | **GPU内存限制** | 防止OOM崩溃 | 多任务共享GPU场景 | | **预热推理** | 稳定推理时间 | 需要稳定延迟的应用 | ## 典型应用场景 1. **大语言模型推理** ```python # 70B参数量的大模型示例 input_data = { "input_ids": np.array([[1, 2304, 399, 5700]]).astype(np.int64), "attention_mask": np.array([[1, 1, 1, 1]]).astype(np.int64) } result = engine.predict(input_data) ``` 2. **大规模视觉模型** ```python # 4K输入图像处理 input_data = { "pixel_values": np.random.rand(1, 3, 3840, 2160).astype(np.float32) } result = engine.predict(input_data) ``` ## 使用说明 1. 初始化模型时优先选择GPU： ```python engine = LargeModelONNXRuntime("large_model.onnx", provider="CUDAExecutionProvider") ``` 2. 若遇到GPU内存不足： ```python # 启用CPU+内存映射方案 engine = LargeModelONNXRuntime("huge_model.onnx", provider="CPUExecutionProvider") ``` 3. 监控内存使用确保不超限： ``` [模型初始化后] 内存使用: 0.72 GB [推理后] 内存使用: 1.25 GB ```