onnxruntime-directml 1.24.1__cp314-cp314-win_amd64.whl

onnxruntime/transformers/shape_optimizer.py

@@ -0,0 +1,400 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+# This tool is not used directly in bert optimization. It could assist developing the optimization script on the following scenarios:
+# (1) It could simplify graph by removing many sub-graphs related to reshape.
+# (2) It could reduce extra inputs and outputs to fit other tools. The script compare_bert_results.py or bert_perf_test.py requires 3 inputs.
+
+import argparse
+import logging
+import os
+import re  # noqa: F401
+import sys
+import tempfile
+from collections import deque  # noqa: F401
+from datetime import datetime
+from pathlib import Path  # noqa: F401
+
+import numpy as np
+import onnx
+from onnx import ModelProto, TensorProto, numpy_helper
+from onnx_model import OnnxModel
+
+import onnxruntime
+
+logger = logging.getLogger(__name__)
+
+CONSTANT_SHAPE_NAME_PREFIX = "constant_shape_opt__"
+RESHAPE_INPUT_SHAPE_PREFIX = "reshape_input_shape__"
+
+
+class BertOnnxModelShapeOptimizer(OnnxModel):
+    """
+    This optimizer will replace Shape output or the shape input of Reshape node by initializer. Currently, it requires
+    model inputs to have static shape.
+    """
+
+    def __init__(self, onnx_model):
+        super().__init__(onnx_model.model)
+
+    def add_shape_initializer(self, shape):
+        """
+        Add an initializer for constant shape.
+        """
+        shape_value = np.asarray(shape, dtype=np.int64)
+        constant_shape_name = self.create_node_name("Constant", CONSTANT_SHAPE_NAME_PREFIX)
+        tensor = onnx.helper.make_tensor(
+            name=constant_shape_name,
+            data_type=TensorProto.INT64,
+            dims=shape_value.shape,
+            vals=shape_value,
+        )
+        self.add_initializer(tensor)
+        return tensor
+
+    def get_shape_outputs(self):
+        """
+        Returns a list of output names of all Shape nodes.
+        """
+        input_name_to_nodes = self.input_name_to_nodes()
+
+        outputs = []
+        for node in self.model.graph.node:
+            if node.op_type == "Shape":
+                if node.output[0] in input_name_to_nodes:
+                    outputs.append(node.output[0])
+
+        return outputs
+
+    def get_reshape_shape_inputs(self):
+        """
+        Returns a list of shape input names of Reshape nodes.
+        """
+        self.output_name_to_node()
+
+        shape_inputs = []
+        for node in self.model.graph.node:
+            if node.op_type == "Reshape":
+                shape_inputs.append(node.input[1])
+
+        return shape_inputs
+
+    def add_shape_for_reshape_input(self):
+        """
+        For each Reshape node, create a Shape node for its first input.
+        Returns the output names of these Shape nodes.
+        """
+        output_names = []
+        nodes_to_add = []
+        for node in self.model.graph.node:
+            if node.op_type == "Reshape":
+                input = node.input[0]
+                output_name = self.create_node_name("Reshape_Input", RESHAPE_INPUT_SHAPE_PREFIX)
+                shape_node = onnx.helper.make_node("Shape", inputs=[input], outputs=[output_name])
+                nodes_to_add.append(shape_node)
+                output_names.append(output_name)
+
+        self.add_nodes(nodes_to_add)
+        return output_names
+
+    def add_extra_graph_output(self, extra_outputs):
+        """
+        Add a list of output names to graph output.
+        """
+        names_to_evaluate = []
+        output_names = [output.name for output in self.model.graph.output]
+        for name in extra_outputs:
+            if self.get_initializer(name) is not None:  # already a constant
+                continue
+            names_to_evaluate.append(name)
+
+            if name not in output_names:
+                output_info = onnx.helper.ValueInfoProto()
+                output_info.name = name
+                self.model.graph.output.extend([output_info])
+                output_names.append(name)
+
+        return names_to_evaluate
+
+    # Update input and output shape to be static
+    def use_static_input(self, inputs, batch_size=1, max_seq_len=128):
+        """
+        Update the model to use static axes instead of dynamic axes for graph inputs.
+        """
+        for input in self.model.graph.input:
+            if input.name in inputs:
+                dim_proto = input.type.tensor_type.shape.dim[0]
+                dim_proto.dim_value = batch_size
+                dim_proto = input.type.tensor_type.shape.dim[1]
+                if dim_proto.HasField("dim_param"):
+                    dim_proto.dim_value = max_seq_len
+                elif dim_proto.HasField("dim_value") and dim_proto.dim_value != max_seq_len:
+                    raise ValueError(
+                        f"Unable to set dimension value to {max_seq_len} for axis {1} of {input.name}. Contradicts existing dimension value {dim_proto.dim_value}."
+                    )
+
+    def create_dummy_inputs(
+        self,
+        input_ids,
+        segment_ids,
+        input_mask,
+        batch_size,
+        sequence_length,
+        elem_type,
+        dictionary_size=8,
+    ):
+        """
+        Create dummy data for model inputs. If the model has more than 3 inputs, please update this function accordingly before running the tool.
+        """
+        assert elem_type in [1, 6, 7]  # only int32, int64 and float32 are supported.
+
+        # Create dummy inputs
+        input_1 = np.random.randint(dictionary_size, size=(batch_size, sequence_length), dtype=np.int32)
+        input_2 = np.ones((batch_size, sequence_length), dtype=np.int32)
+        input_3 = np.zeros((batch_size, sequence_length), dtype=np.int32)
+
+        # Here we assume that 3 inputs have same data type
+        if elem_type == 1:  # float32
+            input_1 = np.float32(input_1)
+            input_2 = np.float32(input_2)
+            input_3 = np.float32(input_3)
+        elif elem_type == 7:  # int64
+            input_1 = np.int64(input_1)
+            input_2 = np.int64(input_2)
+            input_3 = np.int64(input_3)
+
+        inputs = {input_ids: input_1, input_mask: input_2, segment_ids: input_3}
+        return inputs
+
+    def shape_optimization(
+        self,
+        temp_model_path,
+        input_ids,
+        segment_ids,
+        input_mask,
+        output_names,
+        batch_size,
+        sequence_length,
+        enable_shape_opt,
+        enable_reshape_opt,
+        verbose,
+    ):
+        self.bert_inputs = [input_ids, segment_ids, input_mask]
+
+        extra_outputs = []
+        if enable_shape_opt:
+            extra_outputs.extend(self.get_shape_outputs())
+
+        if enable_reshape_opt:
+            reshape_shape_inputs = self.get_reshape_shape_inputs()
+            reshape_input_shapes = self.add_shape_for_reshape_input()
+            extra_outputs.extend(reshape_shape_inputs)
+            extra_outputs.extend(reshape_input_shapes)
+
+        if len(extra_outputs) == 0:
+            return
+
+        names_to_evaluate = self.add_extra_graph_output(extra_outputs)
+
+        # This tool does not support dynamic axes right now.
+        self.use_static_input(self.bert_inputs, batch_size, sequence_length)
+
+        with open(temp_model_path, "wb") as out:
+            out.write(self.model.SerializeToString())
+        sess_options = onnxruntime.SessionOptions()
+        sess_options.graph_optimization_level = onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL
+        session = onnxruntime.InferenceSession(
+            temp_model_path,
+            sess_options,
+            providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
+        )
+
+        elem_type = 7
+        for input in self.model.graph.input:
+            if input.name == input_ids:
+                elem_type = input.type.tensor_type.elem_type
+        inputs = self.create_dummy_inputs(input_ids, segment_ids, input_mask, batch_size, sequence_length, elem_type)
+
+        outputs = session.run(names_to_evaluate, inputs)
+        shapes = {}
+        for i, name in enumerate(names_to_evaluate):
+            shapes[name] = outputs[i]
+
+        logger.debug(f"shapes={shapes}")
+
+        if enable_reshape_opt:
+            for i, shape_input in enumerate(reshape_shape_inputs):
+                input_shape = reshape_input_shapes[i]
+                self.update_target_shape(shapes, shape_input, input_shape, verbose)
+
+        for name, shape in shapes.items():
+            tensor = self.add_shape_initializer(shape)
+            self.replace_input_of_all_nodes(name, tensor.name)
+
+        # Remove extra outputs, and prune all nodes not linked to output.
+        self.prune_graph(output_names)
+
+    def update_target_shape(self, shapes, shape_input, input_shape, verbose):
+        """
+        Update the target shape to use 0 to represent that dimension value does not change.
+        For example, shape of source data is (2, 5, 8) and target shape is (2, 5, 4, 2), the target shape will be updated to (0, 0, 4, 2).
+        """
+        if shape_input in shapes:
+            target_shape = shapes[shape_input]
+        else:
+            initializer = self.get_initializer(shape_input)
+            assert initializer is not None
+            target_shape = numpy_helper.to_array(initializer)
+
+        if input_shape in shapes:
+            source_shape = shapes[input_shape]
+        else:
+            initializer = self.get_initializer(input_shape)
+            assert initializer is not None
+            source_shape = numpy_helper.to_array(initializer)
+
+        new_target_shape = []
+        for i, dim_value in enumerate(target_shape):
+            if i < len(source_shape) and source_shape[i] == dim_value:
+                new_target_shape.append(0)
+            else:
+                new_target_shape.append(dim_value)
+        shapes[shape_input] = new_target_shape
+
+        logger.debug(f"source_shape={source_shape}, target_shape={target_shape}, new_target_shape={new_target_shape}")
+
+    def validate_input(self, input: str):
+        if not self.find_graph_input(input):
+            valid_names = [input.name for input in self.model.graph.input]
+            raise Exception(f"Input {input} does not exist in the graph inputs: {valid_names}")
+
+    def validate_outputs(self, output_names: list[str]):
+        valid_names = [output.name for output in self.model.graph.output]
+        for name in output_names:
+            if name not in valid_names:
+                raise Exception(f"Output {name} does not exist in the graph outputs: {valid_names}")
+
+    def optimize(
+        self,
+        output_path: str,
+        input_ids: str,
+        segment_ids: str,
+        input_mask: str,
+        enable_shape_opt: bool,
+        enable_reshape_opt: bool,
+        output_names: list[str] | None = None,
+        batch_size=1,
+        sequence_length=128,
+        verbose=False,
+    ):
+        # Skip if shape optimization has been done before.
+        for tensor in self.model.graph.initializer:
+            if tensor.name.startswith(CONSTANT_SHAPE_NAME_PREFIX):
+                logger.info("Skip shape optimization since it has been done before")
+                return
+
+        self.validate_input(input_ids)
+        self.validate_input(segment_ids)
+        self.validate_input(input_mask)
+
+        if output_names is not None:
+            self.validate_outputs(output_names)
+            self.prune_graph(output_names)
+
+        remaining_outputs = [output.name for output in self.model.graph.output]
+
+        if enable_shape_opt or enable_reshape_opt:
+            if len(self.get_graph_inputs_excluding_initializers()) != 3:
+                logger.info("Skip shape optimization since graph input number is not 3")
+                return
+
+            with tempfile.TemporaryDirectory() as temp_dir:
+                temp_file_name = "temp_{}.onnx".format(datetime.now().strftime("%m_%d-%H_%M_%S"))
+                dir = "." if verbose else temp_dir
+                temp_file = os.path.join(dir, temp_file_name)
+                self.shape_optimization(
+                    temp_file,
+                    input_ids,
+                    segment_ids,
+                    input_mask,
+                    remaining_outputs,
+                    batch_size,
+                    sequence_length,
+                    enable_shape_opt,
+                    enable_reshape_opt,
+                    verbose,
+                )
+            logger.debug(f"Temp model with additional outputs: {temp_file}")
+            logger.warning(
+                f"Shape optimization is done. The optimized model might only work for input with batch_size={batch_size} sequence_length={sequence_length}"
+            )
+
+        if output_path is not None:
+            with open(output_path, "wb") as out:
+                out.write(self.model.SerializeToString())
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--input", required=True, type=str)
+    parser.add_argument("--output", required=True, type=str)
+    parser.add_argument("--input_ids", required=True, type=str)
+    parser.add_argument("--segment_ids", required=True, type=str)
+    parser.add_argument("--input_mask", required=True, type=str)
+    parser.add_argument("--output_names", required=False, type=str, default=None)
+    parser.add_argument("--batch_size", required=False, type=int, default=1)
+    parser.add_argument("--sequence_length", required=False, type=int, default=128)
+    parser.add_argument("--enable_shape_opt", required=False, action="store_true")
+    parser.set_defaults(enable_shape_opt=False)
+    parser.add_argument("--enable_reshape_opt", required=False, action="store_true")
+    parser.set_defaults(enable_reshape_opt=False)
+    parser.add_argument("--verbose", required=False, action="store_true")
+    parser.set_defaults(verbose=False)
+    args = parser.parse_args()
+    return args
+
+
+def setup_logging(verbose):
+    log_handler = logging.StreamHandler(sys.stdout)
+    if verbose:
+        log_handler.setFormatter(logging.Formatter("[%(filename)s:%(lineno)s - %(funcName)20s()] %(message)s"))
+        logging_level = logging.DEBUG
+    else:
+        log_handler.setFormatter(logging.Formatter("%(filename)20s: %(message)s"))
+        logging_level = logging.INFO
+    log_handler.setLevel(logging_level)
+    logger.addHandler(log_handler)
+    logger.setLevel(logging_level)
+
+
+def main():
+    args = parse_arguments()
+    setup_logging(args.verbose)
+
+    output_names = None if args.output_names is None else args.output_names.split(";")
+
+    model = ModelProto()
+    with open(args.input, "rb") as input_file:
+        model.ParseFromString(input_file.read())
+    onnx_model = OnnxModel(model)
+
+    optimizer = BertOnnxModelShapeOptimizer(onnx_model)
+
+    optimizer.optimize(
+        args.output,
+        args.input_ids,
+        args.segment_ids,
+        args.input_mask,
+        args.enable_shape_opt,
+        args.enable_reshape_opt,
+        output_names,
+        args.batch_size,
+        args.sequence_length,
+        args.verbose,
+    )
+
+
+if __name__ == "__main__":
+    main()

