开始使用 ONNX IR#
ONNX IR 与 ONNX Script 包一起提供，可以通过 onnxscript.ir 获取。要从 ONNX 文件创建 IR 对象，将其加载为 ModelProto 并调用 ir.from_proto() 或 ir.serde.deserialize_model：
参考：ONNX IR 快速上手
# Define an example model for this example
MODEL_TEXT = r"""
<
ir_version: 8,
opset_import: ["" : 18],
producer_name: "pytorch",
producer_version: "2.0.0"
>
torch_jit (float[5,5,5] input_0) => (float[5,5] val_19, float[5,5] val_6) {
val_1 = Constant <value_int: ints = [1]> ()
val_2 = Shape <start: int = 0> (val_1)
val_3 = Size (val_2)
val_4 = Constant <value: tensor = int64 {0}> ()
val_5 = Equal (val_3, val_4)
val_6 = ReduceMean <keepdims: int = 0, noop_with_empty_axes: int = 0> (input_0, val_1)
val_7 = ReduceMean <keepdims: int = 1, noop_with_empty_axes: int = 0> (input_0, val_1)
val_8 = Shape <start: int = 0> (input_0)
val_9 = Gather <axis: int = 0> (val_8, val_1)
val_10 = ReduceProd <keepdims: int = 0, noop_with_empty_axes: int = 0> (val_9)
val_11 = Sub (input_0, val_7)
val_12 = Mul (val_11, val_11)
val_13 = ReduceMean <keepdims: int = 0, noop_with_empty_axes: int = 0> (val_12, val_1)
val_14 = Cast <to: int = 1> (val_10)
val_15 = Mul (val_13, val_14)
val_16 = Constant <value: tensor = int64 {1}> ()
val_17 = Sub (val_10, val_16)
val_18 = Cast <to: int = 1> (val_17)
val_19 = Div (val_15, val_18)
}
"""
import onnx
from onnxscript import ir
# Load the model as onnx.ModelProto
# You can also load the model from a file using onnx.load("model.onnx")
model_proto = onnx.parser.parse_model(MODEL_TEXT)
# Create an IR object from the model
model = ir.serde.deserialize_model(model_proto)
可以探索 IR 对象：
print(f"The main graph has {len(model.graph)} nodes.")
The main graph has 19 nodes.
输入信息：
print(model.graph.inputs)
[Value('input_0', type=Tensor(FLOAT), shape=[5,5,5], producer=None, index=None)]
输出信息：
print(model.graph.outputs)
[Value('val_19', type=Tensor(FLOAT), shape=[5,5], producer=, index=0), Value('val_6', type=Tensor(FLOAT), shape=[5,5], producer=, index=0)]
使用第一个输入的节点。
print(list(model.graph.inputs[0].uses()))
[(Node(name='', domain='', op_type='ReduceMean', inputs=(Value('input_0', type=Tensor(FLOAT), shape=[5,5,5], producer=None, index=None), Value('val_1', type=None, shape=None, producer=, index=0)), attributes=OrderedDict({'keepdims': Attr('keepdims', INT, 0), 'noop_with_empty_axes': Attr('noop_with_empty_axes', INT, 0)}), overload='', outputs=(Value('val_6', type=Tensor(FLOAT), shape=[5,5], producer=, index=0),), version=None, doc_string=None), 0), (Node(name='', domain='', op_type='ReduceMean', inputs=(Value('input_0', type=Tensor(FLOAT), shape=[5,5,5], producer=None, index=None), Value('val_1', type=None, shape=None, producer=, index=0)), attributes=OrderedDict({'keepdims': Attr('keepdims', INT, 1), 'noop_with_empty_axes': Attr('noop_with_empty_axes', INT, 0)}), overload='', outputs=(Value('val_7', type=None, shape=None, producer=, index=0),), version=None, doc_string=None), 0), (Node(name='', domain='', op_type='Shape', inputs=(Value('input_0', type=Tensor(FLOAT), shape=[5,5,5], producer=None, index=None),), attributes=OrderedDict({'start': Attr('start', INT, 0)}), overload='', outputs=(Value('val_8', type=None, shape=None, producer=, index=0),), version=None, doc_string=None), 0), (Node(name='', domain='', op_type='Sub', inputs=(Value('input_0', type=Tensor(FLOAT), shape=[5,5,5], producer=None, index=None), Value('val_7', type=None, shape=None, producer=, index=0)), attributes=OrderedDict(), overload='', outputs=(Value('val_11', type=None, shape=None, producer=, index=0),), version=None, doc_string=None), 0)]
产生最后一个输出（作为第i个输出）的节点。
print(model.graph.outputs[-1].producer())
print(model.graph.outputs[-1].index())
%"val_6"<FLOAT,[5,5]> ⬅️ ::ReduceMean(%"input_0", %"val_1") {keepdims=0, noop_with_empty_axes=0}
0
打印计算图：
model.graph.display(
page=False
) # Set page=True to use a pager in the terminal so long outputs are scrollable
Show code cell output
graph( name=torch_jit, inputs=( %"input_0"<FLOAT,[5,5,5]> ), outputs=( %"val_19"<FLOAT,[5,5]>, %"val_6"<FLOAT,[5,5]> ), ) { 0 | # :anonymous_node:139832865233744 %"val_1"<?,?> ⬅️ ::Constant() {value_int=[1]} 1 | # :anonymous_node:139832855544752 %"val_2"<?,?> ⬅️ ::Shape(%"val_1") {start=0} 2 | # :anonymous_node:139832855545328 %"val_3"<?,?> ⬅️ ::Size(%"val_2") 3 | # :anonymous_node:139832855545472 %"val_4"<?,?> ⬅️ ::Constant() {value=TensorProtoTensor<INT64,[]>(name='')} 4 | # :anonymous_node:139832855545616 %"val_5"<?,?> ⬅️ ::Equal(%"val_3", %"val_4") 5 | # :anonymous_node:139832855545760 %"val_6"<FLOAT,[5,5]> ⬅️ ::ReduceMean(%"input_0", %"val_1") {keepdims=0, noop_with_empty_axes=0} 6 | # :anonymous_node:139832855545904 %"val_7"<?,?> ⬅️ ::ReduceMean(%"input_0", %"val_1") {keepdims=1, noop_with_empty_axes=0} 7 | # :anonymous_node:139832855546048 %"val_8"<?,?> ⬅️ ::Shape(%"input_0") {start=0} 8 | # :anonymous_node:139832855546192 %"val_9"<?,?> ⬅️ ::Gather(%"val_8", %"val_1") {axis=0} 9 | # :anonymous_node:139832855546336 %"val_10"<?,?> ⬅️ ::ReduceProd(%"val_9") {keepdims=0, noop_with_empty_axes=0} 10 | # :anonymous_node:139832855546480 %"val_11"<?,?> ⬅️ ::Sub(%"input_0", %"val_7") 11 | # :anonymous_node:139832855546624 %"val_12"<?,?> ⬅️ ::Mul(%"val_11", %"val_11") 12 | # :anonymous_node:139832855546768 %"val_13"<?,?> ⬅️ ::ReduceMean(%"val_12", %"val_1") {keepdims=0, noop_with_empty_axes=0} 13 | # :anonymous_node:139832855547056 %"val_14"<?,?> ⬅️ ::Cast(%"val_10") {to=1} 14 | # :anonymous_node:139832855547200 %"val_15"<?,?> ⬅️ ::Mul(%"val_13", %"val_14") 15 | # :anonymous_node:139832855547632 %"val_16"<?,?> ⬅️ ::Constant() {value=TensorProtoTensor<INT64,[]>(name='')} 16 | # :anonymous_node:139832855547776 %"val_17"<?,?> ⬅️ ::Sub(%"val_10", %"val_16") 17 | # :anonymous_node:139832855547920 %"val_18"<?,?> ⬅️ ::Cast(%"val_17") {to=1} 18 | # :anonymous_node:139832855548064 %"val_19"<FLOAT,[5,5]> ⬅️ ::Div(%"val_15", %"val_18") return %"val_19"<FLOAT,[5,5]>, %"val_6"<FLOAT,[5,5]> }
将 IR 对象转换回 ModelProto：
model_proto_back = ir.serde.serialize_model(model)