Instructions to use CatSnowdrop/universal_onnxruntime with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Libraries
- TensorRT
How to use CatSnowdrop/universal_onnxruntime with TensorRT:
# No code snippets available yet for this library. # To use this model, check the repository files and the library's documentation. # Want to help? PRs adding snippets are welcome at: # https://github.com/huggingface/huggingface.js
- Notebooks
- Google Colab
- Kaggle
Download onnxruntime/transformers/models/t5/t5_encoder_decoder_init.py from CatSnowdrop/universal_onnxruntime: direct link, hf CLI and curl.
- Browser
- Download file 15.1 kB
-
https://huggingface.co/CatSnowdrop/universal_onnxruntime/resolve/main/onnxruntime/transformers/models/t5/t5_encoder_decoder_init.py
- Command line
-
hf download hf://CatSnowdrop/universal_onnxruntime/onnxruntime/transformers/models/t5/t5_encoder_decoder_init.py
-
curl -L -o t5_encoder_decoder_init.py https://huggingface.co/CatSnowdrop/universal_onnxruntime/resolve/main/onnxruntime/transformers/models/t5/t5_encoder_decoder_init.py
15.1 kB
| # ------------------------------------------------------------------------- | |
| # Copyright (c) Microsoft Corporation. All rights reserved. | |
| # Licensed under the MIT License. | |
| # ------------------------------------------------------------------------- | |
| import logging | |
| import os | |
| import tempfile | |
| from pathlib import Path | |
| import numpy | |
| import onnx | |
| import torch | |
| from onnx_model import OnnxModel | |
| from past_helper import PastKeyValuesHelper | |
| from t5_decoder import T5DecoderInit | |
| from t5_encoder import T5Encoder, T5EncoderInputs | |
| from torch_onnx_export_helper import torch_onnx_export | |
| from transformers import MT5Config, T5Config | |
| from onnxruntime import InferenceSession | |
| logger = logging.getLogger(__name__) | |
| class T5EncoderDecoderInit(torch.nn.Module): | |
| """A combination of T5Encoder and T5DecoderInit.""" | |
| def __init__( | |
| self, | |
| encoder: torch.nn.Module, | |
| decoder: torch.nn.Module, | |
| lm_head: torch.nn.Linear, | |
| config: T5Config | MT5Config, | |
| decoder_start_token_id: int | None = None, | |
| output_cross_only: bool = False, | |
| ): | |
| super().__init__() | |
| self.config: T5Config | MT5Config = config | |
| self.t5_encoder = T5Encoder(encoder, config) | |
| self.t5_decoder_init = T5DecoderInit(decoder, lm_head, config, decoder_start_token_id) | |
| self.output_cross_only = output_cross_only | |
| def forward( | |
| self, | |
| encoder_input_ids: torch.Tensor, | |
| encoder_attention_mask: torch.Tensor, | |
| decoder_input_ids: torch.Tensor | None = None, | |
| ): | |
| encoder_hidden_states: torch.FloatTensor = self.t5_encoder(encoder_input_ids, encoder_attention_mask) | |
| lm_logits, past_self, past_cross = self.t5_decoder_init( | |
| decoder_input_ids, encoder_attention_mask, encoder_hidden_states | |
| ) | |
| if self.output_cross_only: | |
| return past_cross | |
| else: | |
| return lm_logits, encoder_hidden_states, past_self, past_cross | |
| class T5EncoderDecoderInitInputs: | |
| 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 | None = decoder_input_ids | |
| def create_dummy( | |
| config: T5Config | MT5Config, | |
| batch_size: int, | |
| encode_sequence_length: int, | |
| use_decoder_input_ids: int, | |
| device: torch.device, | |
| use_int32_inputs: bool = False, | |
| ): # -> T5EncoderDecoderInitInputs: | |
| encoder_inputs: T5EncoderInputs = T5EncoderInputs.create_dummy( | |
| batch_size, | |
| encode_sequence_length, | |
| config.vocab_size, | |
| device, | |
| use_int32_inputs=use_int32_inputs, | |
| ) | |
| decoder_input_ids = None | |
| 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 T5EncoderDecoderInitInputs(encoder_inputs.input_ids, encoder_inputs.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 T5EncoderDecoderInitHelper: | |
| def export_onnx( | |
| model: T5EncoderDecoderInit, | |
| 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 (T5EncoderDecoderInit): 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. | |
| use_int32_inputs (bool, optional): use int32 instead of int64 for integer inputs. Defaults to False. | |
| """ | |
| assert isinstance(model, T5EncoderDecoderInit) | |
| # Do not exclude decoder in torch onnx export so that cross can show up. | |
| output_cross_only = model.output_cross_only | |
| model.output_cross_only = False | |
| inputs = T5EncoderDecoderInitInputs.create_dummy( | |
| model.config, | |
| batch_size=2, | |
| encode_sequence_length=3, | |
| use_decoder_input_ids=use_decoder_input_ids, | |
| device=device, | |
| use_int32_inputs=use_int32_inputs, | |
| ) | |
| input_list = inputs.to_list() | |
| present_names = PastKeyValuesHelper.get_past_names(model.config.num_decoder_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 'present_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.num_heads) | |
| hidden_size = str(model.config.d_model) | |
| head_size = str(model.config.d_kv) | |
| 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=12, | |
| do_constant_folding=True, | |
| use_external_data_format=use_external_data_format, | |
| verbose=verbose, | |
| ) | |
| # Restore output_cross_only setting. | |
| model.output_cross_only = output_cross_only | |
| # Workaround as mentioned earlier: change numeric dim_param to dim_value | |
| exported_model: onnx.ModelProto = onnx.load(temp_onnx_model_path) | |
| for tensor in exported_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 | |
| if output_cross_only: | |
| # Rewrite onnx graph to only keep present_[key|value]_cross_* outputs. | |
| onnx_model = OnnxModel(exported_model) | |
| output_name_to_node = onnx_model.output_name_to_node() | |
| for output in exported_model.graph.output: | |
| if "cross" in output.name: | |
| assert output.name in output_name_to_node | |
| transpose_node = output_name_to_node[output.name] | |
| assert transpose_node and transpose_node.op_type == "Transpose" | |
| permutation = OnnxModel.get_node_attribute(transpose_node, "perm") | |
| assert isinstance(permutation, list) | |
| assert permutation == [0, 2, 1, 3] | |
| matched_nodes = onnx_model.match_parent_path( | |
| transpose_node, | |
| ["Reshape", "MatMul"], | |
| [0, 0], | |
| output_name_to_node, | |
| ) | |
| assert matched_nodes is not None | |
| reshape_node, matmul_node = matched_nodes | |
| assert "encoder_hidden_states" in matmul_node.input | |
| if not onnx_model.get_initializer("cross_reshape_shape"): | |
| shape_tensor = onnx.helper.make_tensor( | |
| name="cross_reshape_shape", | |
| data_type=onnx.TensorProto.INT64, | |
| dims=[4], | |
| vals=[0, 0, int(num_heads), int(head_size)], | |
| raw=False, | |
| ) | |
| onnx_model.add_initializer(shape_tensor) | |
| reshape_node.input[1] = "cross_reshape_shape" | |
| cross_outputs = [output.name for output in exported_model.graph.output if "cross" in output.name] | |
| onnx_model.prune_graph(cross_outputs, allow_remove_graph_inputs=True) | |
| OnnxModel.save( | |
| exported_model, | |
| onnx_model_path, | |
| save_as_external_data=use_external_data_format, | |
| all_tensors_to_one_file=True, | |
| ) | |
| def onnxruntime_inference(ort_session, inputs: T5EncoderDecoderInitInputs): | |
| """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 | |
| def verify_onnx( | |
| model: T5EncoderDecoderInit, | |
| 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 = T5EncoderDecoderInitInputs.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 = T5EncoderDecoderInitHelper.onnxruntime_inference(ort_session, inputs) | |
| # Run inference of PyTorch model | |
| input_list = inputs.to_list() | |
| torch_outputs = model(*input_list) | |
| num_decoder_layers = model.config.num_decoder_layers | |
| if not model.output_cross_only: | |
| 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 * num_decoder_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 * num_decoder_layers): | |
| max_diff = numpy.amax( | |
| numpy.abs(torch_outputs[3][i].cpu().numpy() - ort_outputs[2 + 2 * num_decoder_layers + i]) | |
| ) | |
| logger.debug(f"cross attention past state {i} max_diff={max_diff}") | |
| max_diff_all = max(max_diff_all, max_diff) | |
| else: | |
| max_diff_all = -float("inf") | |
| for i in range(2 * num_decoder_layers): | |
| max_diff = numpy.amax(numpy.abs(torch_outputs[i].cpu().numpy() - ort_outputs[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) | |