Python API 概述

完整的 API 請參閱API 參考。

載入 ONNX 模型

import onnx

# onnx_model is an in-memory ModelProto
onnx_model = onnx.load("path/to/the/model.onnx")

可執行的 IPython 筆記本

• load_model.ipynb

使用外部資料載入 ONNX 模型

• [預設] 如果外部資料與模型位於同一個目錄下，只需使用 onnx.load()

import onnx

onnx_model = onnx.load("path/to/the/model.onnx")

• 如果外部資料位於另一個目錄下，請使用 load_external_data_for_model() 來指定目錄路徑，並在使用 onnx.load() 後載入

import onnx
from onnx.external_data_helper import load_external_data_for_model

onnx_model = onnx.load("path/to/the/model.onnx", load_external_data=False)
load_external_data_for_model(onnx_model, "data/directory/path/")
# Then the onnx_model has loaded the external data from the specific directory

將 ONNX 模型轉換為外部資料

from onnx.external_data_helper import convert_model_to_external_data

# onnx_model is an in-memory ModelProto
onnx_model = ...
convert_model_to_external_data(onnx_model, all_tensors_to_one_file=True, location="filename", size_threshold=1024, convert_attribute=False)
# Then the onnx_model has converted raw data as external data
# Must be followed by save

儲存 ONNX 模型

import onnx

# onnx_model is an in-memory ModelProto
onnx_model = ...

# Save the ONNX model
onnx.save(onnx_model, "path/to/the/model.onnx")

可執行的 IPython 筆記本

• save_model.ipynb

將 ONNX 模型轉換並儲存為外部資料

import onnx

# onnx_model is an in-memory ModelProto
onnx_model = ...
onnx.save_model(onnx_model, "path/to/save/the/model.onnx", save_as_external_data=True, all_tensors_to_one_file=True, location="filename", size_threshold=1024, convert_attribute=False)
# Then the onnx_model has converted raw data as external data and saved to specific directory

操作 TensorProto 和 Numpy 陣列

import numpy
import onnx
from onnx import numpy_helper

# Preprocessing: create a Numpy array
numpy_array = numpy.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=float)
print(f"Original Numpy array:\n{numpy_array}\n")

# Convert the Numpy array to a TensorProto
tensor = numpy_helper.from_array(numpy_array)
print(f"TensorProto:\n{tensor}")

# Convert the TensorProto to a Numpy array
new_array = numpy_helper.to_array(tensor)
print(f"After round trip, Numpy array:\n{new_array}\n")

# Save the TensorProto
with open("tensor.pb", "wb") as f:
    f.write(tensor.SerializeToString())

# Load a TensorProto
new_tensor = onnx.TensorProto()
with open("tensor.pb", "rb") as f:
    new_tensor.ParseFromString(f.read())
print(f"After saving and loading, new TensorProto:\n{new_tensor}")

from onnx import TensorProto, helper

# Conversion utilities for mapping attributes in ONNX IR
# The functions below are available after ONNX 1.13
np_dtype = helper.tensor_dtype_to_np_dtype(TensorProto.FLOAT)
print(f"The converted numpy dtype for {helper.tensor_dtype_to_string(TensorProto.FLOAT)} is {np_dtype}.")
storage_dtype = helper.tensor_dtype_to_storage_tensor_dtype(TensorProto.FLOAT)
print(f"The storage dtype for {helper.tensor_dtype_to_string(TensorProto.FLOAT)} is {helper.tensor_dtype_to_string(storage_dtype)}.")
field_name = helper.tensor_dtype_to_field(TensorProto.FLOAT)
print(f"The field name for {helper.tensor_dtype_to_string(TensorProto.FLOAT)} is {field_name}.")
tensor_dtype = helper.np_dtype_to_tensor_dtype(np_dtype)
print(f"The tensor data type for numpy dtype: {np_dtype} is {helper.tensor_dtype_to_string(tensor_dtype)}.")

for tensor_dtype in helper.get_all_tensor_dtypes():
    print(helper.tensor_dtype_to_string(tensor_dtype))

可執行的 IPython 筆記本

• np_array_tensorproto.ipynb

使用輔助函式建立 ONNX 模型