onnxruntime/transformers/torch_onnx_export_helper.py

@@ -0,0 +1,74 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+import torch
+from torch._C._onnx import OperatorExportTypes
+
+TrainingMode = torch.onnx.TrainingMode
+from packaging.version import Version  # noqa: E402
+
+
+def torch_onnx_export(
+    model,
+    args,
+    f,
+    export_params=True,
+    verbose=False,
+    training=TrainingMode.EVAL,
+    input_names=None,
+    output_names=None,
+    operator_export_type=OperatorExportTypes.ONNX,
+    opset_version=None,
+    _retain_param_name=None,
+    do_constant_folding=True,
+    example_outputs=None,
+    strip_doc_string=None,
+    dynamic_axes=None,
+    keep_initializers_as_inputs=None,
+    custom_opsets=None,
+    enable_onnx_checker=None,
+    use_external_data_format=None,
+    export_modules_as_functions=False,
+):
+    if Version(torch.__version__) >= Version("1.11.0"):
+        torch.onnx.export(
+            model=model,
+            args=args,
+            f=f,
+            export_params=export_params,
+            verbose=verbose,
+            training=training,
+            input_names=input_names,
+            output_names=output_names,
+            operator_export_type=operator_export_type,
+            opset_version=opset_version,
+            do_constant_folding=do_constant_folding,
+            dynamic_axes=dynamic_axes,
+            keep_initializers_as_inputs=keep_initializers_as_inputs,
+            custom_opsets=custom_opsets,
+            export_modules_as_functions=export_modules_as_functions,
+        )
+    else:
+        torch.onnx.export(
+            model=model,
+            args=args,
+            f=f,
+            export_params=export_params,
+            verbose=verbose,
+            training=training,
+            input_names=input_names,
+            output_names=output_names,
+            operator_export_type=operator_export_type,
+            opset_version=opset_version,
+            _retain_param_name=_retain_param_name,
+            do_constant_folding=do_constant_folding,
+            example_outputs=example_outputs,
+            strip_doc_string=strip_doc_string,
+            dynamic_axes=dynamic_axes,
+            keep_initializers_as_inputs=keep_initializers_as_inputs,
+            custom_opsets=custom_opsets,
+            enable_onnx_checker=enable_onnx_checker,
+            use_external_data_format=use_external_data_format,
+        )

