本地Hugging Face 模型导出为 ONNX

本地Hugging Face 模型导出为 ONNX

安装依赖

pip install optimum[exporters]
# 如果你需要 GPU 支持
pip install optimum[onnxruntime-gpu]

optimum-cli

# 执行转换命令
optimum-cli export onnx --model local_model_path --task task_type  save_onnx_path
e.g:
optimum-cli export onnx --model /home/models/deepseek-ai/DeepSeek-R1-0528-Qwen3-8B --task text-generation  /home/models/deepseek-ai/DeepSeek-R1-0528-Qwen3-8B_ONNX/
##特别说明：
task类型需要根据不同的模型文件指定
## 其他参数说明
--framework pt   pt：表示 PyTorch（torch）；tf：表示 TensorFlow
--sequence_length 2048  指定序列长度（token数量）;通常支持 32K 上下文，但导出时可设置较小值以节省内存
--opset 18 指定 ONNX 算子集的版本;11-13：基础算子，兼容性好;14-16：支持更多优化;17+：支持最新的算子
--device DEVICE       The device to use to do the export. Defaults to "cpu"
--dtype {fp32,fp16,bf16} 
--batch_size BATCH_SIZE  指定批处理大小,--batch_size 1 表示固定 batch_size=1;--batch_size 1 2 4 8 表示支持多个 batch_size

ONNX模型加载验证

import onnxruntime as ort
import numpy as np

model_path = "/home/models/deepseek-ai/DeepSeek-R1-0528-Qwen3-8B_ONNX/model.onnx"

# 方式1：最简单加载
#session = ort.InferenceSession(model_path)

# 方式2：指定执行提供者（CPU或GPU）
# 查看可用提供者
print("可用提供者:", ort.get_available_providers())
# ['CPUExecutionProvider', 'CUDAExecutionProvider', ...]

# 指定使用CPU
#session_cpu = ort.InferenceSession(model_path, providers=['CPUExecutionProvider'])

# 指定使用GPU（如果有）
session_gpu = ort.InferenceSession(model_path, providers=['CUDAExecutionProvider'])


# 获取输入信息
for input_info in session_gpu.get_inputs():
    print(f"输入名称: {input_info.name}")
    print(f"输入形状: {input_info.shape}")  # 例如: ['batch_size', 'sequence_length']
    print(f"输入类型: {input_info.type}")
    print("-" * 40)

# 获取输出信息
for output_info in session_gpu.get_outputs():
    print(f"输出名称: {output_info.name}")
    print(f"输出形状: {output_info.shape}")
    print(f"输出类型: {output_info.type}")

# NLP模型（如你遇到的DeepSeek）
def prepare_text_input(token_ids, attention_mask=None):
    """准备文本模型输入"""
    input_ids = np.array([token_ids], dtype=np.int64)

    if attention_mask is None:
        attention_mask = np.ones_like(input_ids, dtype=np.int64)

    # 对于需要position_ids的模型
    seq_length = input_ids.shape[1]
    position_ids = np.arange(seq_length, dtype=np.int64).reshape(1, -1)

    return {
        'input_ids': input_ids,
        'attention_mask': attention_mask,
        'position_ids': position_ids
    }

# 方法1：使用run()方法
# tokenize后的输入
inputs = prepare_text_input([101, 2023, 2003, 1037,2317,3185])
outputs = session_gpu.run(None, inputs)  # None表示获取所有输出
print(f'outputs : {outputs}')

# 方法2：指定输出名称
output_names = [output.name for output in session_gpu.get_outputs()]
outputs = session_gpu.run(output_names, inputs)
print(f'outputs : {outputs}')