import onnx
from onnx import helper
from onnx import AttributeProto, TensorProto, GraphProto


# The protobuf definition can be found here:
# https://github.com/onnx/onnx/blob/main/onnx/onnx.proto


# Create one input (ValueInfoProto)
X = helper.make_tensor_value_info("X", TensorProto.FLOAT, [3, 2])
pads = helper.make_tensor_value_info("pads", TensorProto.FLOAT, [1, 4])

value = helper.make_tensor_value_info("value", AttributeProto.FLOAT, [1])


# Create one output (ValueInfoProto)
Y = helper.make_tensor_value_info("Y", TensorProto.FLOAT, [3, 4])

# Create a node (NodeProto) - This is based on Pad-11
node_def = helper.make_node(
    "Pad",                  # name
    ["X", "pads", "value"], # inputs
    ["Y"],                  # outputs
    mode="constant",        # attributes
)

# Create the graph (GraphProto)
graph_def = helper.make_graph(
    [node_def],        # nodes
    "test-model",      # name
    [X, pads, value],  # inputs
    [Y],               # outputs
)

# Create the model (ModelProto)
model_def = helper.make_model(graph_def, producer_name="onnx-example")

print(f"The model is:\n{model_def}")
onnx.checker.check_model(model_def)
print("The model is checked!")

可執行的 IPython 筆記本

• make_model.ipynb

• Protobufs.ipynb

用於對應 ONNX IR 中屬性的轉換工具

from onnx import TensorProto, helper

np_dtype = helper.tensor_dtype_to_np_dtype(TensorProto.FLOAT)
print(f"The converted numpy dtype for {helper.tensor_dtype_to_string(TensorProto.FLOAT)} is {np_dtype}.")

field_name = helper.tensor_dtype_to_field(TensorProto.FLOAT)
print(f"The field name for {helper.tensor_dtype_to_string(TensorProto.FLOAT)} is {field_name}.")

# There are other useful conversion utilities. Please checker onnx.helper

檢查 ONNX 模型

import onnx

# Preprocessing: load the ONNX model
model_path = "path/to/the/model.onnx"
onnx_model = onnx.load(model_path)

print(f"The model is:\n{onnx_model}")

# Check the model
try:
    onnx.checker.check_model(onnx_model)
except onnx.checker.ValidationError as e:
    print(f"The model is invalid: {e}")
else:
    print("The model is valid!")

可執行的 IPython 筆記本

• check_model.ipynb

檢查大於 2GB 的大型 ONNX 模型

目前的檢查器支援檢查具有外部資料的模型，但對於大於 2GB 的模型，請使用 onnx.checker 的模型路徑，且外部資料必須位於同一個目錄下。

import onnx

onnx.checker.check_model("path/to/the/model.onnx")
# onnx.checker.check_model(loaded_onnx_model) will fail if given >2GB model

在 ONNX 模型上執行形狀推斷

import onnx
from onnx import helper, shape_inference
from onnx import TensorProto


# Preprocessing: create a model with two nodes, Y"s shape is unknown
node1 = helper.make_node("Transpose", ["X"], ["Y"], perm=[1, 0, 2])
node2 = helper.make_node("Transpose", ["Y"], ["Z"], perm=[1, 0, 2])

graph = helper.make_graph(
    [node1, node2],
    "two-transposes",
    [helper.make_tensor_value_info("X", TensorProto.FLOAT, (2, 3, 4))],
    [helper.make_tensor_value_info("Z", TensorProto.FLOAT, (2, 3, 4))],
)

original_model = helper.make_model(graph, producer_name="onnx-examples")

# Check the model and print Y"s shape information
onnx.checker.check_model(original_model)
print(f"Before shape inference, the shape info of Y is:\n{original_model.graph.value_info}")

# Apply shape inference on the model
inferred_model = shape_inference.infer_shapes(original_model)

# Check the model and print Y"s shape information
onnx.checker.check_model(inferred_model)
print(f"After shape inference, the shape info of Y is:\n{inferred_model.graph.value_info}")

可執行的 IPython 筆記本

• shape_inference.ipynb

對大於 2GB 的大型 ONNX 模型進行形狀推斷

