onnxruntime-directml 1.20.0__cp313-cp313-win_amd64.whl

onnxruntime/transformers/models/whisper/whisper_decoder.py

@@ -0,0 +1,402 @@
+# -------------------------------------------------------------------------
+# 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 tempfile
+from pathlib import Path
+from typing import List, Optional, Union
+
+import numpy
+import onnx
+import torch
+from io_binding_helper import TypeHelper
+from models.t5.past_helper import PastKeyValuesHelper
+from onnx_model import OnnxModel
+from torch_onnx_export_helper import torch_onnx_export
+from transformers import WhisperConfig, file_utils
+from whisper_openai_helper import WhisperDecoderInitOpenai
+
+from onnxruntime import InferenceSession
+
+logger = logging.getLogger(__name__)
+
+
+class WhisperDecoderInit(torch.nn.Module):
+    """A Whisper decoder to create initial past key values.
+    This model is only called once during starting decoding.
+    """
+
+    def __init__(
+        self,
+        decoder: torch.nn.Module,
+        config: WhisperConfig,
+        decoder_start_token_id: Optional[int] = None,
+    ):
+        super().__init__()
+        self.decoder = decoder
+        self.config = config
+        self.decoder_start_token_id = (
+            decoder_start_token_id if decoder_start_token_id is not None else self.config.decoder_start_token_id
+        )
+
+    def forward(
+        self,
+        decoder_input_ids: torch.Tensor,
+        encoder_hidden_states: torch.FloatTensor,
+    ):
+        encoder_outputs = file_utils.ModelOutput()
+        encoder_outputs["last_hidden_state"] = encoder_hidden_states
+        encoder_outputs["hidden_states"] = None
+        encoder_outputs["attentions"] = None
+
+        out = self.decoder.model(
+            None,
+            encoder_outputs=encoder_outputs,
+            decoder_input_ids=decoder_input_ids,
+            past_key_values=None,
+            use_cache=True,
+            return_dict=True,
+        )
+        logits = self.decoder.proj_out(out[0])
+        return logits, out.past_key_values, out.encoder_last_hidden_state
+
+
+class WhisperDecoder(torch.nn.Module):
+    """A Whisper decoder with past key values"""
+
+    def __init__(self, decoder, config, model_impl: str = "hf", model: torch.nn.Module = None):
+        super().__init__()
+        self.decoder = decoder
+        self.config = config
+        self.model_impl = model_impl
+        if model is not None:
+            self.whisper_decoder_openai_init = WhisperDecoderInitOpenai(model, decoder)
+
+    def forward(self, decoder_input_ids, *past):
+        encoder_outputs = file_utils.ModelOutput()
+        dummy_encoder_hidden_states = torch.randn((decoder_input_ids.shape[0], 3000, int(self.config.d_model)))
+        encoder_outputs["last_hidden_state"] = dummy_encoder_hidden_states
+        encoder_outputs["hidden_states"] = dummy_encoder_hidden_states
+        encoder_outputs["attentions"] = None
+
+        if self.model_impl == "openai":
+            dummy_encoder_hidden_states.unsqueeze(0)
+            dec_out, present = self.whisper_decoder_openai_init(
+                decoder_input_ids, dummy_encoder_hidden_states, past=past
+            )
+            return dec_out, present
+
+        if len(past) == 0:
+            past_key_values = None
+        else:
+            past_key_values = PastKeyValuesHelper.back_group_by_layer(past)
+
+        decoder_out = self.decoder(
+            None,
+            encoder_outputs=encoder_outputs,
+            decoder_input_ids=decoder_input_ids,
+            past_key_values=past_key_values,
+            use_cache=True,
+            return_dict=True,
+        )
+        logits = decoder_out[0]
+        present_self, _ = PastKeyValuesHelper.group_by_self_and_cross(decoder_out.past_key_values)
+        return logits, present_self
+
+
+class WhisperDecoderInputs:
+    def __init__(
+        self,
+        decoder_input_ids,
+        past_key_values=None,
+    ):
+        self.decoder_input_ids: torch.LongTensor = decoder_input_ids
+        self.past_key_values: Union[List[torch.FloatTensor], List[torch.HalfTensor], None] = past_key_values
+
+    @staticmethod
+    def create_dummy(
+        config: WhisperConfig,
+        batch_size: int,
+        encode_sequence_length: int,
+        past_decode_sequence_length: int,
+        device: torch.device,
+        float16: bool = False,
+        use_int32_inputs: bool = False,
+        model_impl: str = "hf",
+    ):  # -> WhisperDecoderInputs:
+        """Create dummy inputs for WhisperDecoder.
+
+        Args:
+            decoder: decoder
+            batch_size (int): batch size
+            encode_sequence_length (int): sequence length of input_ids for encoder
+            past_decode_sequence_length (int): past sequence length of input_ids for decoder
+            device (torch.device): device of output tensors
+            float16 (bool): whether the model uses float32 or float16 in input
+            use_int32_inputs(bool): whether use int32 instead of int64 for some inputs
+
+        Returns:
+            WhisperDecoderInputs: dummy inputs for decoder
+        """
+        num_attention_heads: int = config.encoder_attention_heads
+        num_layers: int = config.decoder_layers  # + config.encoder_layers
+        vocab_size: int = config.vocab_size
+
+        # Use head_size, use hidden_size / num_attention_heads here.
+        # For example, whisper-large, d_model=1280 and num_heads=20
+        head_size: int = config.d_model // config.encoder_attention_heads
+
+        sequence_length: int = 1  # fixed for decoding
+        decoder_input_ids = torch.randint(
+            low=0,
+            high=vocab_size - 1,
+            size=(batch_size, sequence_length),
+            dtype=(torch.int32 if use_int32_inputs else torch.int64),
+            device=device,
+        )
+
+        float_type = torch.float16 if float16 else torch.float32
+
+        if past_decode_sequence_length > 0:
+            self_attention_past_shape = [
+                batch_size,
+                num_attention_heads,
+                past_decode_sequence_length,
+                head_size,
+            ]
+            cross_attention_past_shape = [
+                batch_size,
+                num_attention_heads,
+                encode_sequence_length if model_impl == "hf" else past_decode_sequence_length,
+                head_size,
+            ]
+
+            past = []
+            for _ in range(2 * num_layers):
+                past.append(torch.rand(self_attention_past_shape, dtype=float_type, device=device))
+
+            for _ in range(2 * num_layers):
+                past.append(torch.rand(cross_attention_past_shape, dtype=float_type, device=device))
+        else:
+            past = None
+
+        return WhisperDecoderInputs(decoder_input_ids, past)
+
+    def to_list(self) -> List:
+        input_list = [self.decoder_input_ids]
+        if self.past_key_values:
+            input_list.extend(self.past_key_values)
+        return input_list
+
+    def to_fp32(self):
+        past = [p.to(dtype=torch.float32) for p in self.past_key_values] if self.past_key_values else None
+        return WhisperDecoderInputs(
+            self.decoder_input_ids.clone(),
+            past,
+        )
+
+
+class WhisperDecoderHelper:
+    @staticmethod
+    def export_onnx(
+        decoder: WhisperDecoder,
+        device: torch.device,
+        onnx_model_path: str,
+        verbose: bool = True,
+        use_external_data_format: bool = False,
+        use_int32_inputs: bool = False,
+    ):
+        """Export decoder to ONNX
+
+        Args:
+            decoder (Union[WhisperDecoder, WhisperDecoderNoPastState]): decoder object
+            device (torch.device): device of decoder object
+            onnx_model_path (str): onnx path
+            verbose (bool, optional): print verbose information. Defaults to True.
+            use_external_data_format (bool, optional): use external data format or not. Defaults to False.
+            use_int32_inputs (bool, optional): use int32 inputs
+        """
+        assert isinstance(decoder, (WhisperDecoder, WhisperDecoderInit))
+
+        inputs = WhisperDecoderInputs.create_dummy(
+            decoder.config,
+            batch_size=2,
+            encode_sequence_length=3000,
+            past_decode_sequence_length=6 if isinstance(decoder, WhisperDecoder) else 0,
+            device=device,
+            use_int32_inputs=use_int32_inputs,
+            model_impl=decoder.model_impl,
+        )
+        input_list = inputs.to_list()
+
+        # Fix past disappearing bug - duplicate first past entry
+        # input_list.insert(2, input_list[2])
+
+        past_names = PastKeyValuesHelper.get_past_names(decoder.config.decoder_layers, present=False)
+        present_names = PastKeyValuesHelper.get_past_names(decoder.config.decoder_layers, present=True)
+        present_self_names = present_names[: 2 * decoder.config.decoder_layers]
+
+        input_past_names = past_names if isinstance(decoder, WhisperDecoder) else []
+        output_present_names = present_self_names if isinstance(decoder, WhisperDecoder) else present_names
+        output_names = ["logits", *output_present_names]
+
+        # Shape of input tensors (sequence_length==1):
+        #    input_ids: (batch_size, sequence_length)
+        #    past_self_*: (batch_size, num_heads, past_decode_sequence_length, head_size)
+        #    past_cross_*: (batch_size, num_heads, encode_sequence_length, head_size)
+
+        # Shape of output tensors:
+        #    logits: (batch_size, sequence_length, vocab_size)
+        #    past_self_*: (batch_size, num_heads, past_decode_sequence_length + sequence_length, head_size)
+        #    past_cross_*: (batch_size, num_heads, encode_sequence_length, head_size)
+
+        input_names = ["input_ids"]
+        input_names.extend(input_past_names)
+
+        dynamic_axes = {
+            "input_ids": {0: "batch_size"},
+            "encoder_hidden_states": {0: "batch_size", 1: "encode_sequence_length / 2"},
+            "logits": {0: "batch_size", 1: "sequence_length"},
+        }
+
+        for name in input_past_names:
+            dynamic_axes[name] = {
+                0: "batch_size",
+                2: "past_decode_sequence_length" if "self" in name else "encode_sequence_length",
+            }
+
+        for name in output_present_names:
+            if "cross" in name:
+                dynamic_axes[name] = {0: "batch_size", 2: "encode_sequence_length"}
+            else:  # self attention past state
+                if isinstance(decoder, WhisperDecoder):
+                    dynamic_axes[name] = {
+                        0: "batch_size",
+                        2: "past_decode_sequence_length + 1",
+                    }
+                else:
+                    dynamic_axes[name] = {
+                        0: "batch_size",
+                        # 2: 'sequence_length'
+                    }
+
+        Path(onnx_model_path).parent.mkdir(parents=True, exist_ok=True)
+
+        with tempfile.TemporaryDirectory() as tmp_dir_name:
+            temp_onnx_model_path = os.path.join(tmp_dir_name, "decoder.onnx")
+            Path(temp_onnx_model_path).parent.mkdir(parents=True, exist_ok=True)
+            torch_onnx_export(
+                decoder,
+                args=tuple(input_list),
+                f=temp_onnx_model_path if use_external_data_format else onnx_model_path,
+                export_params=True,
+                input_names=input_names,
+                output_names=output_names,
+                dynamic_axes=dynamic_axes,
+                opset_version=17,
+                do_constant_folding=True,
+                use_external_data_format=use_external_data_format,
+                verbose=verbose,
+            )
+
+            if use_external_data_format:
+                model = onnx.load_model(temp_onnx_model_path, load_external_data=True)
+                OnnxModel.save(
+                    model,
+                    onnx_model_path,
+                    save_as_external_data=True,
+                    all_tensors_to_one_file=True,
+                )
+
+    @staticmethod
+    def onnxruntime_inference(ort_session, inputs: WhisperDecoderInputs):
+        """Run inference of ONNX model."""
+        logger.debug("start onnxruntime_inference")
+
+        ort_inputs = {
+            "input_ids": numpy.ascontiguousarray(inputs.decoder_input_ids.cpu().numpy()),
+        }
+
+        if inputs.past_key_values:
+            assert len(inputs.past_key_values) % 4 == 0
+            num_layers = int(len(inputs.past_key_values) / 4)
+            past_names = PastKeyValuesHelper.get_past_names(num_layers)
+            for i, past_tensor in enumerate(inputs.past_key_values):
+                ort_inputs[past_names[i]] = numpy.ascontiguousarray(past_tensor.cpu().numpy())
+
+        ort_outputs = ort_session.run(None, ort_inputs)
+        return ort_outputs
+
+    @staticmethod
+    def verify_onnx(
+        model: Union[WhisperDecoder, WhisperDecoderInit],
+        ort_session: InferenceSession,
+        device: torch.device,
+        use_int32_inputs: bool,
+        max_cases: int = 4,
+    ):
+        """Compare the result from PyTorch and OnnxRuntime to verify the ONNX model is good."""
+        float16: bool = TypeHelper.get_input_type(ort_session, "past_key_self_0") == "tensor(float16)"
+
+        test_cases = [(4, 11, 3), (1, 2, 5), (3, 1, 1), (8, 5, 2)]
+        test_cases_max_diff = []
+        for (
+            batch_size,
+            encode_sequence_length,
+            past_decode_sequence_length,
+        ) in test_cases[:max_cases]:
+            if isinstance(model, WhisperDecoderInit):
+                dec_seq_len = 0
+            else:
+                dec_seq_len = past_decode_sequence_length
+
+            inputs = WhisperDecoderInputs.create_dummy(
+                model.config,
+                batch_size,
+                encode_sequence_length,
+                dec_seq_len,
+                device=device,
+                float16=float16,
+                use_int32_inputs=use_int32_inputs,
+            )
+
+            # We use fp32 PyTroch model as baseline even when ONNX model is fp16
+            input_list = inputs.to_fp32().to_list()
+
+            # Run inference of PyTorch model
+            with torch.no_grad():
+                torch_outputs = model(*input_list)
+
+            ort_outputs = WhisperDecoderHelper.onnxruntime_inference(ort_session, inputs)
+
+            max_diff = numpy.amax(numpy.abs(torch_outputs[0].cpu().numpy() - ort_outputs[0]))
+            max_diff_all = max_diff
+            logger.debug(f"logits max_diff={max_diff}")
+
+            for i in range(2 * model.config.num_layers):
+                max_diff = numpy.amax(numpy.abs(torch_outputs[1][i].cpu().numpy() - ort_outputs[1 + i]))
+                logger.debug(f"self attention past state {i} max_diff={max_diff}")
+                max_diff_all = max(max_diff_all, max_diff)
+
+            if isinstance(model, WhisperDecoderInit):
+                for i in range(2 * model.config.num_layers):
+                    max_diff = numpy.amax(
+                        numpy.abs(torch_outputs[2][i].cpu().numpy() - ort_outputs[1 + 2 * model.config.num_layers + i])
+                    )
+                    logger.debug(f"cross attention past state {i} max_diff={max_diff}")
+                    max_diff_all = max(max_diff_all, max_diff)
+
+            test_cases_max_diff.append(max_diff_all)
+            logger.info(
+                "batch_size=%s, encode_sequence_length=%s, past_decode_sequence_length=%s, max_diff=%s",
+                batch_size,
+                encode_sequence_length,
+                past_decode_sequence_length,
+                max_diff_all,
+            )
+
+        return max_diff_all

