onnxruntime-directml 1.20.0__cp313-cp313-win_amd64.whl

onnxruntime/transformers/quantize_helper.py

@@ -0,0 +1,76 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.  See License.txt in the project root for
+# license information.
+# --------------------------------------------------------------------------
+
+import logging
+import os
+
+import onnx
+import torch
+from transformers.modeling_utils import Conv1D
+
+logger = logging.getLogger(__name__)
+
+
+def _conv1d_to_linear(module):
+    in_size, out_size = module.weight.shape
+    linear = torch.nn.Linear(in_size, out_size)
+    linear.weight.data = module.weight.data.T.contiguous()
+    linear.bias.data = module.bias.data
+    return linear
+
+
+def conv1d_to_linear(model):
+    """in-place
+    This is for Dynamic Quantization, as Conv1D is not recognized by PyTorch, convert it to nn.Linear
+    """
+    logger.debug("replace Conv1D with Linear")
+    for name in list(model._modules):
+        module = model._modules[name]
+        if isinstance(module, Conv1D):
+            linear = _conv1d_to_linear(module)
+            model._modules[name] = linear
+        else:
+            conv1d_to_linear(module)
+
+
+def _get_size_of_pytorch_model(model):
+    torch.save(model.state_dict(), "temp.p")
+    size = os.path.getsize("temp.p") / (1024 * 1024)
+    os.remove("temp.p")
+    return size
+
+
+class QuantizeHelper:
+    @staticmethod
+    def quantize_torch_model(model, dtype=torch.qint8):
+        """
+        Usage: model = quantize_model(model)
+
+        TODO: mix of in-place and return, but results are different
+        """
+        conv1d_to_linear(model)
+        quantized_model = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=dtype)
+        logger.info(f"Size of full precision Torch model(MB):{_get_size_of_pytorch_model(model)}")
+        logger.info(f"Size of quantized Torch model(MB):{_get_size_of_pytorch_model(quantized_model)}")
+        return quantized_model
+
+    @staticmethod
+    def quantize_onnx_model(onnx_model_path, quantized_model_path, use_external_data_format=False):
+        from pathlib import Path
+
+        from onnxruntime.quantization import quantize_dynamic
+
+        Path(quantized_model_path).parent.mkdir(parents=True, exist_ok=True)
+        logger.info(f"Size of full precision ONNX model(MB):{os.path.getsize(onnx_model_path)/(1024*1024)}")
+        quantize_dynamic(
+            onnx_model_path,
+            quantized_model_path,
+            use_external_data_format=use_external_data_format,
+            extra_options={"DefaultTensorType": onnx.TensorProto.FLOAT},
+        )
+        logger.info(f"quantized model saved to:{quantized_model_path}")
+        # TODO: inlcude external data in total model size.
+        logger.info(f"Size of quantized ONNX model(MB):{os.path.getsize(quantized_model_path)/(1024*1024)}")

onnxruntime/transformers/shape_infer_helper.py

@@ -0,0 +1,122 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+import logging
+import os
+import sys
+from typing import Dict
+
+# In ORT Package the symbolic_shape_infer.py is in ../tools
+file_path = os.path.dirname(__file__)
+if os.path.exists(os.path.join(file_path, "../tools/symbolic_shape_infer.py")):
+    sys.path.append(os.path.join(file_path, "../tools"))
+else:
+    sys.path.append(os.path.join(file_path, ".."))
+
+from symbolic_shape_infer import SymbolicShapeInference, get_shape_from_type_proto, sympy  # noqa: E402
+
+logger = logging.getLogger(__name__)
+
+
+class SymbolicShapeInferenceHelper(SymbolicShapeInference):
+    def __init__(self, model, verbose=0, int_max=2**31 - 1, auto_merge=True, guess_output_rank=False):
+        super().__init__(int_max, auto_merge, guess_output_rank, verbose)
+        self.model_ = model
+        self.all_shapes_inferred_: bool = False
+        self.is_inferred_: bool = False
+        self.dynamic_axis_mapping_: Dict[str, int] = {}
+
+    def infer(self, dynamic_axis_mapping: Dict[str, int], max_runs: int = 200):
+        """Run shape inference, and try replace dynamic axis from string to integer when mapping is provided.
+
+        Args:
+            dynamic_axis_mapping (_type_): a dictionary with name of dynamic axis as key, like {"batch_size" : 4}
+            max_runs (int, optional): limit maximum number of runs to avoid infinite loop. Defaults to 200.
+
+        Returns:
+            bool: whether all shapes has been inferred or not.
+        """
+        assert dynamic_axis_mapping is not None
+
+        if self.is_inferred_ and self.dynamic_axis_mapping_ == dynamic_axis_mapping:
+            return self.all_shapes_inferred_
+
+        self.dynamic_axis_mapping_ = dynamic_axis_mapping
+
+        self._preprocess(self.model_)
+
+        count = 0
+        while self.run_:
+            logger.debug(f"shape infer run {count}")
+            self.all_shapes_inferred_ = self._infer_impl()
+            count += 1
+            if max_runs > 0 and count >= max_runs:
+                break
+
+        self.is_inferred_ = True
+        return self.all_shapes_inferred_
+
+    def _get_sympy_shape(self, node, idx):
+        """Override it to ensure shape inference by giving the actual value of dynamic axis."""
+        sympy_shape = []
+
+        shape = self._get_shape(node, idx)
+        if shape:
+            for dim in shape:
+                if isinstance(dim, str):
+                    if dim in self.dynamic_axis_mapping_:
+                        sympy_shape.append(self.dynamic_axis_mapping_[dim])
+                    elif dim in self.symbolic_dims_:
+                        sympy_shape.append(self.symbolic_dims_[dim])
+                    else:
+                        sympy_shape.append(sympy.Symbol(dim, integer=True))
+                else:
+                    assert dim is not None
+                    sympy_shape.append(dim)
+        return sympy_shape
+
+    def get_edge_shape(self, edge):
+        """Get shape of an edge.
+
+        Args:
+            edge (str): name of edge
+
+        Returns:
+            Optional[List[int]]: the shape, or None if shape is unknown
+        """
+        assert self.all_shapes_inferred_
+        if edge not in self.known_vi_:
+            print("Cannot retrieve the shape of " + str(edge))
+            return None
+
+        type_proto = self.known_vi_[edge].type
+        shape = get_shape_from_type_proto(type_proto)
+
+        if shape is not None:
+            for i, dim in enumerate(shape):
+                if isinstance(dim, str) and dim in self.dynamic_axis_mapping_:
+                    shape[i] = self.dynamic_axis_mapping_[dim]
+
+        return shape
+
+    def compare_shape(self, edge, edge_other):
+        """Compare shape of two edges.
+
+        Args:
+            edge (str): name of edge
+            edge_other (str): name of another edge
+
+        Raises:
+            Exception: At least one shape is missed for edges to compare
+
+        Returns:
+            bool: whether the shape is same or not
+        """
+        assert self.all_shapes_inferred_
+        shape = self.get_edge_shape(edge)
+        shape_other = self.get_edge_shape(edge_other)
+        if shape is None or shape_other is None:
+            raise Exception("At least one shape is missed for edges to compare")
+        return shape == shape_other

onnxruntime/transformers/shape_optimizer.py

@@ -0,0 +1,401 @@
+# -------------------------------------------------------------------------
+# 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
+from typing import List, Optional
+
+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: Optional[List[str]] = 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,
+        )