目前的 shape_inference 支援具有外部資料的模型，但對於大於 2GB 的模型，請使用 onnx.shape_inference.infer_shapes_path 的模型路徑，且外部資料必須位於同一個目錄下。您可以指定用於儲存推斷模型的輸出路徑；否則，預設輸出路徑與原始模型路徑相同。

import onnx

# output the inferred model to the original model path
onnx.shape_inference.infer_shapes_path("path/to/the/model.onnx")

# output the inferred model to the specified model path
onnx.shape_inference.infer_shapes_path("path/to/the/model.onnx", "output/inferred/model.onnx")

# inferred_model = onnx.shape_inference.infer_shapes(loaded_onnx_model) will fail if given >2GB model

在 ONNX 函式上執行類型推斷

import onnx
import onnx.helper
import onnx.parser
import onnx.shape_inference

function_text = """
    <opset_import: [ "" : 18 ], domain: "local">
    CastTo <dtype> (x) => (y) {
        y = Cast <to : int = @dtype> (x)
    }
"""
function = onnx.parser.parse_function(function_text)

# The function above has one input-parameter x, and one attribute-parameter dtype.
# To apply type-and-shape-inference to this function, we must supply the type of
# input-parameter and an attribute value for the attribute-parameter as below:

float_type_ = onnx.helper.make_tensor_type_proto(1, None)
dtype_6 = onnx.helper.make_attribute("dtype", 6)
result = onnx.shape_inference.infer_function_output_types(
    function, [float_type_], [dtype_6]
)
print(result) # a list containing the (single) output type

在預設網域（“”/“ai.onnx”）中轉換 ONNX 模型的版本

import onnx
from onnx import version_converter, helper

# Preprocessing: load the model to be converted.
model_path = "path/to/the/model.onnx"
original_model = onnx.load(model_path)

print(f"The model before conversion:\n{original_model}")

# A full list of supported adapters can be found here:
# https://github.com/onnx/onnx/blob/main/onnx/version_converter.py#L21
# Apply the version conversion on the original model
converted_model = version_converter.convert_version(original_model, <int target_version>)

print(f"The model after conversion:\n{converted_model}")

公用程式函式

使用輸入輸出張量名稱提取子模型

函式 extract_model() 從 ONNX 模型中提取子模型。子模型由輸入和輸出張量的名稱確切地定義。

import onnx

input_path = "path/to/the/original/model.onnx"
output_path = "path/to/save/the/extracted/model.onnx"
input_names = ["input_0", "input_1", "input_2"]
output_names = ["output_0", "output_1"]

onnx.utils.extract_model(input_path, output_path, input_names, output_names)

注意：對於控制流程運算符，例如 If 和 Loop，由輸入和輸出張量定義的子模型邊界不應切穿連線到主圖形的子圖，作為這些運算符的屬性。

ONNX Compose

onnx.compose 模組提供工具來建立組合模型。

onnx.compose.merge_models 可用於合併兩個模型，方法是將第一個模型的某些輸出與第二個模型的輸入連接。預設情況下，io_map 參數中不存在的輸入/輸出將保留為組合模型的輸入/輸出。

在此範例中，我們透過將第一個模型的每個輸出連接到第二個模型的輸入來合併兩個模型。產生的模型將具有與第一個模型相同的輸入，以及與第二個模型相同的輸出

import onnx

model1 = onnx.load("path/to/model1.onnx")
# agraph (float[N] A, float[N] B) => (float[N] C, float[N] D)
#   {
#      C = Add(A, B)
#      D = Sub(A, B)
#   }

model2 = onnx.load("path/to/model2.onnx")
#   agraph (float[N] X, float[N] Y) => (float[N] Z)
#   {
#      Z = Mul(X, Y)
#   }

combined_model = onnx.compose.merge_models(
    model1, model2,
    io_map=[("C", "X"), ("D", "Y")]
)

此外，使用者可以指定要包含在組合模型中的 inputs/outputs 清單，有效地捨棄不貢獻於組合模型輸出的圖形部分。在以下範例中，我們只將第一個模型中的兩個輸出之一連接到第二個模型中的兩個輸入。透過明確指定組合模型的輸出，我們正在捨棄第一個模型中未使用的輸出，以及圖形的相關部分

import onnx

