ONNX Format Guide

    开放神经网络交换格式,用于跨平台推理

    Model Weights

    Specification

    ONNX(Open Neural Network Exchange,开放神经网络交换)是一种用于表示机器学习模型的开放标准格式。该格式最初由 Microsoft 和 Facebook 于 2017 年联合开发,现由 Linux Foundation 管理。ONNX 定义了一组通用的算子、数据类型和计算图格式,使在一个框架中训练的模型能够在另一个框架中部署。该格式使用 Protocol Buffers(protobuf)进行序列化,支持包括卷积神经网络、循环网络、Transformer 和经典机器学习模型在内的广泛模型架构。

    ONNX 模型表示为有向无环图(DAG),其中节点代表运算操作(卷积、矩阵乘法、激活函数等),边代表在操作之间流动的张量,图具有定义的输入和输出。ONNX 算子集(opset)采用版本管理,允许模型指定所需的算子定义版本。截至 opset 21,ONNX 定义了超过 200 个算子,涵盖神经网络运算、数学函数、张量操作和控制流。

    该格式包含可扩展的元数据系统,用于模型文档、训练信息和自定义属性。ONNX 还支持量化数据类型(INT8、UINT8)和混合精度表示,支持对资源受限设备上的模型部署进行优化。ONNX Model Zoo 提供了一系列 ONNX 格式的预训练模型,涵盖图像分类、目标检测、NLP、语音识别等常见任务。

    When to Use ONNX

    当你需要跨平台模型部署时——在 PyTorch 或 TensorFlow 中训练并在不同硬件或运行时环境中部署——ONNX 是正确的选择。ONNX Runtime(ORT)提供跨 CPU(x86、ARM)、GPU(NVIDIA、AMD、Intel)、NPU 和专用加速器的优化推理。如果你的部署目标包括 Windows 应用程序(通过 DirectML)、移动设备(通过 ONNX Runtime Mobile)、Web 浏览器(通过 ONNX Runtime Web/WebAssembly)或边缘设备,ONNX 提供了适用于所有平台的统一模型格式。

    当推理性能至关重要且你希望利用 ONNX Runtime 丰富的图优化功能时,应选择 ONNX,这些优化包括算子融合、常量折叠、布局优化和量化。ONNX Runtime 在各种模型架构和硬件平台上始终位列最快的推理引擎之一。它在用于分类、NER 和嵌入任务的 Transformer 编码器模型(BERT、RoBERTa、DeBERTa)方面表现尤为出色。

    ONNX 不太适合超大型生成式语言模型(超过 7B 参数),此类场景下 vLLM、TensorRT-LLM 或 llama.cpp 等专用推理引擎通过连续批处理、PagedAttention 和投机解码等技术能提供更好的性能。ONNX 也存在转换差距——并非所有 PyTorch 操作都被 ONNX 导出器支持,特别是动态控制流和某些自定义算子,导出时可能需要变通处理。

    Schema / Structure

    protobuf
    // ONNX Model structure (simplified from onnx.proto3)
    message ModelProto {
      int64 ir_version = 1;           // IR version (currently 9)
      repeated OperatorSetIdProto opset_import = 8;
      string producer_name = 2;       // e.g., "pytorch"
      string producer_version = 3;
      string domain = 4;
      int64 model_version = 5;
      string doc_string = 6;
      GraphProto graph = 7;           // The computation graph
      repeated StringStringEntryProto metadata_props = 14;
    }
    
    message GraphProto {
      string name = 1;
      repeated NodeProto node = 2;      // Operations in the graph
      repeated TensorProto initializer = 5; // Pretrained weights
      repeated ValueInfoProto input = 11;
      repeated ValueInfoProto output = 12;
    }
    
    message NodeProto {
      repeated string input = 1;
      repeated string output = 2;
      string op_type = 4;              // e.g., "Conv", "MatMul", "Relu"
      repeated AttributeProto attribute = 5;
    }
    简化的 ONNX protobuf 模式,展示模型、图和节点结构

    Example Data

    python
    import torch
    import onnx
    import onnxruntime as ort
    from transformers import AutoTokenizer, AutoModelForSequenceClassification
    
    # Export a HuggingFace model to ONNX
    model_name = "distilbert-base-uncased-finetuned-sst-2-english"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(model_name)
    model.eval()
    
    dummy_input = tokenizer("This is a test", return_tensors="pt")
    torch.onnx.export(
        model,
        (dummy_input["input_ids"], dummy_input["attention_mask"]),
        "sentiment_model.onnx",
        input_names=["input_ids", "attention_mask"],
        output_names=["logits"],
        dynamic_axes={"input_ids": {0: "batch", 1: "seq"},
                      "attention_mask": {0: "batch", 1: "seq"},
                      "logits": {0: "batch"}},
        opset_version=17,
    )
    
    # Run inference with ONNX Runtime
    session = ort.InferenceSession("sentiment_model.onnx")
    inputs = tokenizer("Great product, highly recommend!", return_tensors="np")
    outputs = session.run(None, dict(inputs))
    print(f"Logits: {outputs[0]}")  # [[negative_score, positive_score]]
    将 Hugging Face 情感模型导出为 ONNX 并使用 ONNX Runtime 进行推理

    Ertas Support

    Ertas Studio 支持将 ONNX 作为训练模型的导出目标,实现跨 ONNX Runtime 支持的多种运行时环境的部署。通过 Ertas 云训练流水线训练的模型可以导出为 ONNX 格式,并根据目标部署平台(无论是边缘设备上的 CPU 推理、数据中心的 GPU 推理还是通过 WebAssembly 的浏览器端推理)应用相应的优化。

    Related Resources

    Ship AI that runs on your users' devices.

    Early bird pricing starts at $14.50/mo — locked in for life. Plans for builders and agencies.