Python notes

See https://onnx.ai/onnx/api/index.html#l-python-onnx-api

pip install onnx

The data structure is defined in a protocol buffer file

See https://protobuf.dev/reference/python/python-generated/ for Python APIs for protocol buffers.

TensorProto

It is defined in https://github.com/onnx/onnx/blob/main/onnx/onnx.proto#L483

https://github.com/onnx/onnx/blob/main/onnx/mapping.py defines a mapping to map datatypes from onnx to numpy.

./code/tensor-proto.py
 1#!/usr/bin/env python3
 2import onnx
 3import numpy as np
 4
 5
 6def test_dtype():
 7    # the following line is not portable as it uses the internal implementation detail
 8    #  assert onnx.mapping.TENSOR_TYPE_TO_NP_TYPE[onnx.TensorProto.FLOAT] == np.float32
 9    assert onnx.helper.tensor_dtype_to_np_dtype(onnx.TensorProto.FLOAT) == np.float32
10    assert onnx.helper.tensor_dtype_to_np_dtype(onnx.TensorProto.INT8) == np.int8
11
12    # We can convert a dtype to a string
13    assert (
14        onnx.helper.tensor_dtype_to_string(onnx.TensorProto.FLOAT)
15        == "TensorProto.FLOAT"
16    )
17    assert onnx.TensorProto.DataType.Name(onnx.TensorProto.FLOAT) == "FLOAT"
18
19
20def test_make_tensor():
21    v = np.array([[1, 2], [3, 4]])
22    t = onnx.helper.make_tensor(
23        name="my",
24        data_type=onnx.helper.np_dtype_to_tensor_dtype(v.dtype),
25        dims=v.shape,
26        vals=v,
27    )
28    print(t)
29
30    """
31    dims: 2
32    dims: 2
33    data_type: 7
34    int64_data: 1
35    int64_data: 2
36    int64_data: 3
37    int64_data: 4
38    name: "my"
39    """
40
41
42def test_build_tensor_manually():
43    n = np.arange(5, dtype=np.float32)
44    t = onnx.TensorProto()
45    t.name = "my-tensor"
46    t.data_type = onnx.TensorProto.FLOAT
47    t.dims.extend(n.shape)
48    t.float_data.extend(n)
49    print(t)
50    """
51    dims: 5
52    data_type: 1
53    float_data: 0.0
54    float_data: 1.0
55    float_data: 2.0
56    float_data: 3.0
57    float_data: 4.0
58    name: "my-tensor"
59    """
60
61
62def main():
63    test_dtype()
64    test_make_tensor()
65    test_build_tensor_manually()
66
67
68if __name__ == "__main__":
69    main()

TensorShapeProto

It is defined in https://github.com/onnx/onnx/blob/main/onnx/onnx.proto#L661.

./code/tensor-shape-proto.py
 1#!/usr/bin/env python3
 2import onnx
 3import numpy as np
 4
 5
 6def main():
 7    shape_proto = onnx.TensorShapeProto()
 8
 9    # shape_proto.dim is a list of messages, so we need to use add()
10    dim = shape_proto.dim.add()
11    dim.dim_value = 10
12
13    dim2 = shape_proto.dim.add()
14    dim2.dim_param = "N"
15    print(shape_proto)
16    """
17    dim {
18      dim_value: 10
19    }
20    dim {
21      dim_param: "N"
22    }
23    """
24
25    # Only one of dim_value and dim_param can be set
26    for d in shape_proto.dim:
27        which = d.WhichOneof("value")
28        if which == "dim_value":
29            print(d.dim_value)  # 10
30        elif which == "dim_param":
31            print(d.dim_param)  # N
32        else:
33            assert which is None
34
35
36if __name__ == "__main__":
37    main()

One thing to note is that it contains a oneof field. Also note it can contain symbolic names for shapes.

message TensorShapeProto {
  message Dimension {
    oneof value {
      int64 dim_value = 1;
      string dim_param = 2;   // namespace Shape
    };
    // Standard denotation can optionally be used to denote tensor
    // dimensions with standard semantic descriptions to ensure
    // that operations are applied to the correct axis of a tensor.
    // Refer to https://github.com/onnx/onnx/blob/main/docs/DimensionDenotation.md#denotation-definition
    // for pre-defined dimension denotations.
    optional string denotation = 3;
  };
  repeated Dimension dim = 1;
}

ValueInfoProto