# Default case. Include all outputs in the combined model
combined_model = onnx.compose.merge_models(
    model1, model2,
    io_map=[("C", "X"), ("C", "Y")],
)  # outputs: "D", "Z"

# Explicit outputs. "Y" output and the Sub node are not present in the combined model
combined_model = onnx.compose.merge_models(
    model1, model2,
    io_map=[("C", "X"), ("C", "Y")],
    outputs=["Z"],
)  # outputs: "Z"

onnx.compose.add_prefix 允許您在模型中的名稱中新增前置詞，以避免在合併時發生名稱衝突。預設情況下，它會重新命名圖形中的所有名稱：輸入、輸出、邊緣、節點、初始值設定項、稀疏初始值設定項和值資訊。

import onnx

model = onnx.load("path/to/the/model.onnx")
# model - outputs: ["out0", "out1"], inputs: ["in0", "in1"]

new_model = onnx.compose.add_prefix(model, prefix="m1/")
# new_model - outputs: ["m1/out0", "m1/out1"], inputs: ["m1/in0", "m1/in1"]

# Can also be run in-place
onnx.compose.add_prefix(model, prefix="m1/", inplace=True)

onnx.compose.expand_out_dim 可用於連接預期維度數不同的模型，方法是插入範圍為 1 的維度。當將產生樣本的模型與處理樣本批次的模型組合時，這會很有用。

import onnx

# outputs: "out0", shape=[200, 200, 3]
model1 = onnx.load("path/to/the/model1.onnx")

# outputs: "in0", shape=[N, 200, 200, 3]
model2 = onnx.load("path/to/the/model2.onnx")

# outputs: "out0", shape=[1, 200, 200, 3]
new_model1 = onnx.compose.expand_out_dims(model1, dim_idx=0)

# Models can now be merged
combined_model = onnx.compose.merge_models(
    new_model1, model2, io_map=[("out0", "in0")]
)

# Can also be run in-place
onnx.compose.expand_out_dims(model1, dim_idx=0, inplace=True)

工具

使用可變長度更新模型的輸入輸出維度大小

函式 update_inputs_outputs_dims 會將模型的輸入和輸出維度更新為參數中提供的值。您可以使用 dim_param 提供靜態和動態維度大小。有關靜態和動態維度大小的詳細資訊，請查看 張量形狀。

函式會在輸入/輸出大小更新後執行模型檢查器。

import onnx
from onnx.tools import update_model_dims

model = onnx.load("path/to/the/model.onnx")
# Here both "seq", "batch" and -1 are dynamic using dim_param.
variable_length_model = update_model_dims.update_inputs_outputs_dims(model, {"input_name": ["seq", "batch", 3, -1]}, {"output_name": ["seq", "batch", 1, -1]})

ONNX 剖析器

函式 onnx.parser.parse_model 和 onnx.parser.parse_graph 可用於從如下所示的文字表示建立 ONNX 模型或圖形。有關語言語法的更多詳細資訊，請參閱 語言語法。

input = """
   agraph (float[N, 128] X, float[128, 10] W, float[10] B) => (float[N, 10] C)
   {
        T = MatMul(X, W)
        S = Add(T, B)
        C = Softmax(S)
   }
"""
graph = onnx.parser.parse_graph(input)

input = """
   <
     ir_version: 7,
     opset_import: ["" : 10]
   >
   agraph (float[N, 128] X, float[128, 10] W, float[10] B) => (float[N, 10] C)
   {
      T = MatMul(X, W)
      S = Add(T, B)
      C = Softmax(S)
   }
"""
model = onnx.parser.parse_model(input)

ONNX Inliner

函式 onnx.inliner.inline_local_functions 和 inline_selected_functions 可用於將 ONNX 模型中的模型內部函式內聯展開。特別是，inline_local_functions 可用於產生一個不含函式的模型（適用於不處理或不支援函式的後端）。另一方面，inline_selected_functions 可用於內聯展開選定的函式。目前尚不支援內聯展開作為函式的 ONNX 標準運算（也稱為綱要定義函式）。

import onnx
import onnx.inliner

model = onnx.load("path/to/the/model.onnx")
inlined = onnx.inliner.inline_local_functions(model)
onnx.save("path/to/the/inlinedmodel.onnx")