onnxruntime/transformers/models/whisper/whisper_encoder.py

@@ -0,0 +1,164 @@
+# -------------------------------------------------------------------------
+# 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 tempfile
+from pathlib import Path
+from typing import List
+
+import numpy
+import onnx
+import torch
+from onnx_model import OnnxModel
+from torch_onnx_export_helper import torch_onnx_export
+from transformers import WhisperConfig
+
+from onnxruntime import InferenceSession
+
+logger = logging.getLogger(__name__)
+
+
+class WhisperEncoder(torch.nn.Module):
+    """Whisper encoder outputs only the last hidden state"""
+
+    def __init__(self, encoder, config: WhisperConfig, model_impl: str = "hf"):
+        super().__init__()
+        self.encoder = encoder
+        self.config = config
+        self.model_impl = model_impl
+
+    def forward(self, input_features):
+        if self.model_impl == "openai":
+            return self.encoder(input_features)
+        return self.encoder.model.encoder(input_features)[0]
+
+
+class WhisperEncoderInputs:
+    def __init__(self, input_features):
+        self.input_ids: torch.LongTensor = input_features
+
+    @staticmethod
+    def create_dummy(
+        batch_size: int,
+        sequence_length: int,
+        feature_size: int,
+        device: torch.device,
+        use_int32_inputs: bool = False,
+    ):
+        """Create dummy inputs for Whisper encoder.
+
+        Args:
+            batch_size (int): batch size
+            sequence_length (int): sequence length
+            feature_size (int): feature size for spectrogram input
+            device (torch.device): device of output tensors
+
+        Returns:
+            WhisperEncoderInputs: dummy inputs for encoder
+        """
+
+        input_features = torch.randn(
+            size=(batch_size, feature_size, sequence_length),
+            device=device,
+        )
+        return WhisperEncoderInputs(input_features)
+
+    def to_list(self) -> List:
+        if self.input_ids is None:
+            return []
+        return [self.input_ids]
+
+
+class WhisperEncoderHelper:
+    @staticmethod
+    def export_onnx(
+        encoder,
+        device: torch.device,
+        onnx_model_path: str,
+        verbose: bool = True,
+        use_external_data_format: bool = False,
+        use_int32_inputs: bool = False,
+    ):
+        """Export encoder to ONNX
+
+        Args:
+            encoder (WhisperEncoder): encoder object
+            device (torch.device): device of encoder object
+            onnx_model_path (str): onnx path
+            verbose (bool, optional): print verbose information. Defaults to True.
+            use_external_data_format (bool, optional): use external data format or not. Defaults to False.
+        """
+        config = encoder.config
+        encoder_inputs = WhisperEncoderInputs.create_dummy(
+            batch_size=2,
+            sequence_length=3000,
+            feature_size=config.num_mel_bins,
+            device=device,
+            use_int32_inputs=use_int32_inputs,
+        )
+
+        Path(onnx_model_path).parent.mkdir(parents=True, exist_ok=True)
+
+        with tempfile.TemporaryDirectory() as tmp_dir_name:
+            temp_onnx_model_path = os.path.join(tmp_dir_name, "encoder.onnx")
+            Path(temp_onnx_model_path).parent.mkdir(parents=True, exist_ok=True)
+            torch_onnx_export(
+                encoder,
+                args=tuple(encoder_inputs.to_list()),
+                f=temp_onnx_model_path if use_external_data_format else onnx_model_path,
+                export_params=True,
+                input_names=["input_features"],
+                output_names=["hidden_states"],
+                dynamic_axes={
+                    "input_ids": {0: "batch_size", 1: "feature_size", 2: "sequence_length"},
+                    "hidden_states": {0: "batch_size", 1: "sequence_length"},
+                },
+                opset_version=17,
+                do_constant_folding=True,
+                use_external_data_format=use_external_data_format,
+                verbose=verbose,
+            )
+
+            if use_external_data_format:
+                model = onnx.load_model(temp_onnx_model_path, load_external_data=True)
+                OnnxModel.save(
+                    model,
+                    onnx_model_path,
+                    save_as_external_data=True,
+                    all_tensors_to_one_file=True,
+                )
+
+    @staticmethod
+    def onnxruntime_inference(ort_session, inputs: WhisperEncoderInputs):
+        """Run inference of ONNX model."""
+        ort_inputs = {
+            "input_ids": numpy.ascontiguousarray(inputs.input_ids.cpu().numpy()),
+        }
+
+        return ort_session.run(None, ort_inputs)
+
+    @staticmethod
+    def verify_onnx(
+        model: WhisperEncoder, ort_session: InferenceSession, device: torch.device, use_int32_inputs: bool = False
+    ):
+        """Compare the result from PyTorch and OnnxRuntime to verify the ONNX model is good."""
+        inputs = WhisperEncoderInputs.create_dummy(
+            batch_size=4,
+            sequence_length=11,
+            device=device,
+            use_int32_inputs=use_int32_inputs,
+        )
+        input_list = inputs.to_list()
+        torch_outputs = model(*input_list)
+
+        ort_outputs = WhisperEncoderHelper.onnxruntime_inference(ort_session, inputs)
+
+        max_diff = numpy.amax(numpy.abs(torch_outputs.cpu().numpy() - ort_outputs[0]))
+
+        logger.info(f"max_diff={max_diff}")
+
+        return max_diff