mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-23 19:32:23 +00:00
a36caba073
### Description Bump ruff version in CI and fixed new lint errors. - This change enables the flake8-implicit-str-concat rules which helps detect unintended string concatenations: https://beta.ruff.rs/docs/rules/#flake8-implicit-str-concat-isc - Update gitignore to include common python files that we want to exclude. ### Motivation and Context Code quality
131 lines
4.4 KiB
Python
131 lines
4.4 KiB
Python
import argparse
|
|
import os # noqa: F401
|
|
import sys
|
|
import traceback
|
|
from typing import Any, Callable, Dict, List, Optional, Sequence, Text, Tuple, TypeVar, Union, cast # noqa: F401
|
|
|
|
import numpy as np
|
|
import onnx
|
|
from onnx import AttributeProto # noqa: F401
|
|
from onnx import GraphProto # noqa: F401
|
|
from onnx import SparseTensorProto # noqa: F401
|
|
from onnx import mapping # noqa: F401
|
|
from onnx import numpy_helper # noqa: F401
|
|
from onnx import utils # noqa: F401
|
|
from onnx import TensorProto, ValueInfoProto, helper
|
|
from onnx.helper import make_opsetid
|
|
|
|
|
|
def parse_arguments():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--node_name", required=True, type=str, help="Constant Node name")
|
|
parser.add_argument("--output_file", required=True, type=str, help="Model file name to save")
|
|
|
|
args = parser.parse_args()
|
|
return args
|
|
|
|
|
|
# This is availabl in ONNX at a later commit
|
|
def make_sparse_tensor_value_info(
|
|
name, # type: Text
|
|
elem_type, # type: int
|
|
shape, # type: Optional[Sequence[Union[Text, int]]]
|
|
doc_string="", # type: Text
|
|
shape_denotation=None, # type: Optional[List[Text]]
|
|
): # type: (...) -> ValueInfoProto
|
|
"""Makes a ValueInfoProto based on the data type and shape."""
|
|
value_info_proto = ValueInfoProto()
|
|
value_info_proto.name = name
|
|
if doc_string:
|
|
value_info_proto.doc_string = doc_string
|
|
|
|
sparse_tensor_type_proto = value_info_proto.type.sparse_tensor_type
|
|
sparse_tensor_type_proto.elem_type = elem_type
|
|
|
|
sparse_tensor_shape_proto = sparse_tensor_type_proto.shape
|
|
|
|
if shape is not None:
|
|
# You might think this is a no-op (extending a normal Python
|
|
# list by [] certainly is), but protobuf lists work a little
|
|
# differently; if a field is never set, it is omitted from the
|
|
# resulting protobuf; a list that is explicitly set to be
|
|
# empty will get an (empty) entry in the protobuf. This
|
|
# difference is visible to our consumers, so make sure we emit
|
|
# an empty shape!
|
|
sparse_tensor_shape_proto.dim.extend([])
|
|
|
|
if shape_denotation:
|
|
if len(shape_denotation) != len(shape):
|
|
raise ValueError("Invalid shape_denotation. Must be of the same length as shape.")
|
|
|
|
for i, d in enumerate(shape):
|
|
dim = sparse_tensor_shape_proto.dim.add()
|
|
if d is None:
|
|
pass
|
|
elif isinstance(d, int):
|
|
dim.dim_value = d
|
|
elif isinstance(d, str):
|
|
dim.dim_param = d
|
|
else:
|
|
raise ValueError(f"Invalid item in shape: {d}. Needs to be one of `int` or `str`.")
|
|
|
|
if shape_denotation:
|
|
dim.denotation = shape_denotation[i]
|
|
|
|
return value_info_proto
|
|
|
|
|
|
def create_model(constant_node_name, output_file_name):
|
|
dense_shape = [3, 3]
|
|
sparse_values = [1.764052391052246, 0.40015721321105957, 0.978738009929657]
|
|
values_tensor = helper.make_tensor(
|
|
name="Constant",
|
|
data_type=TensorProto.FLOAT,
|
|
dims=[len(sparse_values)],
|
|
vals=np.array(sparse_values).astype(np.float32),
|
|
raw=False,
|
|
)
|
|
|
|
linear_indicies = [2, 3, 5]
|
|
indicies_tensor = helper.make_tensor(
|
|
name="indicies",
|
|
data_type=TensorProto.INT64,
|
|
dims=[len(linear_indicies)],
|
|
vals=np.array(linear_indicies).astype(np.int64),
|
|
raw=False,
|
|
)
|
|
sparse_tensor = helper.make_sparse_tensor(values_tensor, indicies_tensor, dense_shape)
|
|
|
|
# Nodes
|
|
# sparse_attribute = helper.make_attribute('value', sparse_tensor)
|
|
constant_node = helper.make_node(
|
|
constant_node_name,
|
|
inputs=[],
|
|
outputs=["values"],
|
|
name="Constant",
|
|
domain="",
|
|
value=sparse_tensor,
|
|
)
|
|
|
|
# Outputs, a square matrix
|
|
Values_info = make_sparse_tensor_value_info("values", TensorProto.FLOAT, dense_shape) # noqa: N806
|
|
|
|
graph_def = helper.make_graph(
|
|
nodes=[constant_node],
|
|
name="ConstantNodeOutput",
|
|
inputs=[],
|
|
outputs=[Values_info],
|
|
)
|
|
|
|
model_def = helper.make_model(graph_def, producer_name="dmitrism", opset_imports=[make_opsetid("", 12)])
|
|
|
|
onnx.save(model_def, output_file_name)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
args = parse_arguments()
|
|
sys.exit(create_model(args.node_name, args.output_file))
|
|
except Exception as inst:
|
|
print("Exception thrown: ", str(inst))
|
|
print(traceback.format_exc())
|