mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-29 20:14:01 +00:00
Whisper Export (#15247)
### Description Add scripts to export Whisper model to ONNX and integrate the ORT BeamSearch op with the resulting graphs. Example command to execute this script: python convert_to_onnx.py -m openai/whisper-large --output whisper -e --------- Co-authored-by: Peter McAughan <petermca@microsoft.com>
This commit is contained in:
petermcaughan
GitHub
parent
3cf3fa0467
commit
f30e2d4387
8 changed files with 1631 additions and 0 deletions
|
|
@ -5,6 +5,9 @@
|
|||
# --------------------------------------------------------------------------
|
||||
|
||||
import logging
|
||||
from typing import List, Tuple
|
||||
|
||||
import torch
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
|
@ -66,3 +69,82 @@ class PastKeyValuesHelper:
|
|||
]
|
||||
for i in range(num_layers)
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def back_group_by_layer(past_key_values: Tuple[Tuple[torch.Tensor]]):
|
||||
"""Categorize present_key_values from self and cross attention to layer by layer.
|
||||
|
||||
Reorder past state from grouped by self/cross attention to grouped by layer.
|
||||
Before: past_key_self_0, past_value_self_0, past_key_self_1, past_value_self_1, ...,
|
||||
past_key_cross_0, past_value_cross_0, past_key_cross_1, past_value_cross_1, ...
|
||||
After: (past_key_self_0, past_value_self_0, past_key_cross_0, past_value_cross_0),
|
||||
(past_key_self_1, past_value_self_1, past_key_cross_1, past_value_cross_1),
|
||||
|
||||
Args:
|
||||
present_key_values: From past_key_values of a model (group by self and cross attention)
|
||||
|
||||
Returns:
|
||||
past_tuples: present key and values grouped by layer.
|
||||
"""
|
||||
past_tuples = ()
|
||||
half_idx = len(past_key_values) // 2
|
||||
for i in range(len(past_key_values) // 4):
|
||||
idx = 2 * i
|
||||
past_tuples += (
|
||||
(
|
||||
past_key_values[idx],
|
||||
past_key_values[idx + 1],
|
||||
past_key_values[half_idx + idx],
|
||||
past_key_values[half_idx + idx + 1],
|
||||
),
|
||||
)
|
||||
return past_tuples
|
||||
|
||||
@staticmethod
|
||||
def group_by_self_and_cross(present_key_values: Tuple[torch.Tensor], concat: bool = False):
|
||||
"""Categorize present_key_values into self and cross attention.
|
||||
|
||||
Split present state from grouped by layer to grouped by self/cross attention.
|
||||
Before: (past_key_self_0, past_value_self_0, past_key_cross_0, past_value_cross_0),
|
||||
(past_key_self_1, past_value_self_1, past_key_cross_1, past_value_cross_1), ...
|
||||
After: (past_key_self_0, past_value_self_0, past_key_self_1, past_value_self_1, ...),
|
||||
(past_key_cross_0, past_value_cross_0, past_key_cross_1, past_value_cross_1, ...)
|
||||
|
||||
Args:
|
||||
present_key_values: From past_key_values of a model (group by layer)
|
||||
concat: If concat self attention with cross attention key/value to return
|
||||
|
||||
Returns:
|
||||
present_self (Tuple[torch.Tensor]): present key and values from self attention
|
||||
present_cross (Tuple[torch.Tensor]): present key and values from cross attention
|
||||
"""
|
||||
present_self: List[torch.Tensor] = []
|
||||
present_cross: List[torch.Tensor] = []
|
||||
for _, present_layer_i in enumerate(present_key_values):
|
||||
assert len(present_layer_i) == 4, f"Expected to have four items. Got {len(present_layer_i)}"
|
||||
present_key_self, present_value_self, present_key_cross, present_value_cross = present_layer_i
|
||||
present_self.extend([present_key_self, present_value_self])
|
||||
present_cross.extend([present_key_cross, present_value_cross])
|
||||
if concat:
|
||||
return present_self + present_cross
|
||||
else:
|
||||
return present_self, present_cross
|
||||
|
||||
@staticmethod
|
||||
def get_input_names(past_key_values: Tuple[Tuple[torch.Tensor]], encoder=True):
|
||||
"""Process input names of model wrapper.
|
||||
|
||||
Args:
|
||||
past_key_values: Consider `self` and `cross` past_key_values
|
||||
|
||||
Returns:
|
||||
names (List[string]): input names
|
||||
"""
|
||||
names = []
|
||||
num_layers = len(past_key_values) // 4 if encoder else len(past_key_values)
|
||||
prefix = "past_" if not encoder else "present_"
|
||||
for i in range(num_layers):
|
||||
names.extend([prefix + s for s in [f"key_self_{i}", f"value_self_{i}"]])
|
||||
for i in range(num_layers):
|
||||
names.extend([prefix + s for s in [f"key_cross_{i}", f"value_cross_{i}"]])
|
||||
return names
|
||||
|
|
|
|||
|
|
@ -0,0 +1,4 @@
|
|||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License.
|
||||
# --------------------------------------------------------------------------
|
||||
|
|
@ -0,0 +1,294 @@
|
|||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import logging
|
||||
import os
|
||||
import sys
|
||||
|
||||
import torch
|
||||
from whisper_chain import chain_model
|
||||
from whisper_helper import PRETRAINED_WHISPER_MODELS, WhisperHelper
|
||||
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
from benchmark_helper import Precision, create_onnxruntime_session, prepare_environment, setup_logger # noqa: E402
|
||||
|
||||
logger = logging.getLogger("")
|
||||
|
||||
|
||||
def parse_arguments():
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
pretrained_models = PRETRAINED_WHISPER_MODELS
|
||||
parser.add_argument(
|
||||
"-m",
|
||||
"--model_name_or_path",
|
||||
required=False,
|
||||
default=PRETRAINED_WHISPER_MODELS[0],
|
||||
type=str,
|
||||
help="Model path, or pretrained model name in the list: " + ", ".join(pretrained_models),
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--cache_dir",
|
||||
required=False,
|
||||
type=str,
|
||||
default=os.path.join(".", "cache_models"),
|
||||
help="Directory to cache pre-trained models",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
required=False,
|
||||
type=str,
|
||||
default=os.path.join(".", "onnx_models"),
|
||||
help="Output directory",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"-o",
|
||||
"--optimize_onnx",
|
||||
required=False,
|
||||
action="store_true",
|
||||
help="Use optimizer.py to optimize onnx model",
|
||||
)
|
||||
parser.set_defaults(optimize_onnx=False)
|
||||
|
||||
parser.add_argument("--use_gpu", required=False, action="store_true", help="use GPU for inference")
|
||||
parser.set_defaults(use_gpu=False)
|
||||
|
||||
parser.add_argument(
|
||||
"-p",
|
||||
"--precision",
|
||||
required=False,
|
||||
type=Precision,
|
||||
default=Precision.FLOAT32,
|
||||
choices=[Precision.FLOAT32, Precision.FLOAT16],
|
||||
help="Precision of model to run. fp32 for full precision, fp16 for half precision",
|
||||
)
|
||||
|
||||
parser.add_argument("--verbose", required=False, action="store_true")
|
||||
parser.set_defaults(verbose=False)
|
||||
|
||||
parser.add_argument("-e", "--use_external_data_format", required=False, action="store_true")
|
||||
parser.set_defaults(use_external_data_format=False)
|
||||
|
||||
parser.add_argument(
|
||||
"-s",
|
||||
"--use_decoder_start_token",
|
||||
required=False,
|
||||
action="store_true",
|
||||
help="Use config.decoder_start_token_id. Otherwise, add an extra graph input for decoder_input_ids.",
|
||||
)
|
||||
parser.set_defaults(use_decoder_start_token=False)
|
||||
|
||||
parser.add_argument(
|
||||
"-w",
|
||||
"--overwrite",
|
||||
required=False,
|
||||
action="store_true",
|
||||
help="overwrite existing ONNX model",
|
||||
)
|
||||
parser.set_defaults(overwrite=False)
|
||||
|
||||
parser.add_argument(
|
||||
"--disable_auto_mixed_precision",
|
||||
required=False,
|
||||
action="store_true",
|
||||
help="use pure fp16 instead of mixed precision",
|
||||
)
|
||||
parser.set_defaults(disable_auto_mixed_precision=False)
|
||||
|
||||
parser.add_argument(
|
||||
"--separate_encoder_and_decoder_init",
|
||||
required=False,
|
||||
action="store_true",
|
||||
help="Do not merge encode and decoder init. Output 3 instead of 2 onnx models.",
|
||||
)
|
||||
parser.set_defaults(separate_encoder_and_decoder_init=False)
|
||||
|
||||
parser.add_argument(
|
||||
"--use_int64_inputs",
|
||||
required=False,
|
||||
action="store_true",
|
||||
help="Use int64 instead of int32 for input_ids, position_ids and attention_mask.",
|
||||
)
|
||||
parser.set_defaults(use_int64_inputs=False)
|
||||
|
||||
parser.add_argument(
|
||||
"--chain_model",
|
||||
required=False,
|
||||
action="store_true",
|
||||
help="Produce beam search model with chained encdecinit and decoder.",
|
||||
)
|
||||
parser.set_defaults(chain_model=True)
|
||||
|
||||
parser.add_argument(
|
||||
"--beam_output_model",
|
||||
type=str,
|
||||
default="whisper_beamsearch.onnx",
|
||||
help="default name is whisper_beamsearch.onnx.",
|
||||
)
|
||||
|
||||
parser.add_argument("--no_repeat_ngram_size", type=int, default=3, help="default to 3")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
return args
|
||||
|
||||
|
||||
def export_onnx_models(
|
||||
model_name_or_path,
|
||||
cache_dir,
|
||||
output_dir,
|
||||
use_gpu,
|
||||
use_external_data_format,
|
||||
optimize_onnx,
|
||||
precision,
|
||||
verbose,
|
||||
use_decoder_start_token: bool = False,
|
||||
merge_encoder_and_decoder_init: bool = True,
|
||||
overwrite: bool = False,
|
||||
disable_auto_mixed_precision: bool = False,
|
||||
use_int32_inputs: bool = True,
|
||||
):
|
||||
device = torch.device("cuda:0" if use_gpu else "cpu")
|
||||
|
||||
models = WhisperHelper.load_model(model_name_or_path, cache_dir, device, merge_encoder_and_decoder_init)
|
||||
config = models["decoder"].config
|
||||
|
||||
if (not use_external_data_format) and (config.num_layers > 24):
|
||||
logger.info("Try use_external_data_format when model size > 2GB")
|
||||
|
||||
output_paths = []
|
||||
for name, model in models.items():
|
||||
model.to(device)
|
||||
filename_suffix = "_" + name
|
||||
|
||||
onnx_path = WhisperHelper.get_onnx_path(
|
||||
output_dir,
|
||||
model_name_or_path,
|
||||
suffix=filename_suffix,
|
||||
new_folder=False,
|
||||
)
|
||||
|
||||
if overwrite or not os.path.exists(onnx_path):
|
||||
logger.info(f"Exporting ONNX model to {onnx_path}")
|
||||
# We have to clone model before exporting onnx, otherwise verify_onnx will report large difference.
|
||||
cloned_model = copy.deepcopy(model).to(device)
|
||||
WhisperHelper.export_onnx(
|
||||
cloned_model,
|
||||
device,
|
||||
onnx_path,
|
||||
verbose,
|
||||
use_external_data_format,
|
||||
use_decoder_input_ids=not use_decoder_start_token,
|
||||
use_int32_inputs=use_int32_inputs,
|
||||
)
|
||||
else:
|
||||
logger.info(f"Skip exporting: existed ONNX model {onnx_path}")
|
||||
|
||||
# Optimize ONNX graph. Note that we have not implemented graph optimization for Whisper yet.
|
||||
if optimize_onnx or precision != Precision.FLOAT32:
|
||||
output_path = WhisperHelper.get_onnx_path(
|
||||
output_dir,
|
||||
model_name_or_path,
|
||||
suffix=filename_suffix + "_" + str(precision),
|
||||
new_folder=False,
|
||||
)
|
||||
|
||||
if overwrite or not os.path.exists(output_path):
|
||||
logger.info(f"Optimizing model to {output_path}")
|
||||
WhisperHelper.optimize_onnx(
|
||||
onnx_path,
|
||||
output_path,
|
||||
precision == Precision.FLOAT16,
|
||||
config.encoder_attention_heads,
|
||||
config.d_model,
|
||||
use_external_data_format,
|
||||
auto_mixed_precision=not disable_auto_mixed_precision,
|
||||
use_gpu=use_gpu,
|
||||
)
|
||||
else:
|
||||
logger.info(f"Skip optimizing: existed ONNX model {onnx_path}")
|
||||
else:
|
||||
output_path = onnx_path
|
||||
|
||||
ort_session = create_onnxruntime_session(
|
||||
output_path,
|
||||
use_gpu=use_gpu,
|
||||
provider=["CUDAExecutionProvider", "CPUExecutionProvider"] if use_gpu else ["CPUExecutionProvider"],
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
max_diff = WhisperHelper.verify_onnx(model, ort_session, device, use_int32_inputs)
|
||||
logger.info(f"PyTorch and OnnxRuntime results max difference = {max_diff}")
|
||||
if max_diff > 1e-4:
|
||||
logger.warning("PyTorch and OnnxRuntime results are NOT close")
|
||||
|
||||
output_paths.append(output_path)
|
||||
|
||||
return output_paths
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_arguments()
|
||||
|
||||
setup_logger(args.verbose)
|
||||
|
||||
logger.info(f"Arguments:{args}")
|
||||
|
||||
cache_dir = args.cache_dir
|
||||
output_dir = args.output if not args.output.endswith(".onnx") else os.path.dirname(args.output)
|
||||
prepare_environment(cache_dir, output_dir, args.use_gpu)
|
||||
|
||||
if args.precision != Precision.FLOAT32:
|
||||
assert args.optimize_onnx, "fp16/int8 requires --optimize_onnx"
|
||||
|
||||
if args.precision == Precision.FLOAT16:
|
||||
assert args.use_gpu, "fp16 requires --use_gpu"
|
||||
|
||||
if args.optimize_onnx:
|
||||
logger.warning("Graph optimization for Whisper is not implemented yet.")
|
||||
|
||||
output_paths = export_onnx_models(
|
||||
args.model_name_or_path,
|
||||
cache_dir,
|
||||
output_dir,
|
||||
args.use_gpu,
|
||||
args.use_external_data_format,
|
||||
args.optimize_onnx,
|
||||
args.precision,
|
||||
args.verbose,
|
||||
args.use_decoder_start_token,
|
||||
not args.separate_encoder_and_decoder_init,
|
||||
args.overwrite,
|
||||
args.disable_auto_mixed_precision,
|
||||
not args.use_int64_inputs,
|
||||
)
|
||||
|
||||
if args.chain_model:
|
||||
logger.info("Chaining model ... :")
|
||||
args.beam_model_output_dir = WhisperHelper.get_onnx_path(
|
||||
output_dir,
|
||||
args.model_name_or_path,
|
||||
suffix="_beamsearch",
|
||||
new_folder=False,
|
||||
)
|
||||
for path in output_paths:
|
||||
if "encoder_decoder" in path:
|
||||
args.encoder_path = path
|
||||
elif "decoder" in path:
|
||||
args.decoder_path = path
|
||||
chain_model(args)
|
||||
output_paths.append(args.beam_model_output_dir)
|
||||
|
||||
logger.info(f"Done! Outputs: {output_paths}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
|
|
@ -0,0 +1,101 @@
|
|||
import onnx
|
||||
from onnx import TensorProto, helper
|
||||
from transformers import WhisperConfig
|
||||
|
||||
|
||||
def add_attention_mask(model):
|
||||
# Add attention mask - required by BeamSearch but unused in Pytorch
|
||||
mask = helper.make_tensor_value_info(
|
||||
"encoder_attention_mask", TensorProto.INT32, shape=["batch", "feature_size", "sequence"]
|
||||
)
|
||||
model.graph.input.insert(1, mask)
|
||||
|
||||
|
||||
def chain_model(args):
|
||||
# Load encoder/decoder and insert necessary (but unused) graph inputs expected by BeamSearch op
|
||||
encoder_model = onnx.load(args.encoder_path, load_external_data=True)
|
||||
encoder_model.graph.name = "encoderdecoderinit subgraph"
|
||||
add_attention_mask(encoder_model)
|
||||
|
||||
decoder_model = onnx.load(args.decoder_path, load_external_data=True)
|
||||
decoder_model.graph.name = "decoder subgraph"
|
||||
add_attention_mask(decoder_model)
|
||||
|
||||
config = WhisperConfig.from_pretrained(args.model_name_or_path)
|
||||
pad_token_id = config.pad_token_id
|
||||
|
||||
beam_inputs = [
|
||||
"input_features",
|
||||
"max_length",
|
||||
"min_length",
|
||||
"num_beams",
|
||||
"num_return_sequences",
|
||||
"length_penalty",
|
||||
"repetition_penalty",
|
||||
"",
|
||||
"",
|
||||
"attention_mask",
|
||||
]
|
||||
beam_outputs = ["sequences"]
|
||||
|
||||
node = helper.make_node("BeamSearch", inputs=beam_inputs, outputs=beam_outputs, name="BeamSearch_zcode")
|
||||
node.domain = "com.microsoft"
|
||||
node.attribute.extend(
|
||||
[
|
||||
helper.make_attribute("eos_token_id", 50256),
|
||||
helper.make_attribute("pad_token_id", pad_token_id),
|
||||
helper.make_attribute("decoder_start_token_id", 50257),
|
||||
helper.make_attribute("no_repeat_ngram_size", args.no_repeat_ngram_size),
|
||||
helper.make_attribute("early_stopping", True),
|
||||
helper.make_attribute("model_type", 2),
|
||||
helper.make_attribute("decoder", decoder_model.graph),
|
||||
helper.make_attribute("encoder", encoder_model.graph),
|
||||
]
|
||||
)
|
||||
|
||||
# beam graph inputs
|
||||
input_features = helper.make_tensor_value_info(
|
||||
"input_features", TensorProto.FLOAT, ["batch_size", "feature_size", "sequence_length"]
|
||||
)
|
||||
max_length = helper.make_tensor_value_info("max_length", TensorProto.INT32, [1])
|
||||
min_length = helper.make_tensor_value_info("min_length", TensorProto.INT32, [1])
|
||||
num_beams = helper.make_tensor_value_info("num_beams", TensorProto.INT32, [1])
|
||||
num_return_sequences = helper.make_tensor_value_info("num_return_sequences", TensorProto.INT32, [1])
|
||||
length_penalty = helper.make_tensor_value_info("length_penalty", TensorProto.FLOAT, [1])
|
||||
repetition_penalty = helper.make_tensor_value_info("repetition_penalty", TensorProto.FLOAT, [1])
|
||||
attention_mask = helper.make_tensor_value_info(
|
||||
"attention_mask", TensorProto.INT32, ["batch_size", "feature_size", "sequence_length"]
|
||||
)
|
||||
|
||||
graph_inputs = [
|
||||
input_features,
|
||||
max_length,
|
||||
min_length,
|
||||
num_beams,
|
||||
num_return_sequences,
|
||||
length_penalty,
|
||||
repetition_penalty,
|
||||
attention_mask,
|
||||
]
|
||||
|
||||
# graph outputs
|
||||
sequences = helper.make_tensor_value_info(
|
||||
"sequences", TensorProto.INT32, ["batch_size", "num_return_sequences", "max_length"]
|
||||
)
|
||||
graph_outputs = [sequences]
|
||||
|
||||
# Initializers/opsets
|
||||
initializers = []
|
||||
opset_import = [helper.make_opsetid(domain="com.microsoft", version=1), helper.make_opsetid(domain="", version=17)]
|
||||
|
||||
beam_graph = helper.make_graph([node], "beam-search-test", graph_inputs, graph_outputs, initializers)
|
||||
beam_model = helper.make_model(beam_graph, producer_name="pytorch", opset_imports=opset_import)
|
||||
|
||||
onnx.save(
|
||||
beam_model,
|
||||
args.beam_model_output_dir,
|
||||
save_as_external_data=True,
|
||||
all_tensors_to_one_file=True,
|
||||
convert_attribute=True,
|
||||
)
|
||||
onnx.checker.check_model(args.beam_model_output_dir, full_check=True)
|
||||
|
|
@ -0,0 +1,416 @@
|
|||
# -------------------------------------------------------------------------
|
||||
# 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 sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import List, Union
|
||||
|
||||
import numpy
|
||||
import onnx
|
||||
import torch
|
||||
from transformers import WhisperConfig, file_utils
|
||||
from whisper_encoder import WhisperEncoderInputs
|
||||
|
||||
from onnxruntime import InferenceSession
|
||||
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
from io_binding_helper import TypeHelper # noqa: E402
|
||||
from models.t5.past_helper import PastKeyValuesHelper # noqa: E402
|
||||
from onnx_model import OnnxModel # noqa: E402
|
||||
from torch_onnx_export_helper import torch_onnx_export # noqa: E402
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class WhisperDecoderInit(torch.nn.Module):
|
||||
"""A Whisper decoder with LM head to create initial past key values.
|
||||
This model is only called once during starting decoding.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
decoder: torch.nn.Module,
|
||||
lm_head: torch.nn.Module,
|
||||
config: WhisperConfig,
|
||||
decoder_start_token_id: int = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.decoder = decoder
|
||||
self.lm_head = lm_head
|
||||
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_attention_mask: 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,
|
||||
head_mask=encoder_attention_mask,
|
||||
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 LM head and past key values"""
|
||||
|
||||
def __init__(self, decoder, lm_head, config):
|
||||
super().__init__()
|
||||
self.decoder = decoder
|
||||
self.lm_head = lm_head
|
||||
self.config = config
|
||||
|
||||
def forward(self, decoder_input_ids, encoder_attention_mask, *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 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,
|
||||
# decoder_attention_mask=encoder_attention_mask,
|
||||
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,
|
||||
encoder_attention_mask=None,
|
||||
past_key_values=None,
|
||||
):
|
||||
self.decoder_input_ids: torch.LongTensor = decoder_input_ids
|
||||
self.encoder_attention_mask: torch.LongTensor = encoder_attention_mask
|
||||
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,
|
||||
): # -> 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,
|
||||
)
|
||||
|
||||
encoder_inputs = WhisperEncoderInputs.create_dummy(
|
||||
batch_size,
|
||||
encode_sequence_length,
|
||||
vocab_size,
|
||||
device,
|
||||
use_int32_inputs=use_int32_inputs,
|
||||
)
|
||||
|
||||
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,
|
||||
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
|
||||
|
||||
encoder_attention_mask = torch.zeros(
|
||||
(encoder_inputs.input_ids.shape[0], 1, encoder_inputs.input_ids.shape[1], encoder_inputs.input_ids.shape[1])
|
||||
).type(torch.int8)
|
||||
return WhisperDecoderInputs(decoder_input_ids, encoder_attention_mask, past)
|
||||
|
||||
def to_list(self) -> List:
|
||||
input_list = [
|
||||
self.decoder_input_ids,
|
||||
self.encoder_attention_mask,
|
||||
]
|
||||
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(),
|
||||
self.encoder_attention_mask.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=5 if isinstance(decoder, WhisperDecoder) else 0,
|
||||
device=device,
|
||||
use_int32_inputs=use_int32_inputs,
|
||||
)
|
||||
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)
|
||||
# encoder_attention_mask: (batch_size, encode_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.append("encoder_attention_mask")
|
||||
input_names.extend(input_past_names)
|
||||
|
||||
dynamic_axes = {
|
||||
"input_ids": {0: "batch_size"},
|
||||
"encoder_attention_mask": {0: "batch_size", 1: "encode_sequence_length"},
|
||||
"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()),
|
||||
"encoder_attention_mask": numpy.ascontiguousarray(inputs.encoder_attention_mask.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(
|
||||
f"batch_size={batch_size}, encode_sequence_length={encode_sequence_length}, "
|
||||
+ f"past_decode_sequence_length={past_decode_sequence_length}, max_diff={max_diff_all}"
|
||||
)
|
||||
|
||||
return max_diff_all
|
||||
|
|
@ -0,0 +1,162 @@
|
|||
# -------------------------------------------------------------------------
|
||||
# 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 sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
|
||||
import numpy
|
||||
import onnx
|
||||
import torch
|
||||
from transformers import WhisperConfig
|
||||
|
||||
from onnxruntime import InferenceSession
|
||||
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
from onnx_model import OnnxModel # noqa: E402
|
||||
from torch_onnx_export_helper import torch_onnx_export # noqa: E402
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class WhisperEncoder(torch.nn.Module):
|
||||
"""Whisper encoder outputs only the last hidden state"""
|
||||
|
||||
def __init__(self, encoder, config: WhisperConfig):
|
||||
super().__init__()
|
||||
self.encoder = encoder
|
||||
self.config = config
|
||||
|
||||
def forward(self, input_features, attention_mask):
|
||||
return self.encoder.model.encoder(input_features)[0]
|
||||
|
||||
|
||||
class WhisperEncoderInputs:
|
||||
def __init__(self, input_features, attention_mask):
|
||||
self.input_ids: torch.LongTensor = input_features
|
||||
# HF Whisper model doesn't support Attention Mask functionality
|
||||
|
||||
@staticmethod
|
||||
def create_dummy(
|
||||
batch_size: int, sequence_length: int, feature_size: int, device: torch.device, use_int32_inputs: bool
|
||||
):
|
||||
"""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
|
||||
"""
|
||||
dtype = torch.float32
|
||||
|
||||
input_features = torch.randn(
|
||||
size=(batch_size, feature_size, sequence_length),
|
||||
device=device,
|
||||
)
|
||||
attention_mask = torch.ones([batch_size, feature_size, sequence_length], dtype=dtype, device=device)
|
||||
return WhisperEncoderInputs(input_features, attention_mask)
|
||||
|
||||
def to_list(self) -> List:
|
||||
if self.input_features is None:
|
||||
return []
|
||||
return [self.input_features]
|
||||
|
||||
|
||||
class WhisperEncoderHelper:
|
||||
@staticmethod
|
||||
def export_onnx(
|
||||
encoder,
|
||||
device: torch.device,
|
||||
onnx_model_path: str,
|
||||
verbose: bool = True,
|
||||
use_external_data_format: 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,
|
||||
)
|
||||
|
||||
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()),
|
||||
"attention_mask": numpy.ascontiguousarray(inputs.attention_mask.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
|
||||
|
|
@ -0,0 +1,307 @@
|
|||
# -------------------------------------------------------------------------
|
||||
# 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 sys
|
||||
import tempfile
|
||||
from pathlib import Path
|
||||
from typing import List, Optional
|
||||
|
||||
import numpy
|
||||
import onnx
|
||||
import torch
|
||||
from transformers import WhisperConfig
|
||||
from whisper_decoder import WhisperDecoderInit
|
||||
from whisper_encoder import WhisperEncoder, WhisperEncoderInputs
|
||||
|
||||
from onnxruntime import InferenceSession
|
||||
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
from models.t5.past_helper import PastKeyValuesHelper # noqa: E402
|
||||
from onnx_model import OnnxModel # noqa: E402
|
||||
from torch_onnx_export_helper import torch_onnx_export # noqa: E402
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class WhisperEncoderDecoderInit(torch.nn.Module):
|
||||
"""A combination of WhisperEncoder and WhisperDecoderInit."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
encoder: torch.nn.Module,
|
||||
decoder: torch.nn.Module,
|
||||
lm_head: torch.nn.Module,
|
||||
config: WhisperConfig,
|
||||
decoder_start_token_id: Optional[int] = None,
|
||||
):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.whisper_encoder = WhisperEncoder(encoder, config)
|
||||
self.whisper_decoder_init = WhisperDecoderInit(decoder, lm_head, config, decoder_start_token_id)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
encoder_input_ids: torch.Tensor,
|
||||
encoder_attention_mask: torch.Tensor,
|
||||
decoder_input_ids: torch.Tensor = None,
|
||||
):
|
||||
encoder_hidden_states: torch.FloatTensor = self.whisper_encoder(encoder_input_ids, None)
|
||||
# Decoder out: (logits, past_key_values, encoder_hidden_state)
|
||||
decinit_out = self.whisper_decoder_init(decoder_input_ids, encoder_attention_mask, encoder_hidden_states)
|
||||
present_self, present_cross = PastKeyValuesHelper.group_by_self_and_cross(decinit_out[1])
|
||||
present = present_self + present_cross
|
||||
return decinit_out[0], encoder_hidden_states, present
|
||||
|
||||
|
||||
class WhisperEncoderDecoderInitInputs:
|
||||
def __init__(self, encoder_input_ids, encoder_attention_mask, decoder_input_ids=None):
|
||||
self.encoder_input_ids: torch.LongTensor = encoder_input_ids
|
||||
self.encoder_attention_mask: torch.LongTensor = encoder_attention_mask
|
||||
self.decoder_input_ids: torch.LongTensor = decoder_input_ids
|
||||
|
||||
@staticmethod
|
||||
def create_dummy(
|
||||
config: WhisperConfig,
|
||||
batch_size: int,
|
||||
encode_sequence_length: int,
|
||||
use_decoder_input_ids: int,
|
||||
device: torch.device,
|
||||
use_int32_inputs: bool = False,
|
||||
): # -> WhisperEncoderDecoderInitInputs:
|
||||
encoder_inputs: WhisperEncoderInputs = WhisperEncoderInputs.create_dummy(
|
||||
batch_size,
|
||||
sequence_length=3000,
|
||||
feature_size=config.num_mel_bins,
|
||||
device=device,
|
||||
use_int32_inputs=use_int32_inputs,
|
||||
)
|
||||
decoder_input_ids = None
|
||||
encoder_attention_mask = torch.zeros(
|
||||
(encoder_inputs.input_ids.shape[0], 1, encoder_inputs.input_ids.shape[1], encoder_inputs.input_ids.shape[1])
|
||||
).type(torch.int8)
|
||||
if use_decoder_input_ids:
|
||||
dtype = torch.int32 if use_int32_inputs else torch.int64
|
||||
decoder_input_ids = torch.ones((batch_size, 1), dtype=dtype, device=device) * config.decoder_start_token_id
|
||||
|
||||
return WhisperEncoderDecoderInitInputs(encoder_inputs.input_ids, encoder_attention_mask, decoder_input_ids)
|
||||
|
||||
def to_list(self) -> List:
|
||||
input_list = [self.encoder_input_ids, self.encoder_attention_mask]
|
||||
if self.decoder_input_ids is not None:
|
||||
input_list.append(self.decoder_input_ids)
|
||||
return input_list
|
||||
|
||||
|
||||
class WhisperEncoderDecoderInitHelper:
|
||||
@staticmethod
|
||||
def export_onnx(
|
||||
model: WhisperEncoderDecoderInit,
|
||||
device: torch.device,
|
||||
onnx_model_path: str,
|
||||
use_decoder_input_ids: bool = True,
|
||||
verbose: bool = True,
|
||||
use_external_data_format: bool = False,
|
||||
use_int32_inputs: bool = False,
|
||||
):
|
||||
"""Export decoder to ONNX
|
||||
|
||||
Args:
|
||||
model (WhisperEncoderDecoderInit): the model to export
|
||||
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.
|
||||
"""
|
||||
assert isinstance(model, WhisperEncoderDecoderInit)
|
||||
|
||||
inputs = WhisperEncoderDecoderInitInputs.create_dummy(
|
||||
model.config,
|
||||
batch_size=2,
|
||||
encode_sequence_length=3000,
|
||||
use_decoder_input_ids=use_decoder_input_ids,
|
||||
device=device,
|
||||
use_int32_inputs=use_int32_inputs,
|
||||
)
|
||||
input_list = inputs.to_list()
|
||||
|
||||
out = model(inputs.encoder_input_ids, inputs.encoder_attention_mask, inputs.decoder_input_ids)
|
||||
present = out[2]
|
||||
# pdb.set_trace()
|
||||
present_names = PastKeyValuesHelper.get_input_names(present, encoder=True)
|
||||
# present_names = PastKeyValuesHelper.get_past_names(model.config.num_layers, present=True)
|
||||
|
||||
output_names = ["logits", "encoder_hidden_states", *present_names]
|
||||
|
||||
# Shape of input tensors (sequence_length==1):
|
||||
# input_ids: (batch_size, sequence_length)
|
||||
# encoder_attention_mask: (batch_size, encode_sequence_length)
|
||||
# encoder_hidden_states: (batch_size, encode_sequence_length, hidden_size)
|
||||
# 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 = ["encoder_input_ids", "encoder_attention_mask"]
|
||||
|
||||
# ONNX exporter might mark dimension like 'Transposepresent_value_self_1_dim_2' in shape inference.
|
||||
# We use a workaround here: first use dim_param "1" for sequence_length, and later change to dim_value.
|
||||
sequence_length = "1"
|
||||
num_heads = str(model.config.encoder_attention_heads)
|
||||
hidden_size = str(model.config.d_model)
|
||||
head_size = str(model.config.d_model // model.config.encoder_attention_heads)
|
||||
dynamic_axes = {
|
||||
"encoder_input_ids": {0: "batch_size", 1: "encode_sequence_length"},
|
||||
"encoder_attention_mask": {0: "batch_size", 1: "encode_sequence_length"},
|
||||
"encoder_hidden_states": {
|
||||
0: "batch_size",
|
||||
1: "encode_sequence_length",
|
||||
2: hidden_size,
|
||||
},
|
||||
"logits": {
|
||||
0: "batch_size",
|
||||
1: sequence_length,
|
||||
},
|
||||
}
|
||||
|
||||
if use_decoder_input_ids:
|
||||
input_names.append("decoder_input_ids")
|
||||
dynamic_axes["decoder_input_ids"] = {
|
||||
0: "batch_size",
|
||||
1: sequence_length,
|
||||
}
|
||||
|
||||
for name in present_names:
|
||||
if "cross" in name:
|
||||
dynamic_axes[name] = {
|
||||
0: "batch_size",
|
||||
1: num_heads,
|
||||
2: "encode_sequence_length",
|
||||
3: head_size,
|
||||
}
|
||||
|
||||
else: # self attention past state
|
||||
dynamic_axes[name] = {
|
||||
0: "batch_size",
|
||||
1: num_heads,
|
||||
2: sequence_length,
|
||||
3: head_size,
|
||||
}
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp_dir_name:
|
||||
temp_onnx_model_path = os.path.join(tmp_dir_name, "encoder_decoder_init.onnx")
|
||||
Path(temp_onnx_model_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
torch_onnx_export(
|
||||
model,
|
||||
args=tuple(input_list),
|
||||
f=temp_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,
|
||||
)
|
||||
|
||||
# Workaround as mentioned earlier: change numeric dim_param to dim_value
|
||||
model = onnx.load(temp_onnx_model_path)
|
||||
for tensor in model.graph.output:
|
||||
for dim_proto in tensor.type.tensor_type.shape.dim:
|
||||
if dim_proto.HasField("dim_param") and dim_proto.dim_param in [
|
||||
sequence_length,
|
||||
num_heads,
|
||||
hidden_size,
|
||||
head_size,
|
||||
]:
|
||||
dim_value = int(dim_proto.dim_param)
|
||||
dim_proto.Clear()
|
||||
dim_proto.dim_value = dim_value
|
||||
|
||||
OnnxModel.save(
|
||||
model,
|
||||
onnx_model_path,
|
||||
save_as_external_data=use_external_data_format,
|
||||
all_tensors_to_one_file=True,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def onnxruntime_inference(ort_session, inputs: WhisperEncoderDecoderInitInputs):
|
||||
"""Run inference of ONNX model."""
|
||||
logger.debug("start onnxruntime_inference")
|
||||
|
||||
ort_inputs = {
|
||||
"encoder_input_ids": numpy.ascontiguousarray(inputs.encoder_input_ids.cpu().numpy()),
|
||||
"encoder_attention_mask": numpy.ascontiguousarray(inputs.encoder_attention_mask.cpu().numpy()),
|
||||
}
|
||||
if inputs.decoder_input_ids is not None:
|
||||
ort_inputs["decoder_input_ids"] = numpy.ascontiguousarray(inputs.decoder_input_ids.cpu().numpy())
|
||||
|
||||
ort_outputs = ort_session.run(None, ort_inputs)
|
||||
return ort_outputs
|
||||
|
||||
@staticmethod
|
||||
def verify_onnx(
|
||||
model: WhisperEncoderDecoderInit,
|
||||
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."""
|
||||
ort_inputs = ort_session.get_inputs()
|
||||
use_decoder_input_ids = len(ort_inputs) == 3
|
||||
|
||||
test_cases = [(4, 11), (1, 2), (3, 1), (8, 5)]
|
||||
test_cases_max_diff = []
|
||||
for batch_size, encode_sequence_length in test_cases[:max_cases]:
|
||||
inputs = WhisperEncoderDecoderInitInputs.create_dummy(
|
||||
model.config,
|
||||
batch_size,
|
||||
encode_sequence_length,
|
||||
use_decoder_input_ids=use_decoder_input_ids,
|
||||
device=device,
|
||||
use_int32_inputs=use_int32_inputs,
|
||||
)
|
||||
|
||||
ort_outputs = WhisperEncoderDecoderInitHelper.onnxruntime_inference(ort_session, inputs)
|
||||
|
||||
# Run inference of PyTorch model
|
||||
input_list = inputs.to_list()
|
||||
torch_outputs = model(*input_list)
|
||||
|
||||
assert torch_outputs[0].cpu().numpy().shape == ort_outputs[0].shape
|
||||
max_diff = numpy.amax(numpy.abs(torch_outputs[0].cpu().numpy() - ort_outputs[0]))
|
||||
logger.debug(f"logits max_diff={max_diff}")
|
||||
max_diff_all = max_diff
|
||||
|
||||
assert torch_outputs[1].cpu().numpy().shape == ort_outputs[1].shape
|
||||
max_diff = numpy.amax(numpy.abs(torch_outputs[1].cpu().numpy() - ort_outputs[1]))
|
||||
logger.debug(f"encoder_hidden_states max_diff={max_diff}")
|
||||
max_diff_all = max(max_diff_all, max_diff)
|
||||
|
||||
for i in range(2 * model.config.num_layers):
|
||||
max_diff = numpy.amax(numpy.abs(torch_outputs[2][i].cpu().numpy() - ort_outputs[2 + i]))
|
||||
logger.debug(f"self attention past state {i} max_diff={max_diff}")
|
||||
|
||||
for i in range(2 * model.config.num_layers):
|
||||
max_diff = numpy.amax(
|
||||
numpy.abs(torch_outputs[3][i].cpu().numpy() - ort_outputs[2 + 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(
|
||||
f"batch_size={batch_size} encode_sequence_length={encode_sequence_length}, max_diff={max_diff_all}"
|
||||
)
|
||||
|
||||
return max(test_cases_max_diff)
|
||||
|
|
@ -0,0 +1,265 @@
|
|||
# -------------------------------------------------------------------------
|
||||
# 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 sys
|
||||
from pathlib import Path
|
||||
from typing import Dict, Tuple, Union
|
||||
|
||||
import torch
|
||||
from transformers import WhisperForConditionalGeneration
|
||||
from whisper_decoder import WhisperDecoder, WhisperDecoderHelper, WhisperDecoderInit
|
||||
from whisper_encoder import WhisperEncoder, WhisperEncoderHelper
|
||||
from whisper_encoder_decoder_init import WhisperEncoderDecoderInit, WhisperEncoderDecoderInitHelper
|
||||
|
||||
from onnxruntime import InferenceSession
|
||||
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), "..", ".."))
|
||||
from float16 import float_to_float16_max_diff # noqa: E402
|
||||
from onnx_model import OnnxModel # noqa: E402
|
||||
from optimizer import optimize_model # noqa: E402
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
PRETRAINED_WHISPER_MODELS = [
|
||||
"whisper-tiny",
|
||||
"whisper-tiny.en",
|
||||
"whisper-small",
|
||||
"whisper-small.en",
|
||||
"whisper-medium",
|
||||
"whisper-medium.en",
|
||||
"whisper-base",
|
||||
"whisper-base.en",
|
||||
"whisper-large",
|
||||
"whisper-large-v2",
|
||||
]
|
||||
|
||||
|
||||
class WhisperHelper:
|
||||
@staticmethod
|
||||
def get_onnx_path(
|
||||
output_dir: str,
|
||||
model_name_or_path: str,
|
||||
suffix: str = "",
|
||||
new_folder: bool = False,
|
||||
) -> str:
|
||||
"""Build onnx path
|
||||
|
||||
Args:
|
||||
output_dir (str): output directory
|
||||
model_name_or_path (str): pretrained model name, or path to the model checkpoint
|
||||
suffix (str, optional): suffix like "_encoder" or "_decoder_fp16" will be appended to file name. Defaults to None.
|
||||
new_folder (bool, optional): create a new directory for the model. Defaults to False.
|
||||
|
||||
Returns:
|
||||
str: path of onnx model
|
||||
"""
|
||||
model_name = model_name_or_path
|
||||
if os.path.isdir(model_name_or_path):
|
||||
model_name = Path(model_name_or_path).parts[-1]
|
||||
else:
|
||||
model_name.split("/")[-1]
|
||||
|
||||
model_name += suffix
|
||||
|
||||
directory = os.path.join(output_dir, model_name) if new_folder else output_dir
|
||||
return os.path.join(directory, model_name + ".onnx")
|
||||
|
||||
@staticmethod
|
||||
def load_model(
|
||||
model_name_or_path: str, cache_dir: str, device: torch.device, merge_encoder_and_decoder_init: bool = True
|
||||
) -> Dict[str, torch.nn.Module]:
|
||||
"""Load model given a pretrained name or path, then build models for ONNX conversion.
|
||||
|
||||
Args:
|
||||
model_name_or_path (str): pretrained model name or path
|
||||
cache_dir (str): cache directory
|
||||
device (torch.device): device to run the model
|
||||
merge_encoder_and_decoder_init (bool, optional): Whether merge encoder and decoder initialization into one ONNX model. Defaults to True.
|
||||
Returns:
|
||||
Dict[str, torch.nn.Module]: mapping from name to modules for ONNX conversion.
|
||||
"""
|
||||
model = WhisperForConditionalGeneration.from_pretrained(model_name_or_path, cache_dir=cache_dir)
|
||||
|
||||
decoder = WhisperDecoder(model, None, model.config)
|
||||
decoder.eval().to(device)
|
||||
|
||||
if merge_encoder_and_decoder_init:
|
||||
encoder_decoder_init = WhisperEncoderDecoderInit(
|
||||
model,
|
||||
model,
|
||||
None,
|
||||
model.config,
|
||||
decoder_start_token_id=None,
|
||||
)
|
||||
return {"encoder_decoder_init": encoder_decoder_init, "decoder": decoder}
|
||||
else:
|
||||
encoder = WhisperEncoder(model.model.encoder, model.config)
|
||||
encoder.eval().to(device)
|
||||
decoder_init = WhisperDecoderInit(model.decoder, None, model.config)
|
||||
decoder_init.eval().to(device)
|
||||
return {
|
||||
"encoder": encoder,
|
||||
"decoder": decoder,
|
||||
"decoder_init": decoder_init,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def export_onnx(
|
||||
model: Union[WhisperEncoder, WhisperDecoder, WhisperDecoderInit, WhisperEncoderDecoderInit],
|
||||
device: torch.device,
|
||||
onnx_model_path: str,
|
||||
verbose: bool = True,
|
||||
use_external_data_format: bool = False,
|
||||
use_decoder_input_ids: bool = True,
|
||||
use_int32_inputs: bool = False,
|
||||
):
|
||||
if isinstance(model, WhisperEncoder):
|
||||
WhisperEncoderHelper.export_onnx(
|
||||
model,
|
||||
device,
|
||||
onnx_model_path,
|
||||
verbose,
|
||||
use_external_data_format,
|
||||
)
|
||||
elif isinstance(model, WhisperEncoderDecoderInit):
|
||||
WhisperEncoderDecoderInitHelper.export_onnx(
|
||||
model,
|
||||
device,
|
||||
onnx_model_path,
|
||||
use_decoder_input_ids,
|
||||
verbose,
|
||||
use_external_data_format,
|
||||
use_int32_inputs,
|
||||
)
|
||||
else:
|
||||
WhisperDecoderHelper.export_onnx(
|
||||
model,
|
||||
device,
|
||||
onnx_model_path,
|
||||
verbose,
|
||||
use_external_data_format,
|
||||
use_int32_inputs,
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def auto_mixed_precision(
|
||||
onnx_model: OnnxModel,
|
||||
op_block_list: Tuple[str] = (
|
||||
"SimplifiedLayerNormalization",
|
||||
"SkipSimplifiedLayerNormalization",
|
||||
"Relu",
|
||||
"Add",
|
||||
),
|
||||
):
|
||||
"""Convert model to mixed precision.
|
||||
It detects whether original model has fp16 precision weights, and set parameters for float16 conversion automatically.
|
||||
Args:
|
||||
onnx_model (OnnxModel): optimized ONNX model
|
||||
op_block_list (List[str], optional): . Defaults to ["SimplifiedLayerNormalization", "SkipSimplifiedLayerNormalization", "Relu", "Add"]
|
||||
Returns:
|
||||
parameters(dict): a dictionary of parameters used in float16 conversion
|
||||
"""
|
||||
op_full_set = set([node.op_type for node in onnx_model.nodes()])
|
||||
fp32_op_set = set(op_block_list)
|
||||
fp16_op_set = op_full_set.difference(fp32_op_set)
|
||||
logger.info(f"fp32 op: {fp32_op_set} fp16 op: {fp16_op_set}")
|
||||
|
||||
# logits is the first output
|
||||
logits_output_name = onnx_model.graph().output[0].name
|
||||
|
||||
# We use the weight in last MatMul node to detect whether the model is stored with float16 weights from training.
|
||||
is_weight_fp16_precision = False
|
||||
output_name_to_node = onnx_model.output_name_to_node()
|
||||
assert logits_output_name in output_name_to_node
|
||||
node = output_name_to_node[logits_output_name]
|
||||
last_matmul_node = None
|
||||
if node.op_type == "MatMul":
|
||||
last_matmul_node = node
|
||||
logger.info(f"Found last MatMul node for logits: {node.name}")
|
||||
initializer = None
|
||||
for input in node.input:
|
||||
initializer = onnx_model.get_initializer(input)
|
||||
if initializer is not None:
|
||||
break
|
||||
|
||||
# when the max difference of value after converting float to float16 is lower than a threshold (1e-6),
|
||||
# we can deduce that the weights are stored in float16 precision.
|
||||
max_diff = float_to_float16_max_diff(initializer)
|
||||
logger.debug(f"max diff of converting weights in last MatMul node {node.name}: {max_diff}")
|
||||
is_weight_fp16_precision = max_diff < 1e-6
|
||||
else:
|
||||
logger.warning(f"Failed to find MatMul node for logits. Found {node.op_type} of node {node.name}")
|
||||
|
||||
keep_io_types = []
|
||||
node_block_list = []
|
||||
if (not is_weight_fp16_precision) and (last_matmul_node is not None):
|
||||
# When original weight is float32 precision, keep logits and last MatMul in float32 could get better precision.
|
||||
keep_io_types = [logits_output_name]
|
||||
node_block_list = [last_matmul_node.name]
|
||||
|
||||
parameters = {
|
||||
"keep_io_types": keep_io_types,
|
||||
"op_block_list": list(op_block_list),
|
||||
"node_block_list": node_block_list,
|
||||
"force_fp16_initializers": is_weight_fp16_precision,
|
||||
}
|
||||
|
||||
logger.info(f"auto_mixed_precision parameters: {parameters}")
|
||||
onnx_model.convert_float_to_float16(use_symbolic_shape_infer=True, **parameters)
|
||||
|
||||
return parameters
|
||||
|
||||
@staticmethod
|
||||
def optimize_onnx(
|
||||
onnx_model_path: str,
|
||||
optimized_model_path: str,
|
||||
is_float16: bool,
|
||||
num_attention_heads: int,
|
||||
hidden_size: int,
|
||||
use_external_data_format: bool = False,
|
||||
auto_mixed_precision: bool = True,
|
||||
use_gpu: bool = False,
|
||||
):
|
||||
"""Optimize ONNX model with an option to convert it to use mixed precision."""
|
||||
|
||||
from fusion_options import FusionOptions
|
||||
|
||||
optimization_options = None
|
||||
if is_float16:
|
||||
optimization_options = FusionOptions("bart")
|
||||
optimization_options.enable_skip_layer_norm = False
|
||||
|
||||
m = optimize_model(
|
||||
onnx_model_path,
|
||||
model_type="bart",
|
||||
num_heads=num_attention_heads,
|
||||
hidden_size=hidden_size,
|
||||
opt_level=2 if not use_external_data_format else None,
|
||||
optimization_options=optimization_options,
|
||||
use_gpu=False,
|
||||
only_onnxruntime=False,
|
||||
)
|
||||
|
||||
if is_float16:
|
||||
if auto_mixed_precision:
|
||||
WhisperHelper.auto_mixed_precision(m)
|
||||
else:
|
||||
m.convert_model_float32_to_float16(cast_input_output=False)
|
||||
|
||||
m.save_model_to_file(optimized_model_path, use_external_data_format, all_tensors_to_one_file=True)
|
||||
|
||||
@staticmethod
|
||||
def verify_onnx(
|
||||
model: Union[WhisperEncoder, WhisperDecoder, WhisperDecoderInit, WhisperEncoderDecoderInit],
|
||||
ort_session: InferenceSession,
|
||||
device: torch.device,
|
||||
use_int32_inputs: bool,
|
||||
):
|
||||
"""Compare the result from PyTorch and OnnxRuntime to verify the ONNX model is good."""
|
||||
# Not implemented for Whisper currently
|
||||
return True
|
||||
Loading…
Reference in a new issue