onnxruntime_directml-1.24.1.dist-info/METADATA

@@ -0,0 +1,216 @@
+Metadata-Version: 2.4
+Name: onnxruntime-directml
+Version: 1.24.1
+Summary: ONNX Runtime is a runtime accelerator for Machine Learning models
+Home-page: https://onnxruntime.ai
+Download-URL: https://github.com/microsoft/onnxruntime/tags
+Author: Microsoft Corporation
+Author-email: onnxruntime@microsoft.com
+License: MIT License
+Keywords: onnx machine learning
+Classifier: Development Status :: 5 - Production/Stable
+Classifier: Intended Audience :: Developers
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: POSIX :: Linux
+Classifier: Operating System :: Microsoft :: Windows
+Classifier: Operating System :: MacOS
+Classifier: Topic :: Scientific/Engineering
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Software Development
+Classifier: Topic :: Software Development :: Libraries
+Classifier: Topic :: Software Development :: Libraries :: Python Modules
+Classifier: Programming Language :: Python
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Programming Language :: Python :: 3.14
+Requires-Python: >=3.10
+Requires-Dist: flatbuffers
+Requires-Dist: numpy>=1.21.6
+Requires-Dist: packaging
+Requires-Dist: protobuf
+Requires-Dist: sympy
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: download-url
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+ONNX Runtime
+============
+
+ONNX Runtime is a performance-focused scoring engine for Open Neural Network Exchange (ONNX) models.
+For more information on ONNX Runtime, please see `aka.ms/onnxruntime <https://aka.ms/onnxruntime/>`_ or the `Github project <https://github.com/microsoft/onnxruntime/>`_.
+
+
+Changes
+-------
+
+1.24.1
+^^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.24.1
+
+1.23.0
+^^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.23.0
+
+1.22.0
+^^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.22.0
+
+1.21.0
+^^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.21.0
+
+1.20.0
+^^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.20.0
+
+1.19.0
+^^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.19.0
+
+1.18.0
+^^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.18.0
+
+1.17.0
+^^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.17.0
+
+1.16.0
+^^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.16.0
+
+1.15.0
+^^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.15.0
+
+1.14.0
+^^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.14.0
+
+1.13.0
+^^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.13.0
+
+1.12.0
+^^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.12.0
+
+1.11.0
+^^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.11.0
+
+1.10.0
+^^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.10.0
+
+1.9.0
+^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.9.0
+
+1.8.2
+^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.8.2
+
+1.8.1
+^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.8.1
+
+1.8.0
+^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.8.0
+
+1.7.0
+^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.7.0
+
+1.6.0
+^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.6.0
+
+1.5.3
+^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.5.3
+
+1.5.2
+^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.5.2
+
+1.5.1
+^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.5.1
+
+
+1.4.0
+^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.4.0
+
+1.3.1
+^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.3.1
+
+1.3.0
+^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.3.0
+
+1.2.0
+^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.2.0
+
+1.1.0
+^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.1.0
+
+1.0.0
+^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v1.0.0
+
+0.5.0
+^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v0.5.0
+
+0.4.0
+^^^^^
+
+Release Notes : https://github.com/Microsoft/onnxruntime/releases/tag/v0.4.0