本地Hugging Face 模型導(dǎo)出為 ONNX

安裝依賴

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

optimum-cli

# 執(zhí)行轉(zhuǎn)換命令
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類型需要根據(jù)不同的模型文件指定
## 其他參數(shù)說明
--framework pt   pt：表示 PyTorch（torch）；tf：表示 TensorFlow
--sequence_length 2048  指定序列長(zhǎng)度（token數(shù)量）;通常支持 32K 上下文，但導(dǎo)出時(shí)可設(shè)置較小值以節(jié)省內(nèi)存
--opset 18 指定 ONNX 算子集的版本;11-13：基礎(chǔ)算子，兼容性好;14-16：支持更多優(yōu)化;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 表示支持多個(gè) batch_size

ONNX模型加載驗(yàn)證

import onnxruntime as ort
import numpy as np

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

# 方式1：最簡(jiǎn)單加載
#session = ort.InferenceSession(model_path)

# 方式2：指定執(zhí)行提供者（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):
    """準(zhǔn)備文本模型輸入"""
    input_ids = np.array([token_ids], dtype=np.int64)

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

    # 對(duì)于需要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}')