mirror of
https://github.com/saymrwulf/onnxruntime.git
synced 2026-07-30 20:18:08 +00:00
Tianlei Wu
GitHub
parent
12eb660415
commit
379b3cdef6
7 changed files with 1087 additions and 0 deletions
|
|
@ -0,0 +1,4 @@
|
|||
import os
|
||||
import sys
|
||||
|
||||
sys.path.append(os.path.dirname(__file__))
|
||||
|
|
@ -0,0 +1,184 @@
|
|||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
import os
|
||||
import sys
|
||||
import argparse
|
||||
import logging
|
||||
import torch
|
||||
import copy
|
||||
from t5_helper import PRETRAINED_T5_MODELS, T5Helper
|
||||
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))
|
||||
from benchmark_helper import setup_logger, prepare_environment, create_onnxruntime_session, Precision
|
||||
|
||||
logger = logging.getLogger('')
|
||||
|
||||
|
||||
def parse_arguments():
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser.add_argument('-m',
|
||||
'--model_name_or_path',
|
||||
required=False,
|
||||
default=PRETRAINED_T5_MODELS[0],
|
||||
type=str,
|
||||
help='Model path, or pretrained model name in the list: ' + ', '.join(PRETRAINED_T5_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 in decoding. 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)
|
||||
|
||||
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 = True,
|
||||
merge_encoder_and_decoder_init: bool = True,
|
||||
overwrite: bool = False):
|
||||
device = torch.device("cuda:0" if use_gpu else "cpu")
|
||||
|
||||
models = T5Helper.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(f"Try use_external_data_format when model size > 2GB")
|
||||
|
||||
output_paths = []
|
||||
for name, model in models.items():
|
||||
filename_suffix = "_" + name
|
||||
|
||||
onnx_path = T5Helper.get_onnx_path(output_dir,
|
||||
model_name_or_path,
|
||||
suffix=filename_suffix,
|
||||
new_folder=use_external_data_format)
|
||||
|
||||
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.
|
||||
T5Helper.export_onnx(copy.deepcopy(model),
|
||||
device,
|
||||
onnx_path,
|
||||
verbose,
|
||||
use_external_data_format,
|
||||
use_decoder_input_ids=not use_decoder_start_token)
|
||||
else:
|
||||
logger.info(f"Skip exporting: existed ONNX model {onnx_path}")
|
||||
|
||||
# Optimize ONNX graph. Note that we have not implemented graph optimization for T5 yet.
|
||||
if optimize_onnx or precision != Precision.FLOAT32:
|
||||
output_path = T5Helper.get_onnx_path(output_dir,
|
||||
model_name_or_path,
|
||||
suffix=filename_suffix + "_" + str(precision),
|
||||
new_folder=use_external_data_format)
|
||||
|
||||
if overwrite or not os.path.exists(output_path):
|
||||
logger.info(f"Optimizing model to {output_path}")
|
||||
T5Helper.optimize_onnx(onnx_path, output_path, precision == Precision.FLOAT16, config.num_heads,
|
||||
config.hidden_size, use_external_data_format)
|
||||
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'])
|
||||
max_diff = T5Helper.verify_onnx(model, ort_session, device)
|
||||
logger.info(f'PyTorch and OnnxRuntime results max difference = {max_diff}')
|
||||
if max_diff > 1e-4:
|
||||
logger.warn(f'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.warn(f'Graph optimization for T5 is not implemented yet.')
|
||||
|
||||
with torch.no_grad():
|
||||
merge_encoder_and_decoder_init = True # Merge encoder and decoder initialization into one model is recommended.
|
||||
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, merge_encoder_and_decoder_init,
|
||||
args.overwrite)
|
||||
|
||||
logger.info(f"Done! Outputs: {output_paths}")
|
||||
|
||||
|
||||
if __name__ == '__main__':
|
||||
main()
|
||||
|
|
@ -0,0 +1,50 @@
|
|||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
import logging
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class PastKeyValuesHelper:
|
||||
""" Helper functions to process past key values for encoder-decoder model"""
|
||||
|
||||
@staticmethod
|
||||
def get_past_names(num_layers, present: bool = False):
|
||||
past_self_names = []
|
||||
past_cross_names = []
|
||||
for i in range(num_layers):
|
||||
past_self_names.extend([f'present_key_self_{i}', f'present_value_self_{i}']
|
||||
if present else [f'past_key_self_{i}', f'past_value_self_{i}'])
|
||||
past_cross_names.extend([f'present_key_cross_{i}', f'present_value_cross_{i}']
|
||||
if present else [f'past_key_cross_{i}', f'past_value_cross_{i}'])
|
||||
return past_self_names + past_cross_names
|
||||
|
||||
@staticmethod
|
||||
def group_by_self_or_cross(present_key_values):
|
||||
"""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, ...)
|
||||
|
||||
"""
|
||||
present_self = []
|
||||
present_cross = []
|
||||
for i, 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])
|
||||
return present_self, present_cross
|
||||
|
||||
@staticmethod
|
||||
def group_by_layer(past, num_layers):
|
||||
"""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),
|
||||
"""
|
||||
assert len(past) == 4 * num_layers
|
||||
return tuple([past[2 * i], past[2 * i + 1], past[2 * num_layers + 2 * i], past[2 * num_layers + 2 * i + 1]]
|
||||
for i in range(num_layers))
|
||||
335
onnxruntime/python/tools/transformers/models/t5/t5_decoder.py
Normal file
335
onnxruntime/python/tools/transformers/models/t5/t5_decoder.py
Normal file
|
|
@ -0,0 +1,335 @@
|
|||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
from pathlib import Path
|
||||
from typing import List, Union
|
||||
import logging
|
||||
import numpy
|
||||
import torch
|
||||
from transformers import T5Config
|
||||
from onnxruntime import InferenceSession
|
||||
from t5_encoder import T5EncoderInputs
|
||||
from past_helper import PastKeyValuesHelper
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class T5DecoderInit(torch.nn.Module):
|
||||
""" A T5 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: T5Config,
|
||||
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):
|
||||
if decoder_input_ids is None:
|
||||
batch_size = encoder_attention_mask.shape[0]
|
||||
decoder_input_ids = torch.ones(
|
||||
(batch_size, 1), dtype=torch.long, device=encoder_attention_mask.device) * self.decoder_start_token_id
|
||||
|
||||
decoder_outputs = self.decoder(input_ids=decoder_input_ids,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
use_cache=True,
|
||||
return_dict=True)
|
||||
|
||||
sequence_output = decoder_outputs.last_hidden_state
|
||||
present_key_values = decoder_outputs.past_key_values
|
||||
|
||||
sequence_output = sequence_output * (self.config.d_model**-0.5)
|
||||
|
||||
lm_logits = self.lm_head(sequence_output)
|
||||
past_self, past_cross = PastKeyValuesHelper.group_by_self_or_cross(present_key_values)
|
||||
return lm_logits, past_self, past_cross
|
||||
|
||||
|
||||
class T5Decoder(torch.nn.Module):
|
||||
""" A T5 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, encoder_hidden_states, *past):
|
||||
|
||||
past_key_values = PastKeyValuesHelper.group_by_layer(past, self.config.num_layers)
|
||||
|
||||
decoder_outputs = self.decoder(input_ids=decoder_input_ids,
|
||||
past_key_values=past_key_values,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
encoder_attention_mask=encoder_attention_mask,
|
||||
use_cache=True,
|
||||
return_dict=True)
|
||||
|
||||
sequence_output = decoder_outputs.last_hidden_state
|
||||
present_key_values = decoder_outputs.past_key_values
|
||||
|
||||
sequence_output = sequence_output * (self.config.d_model**-0.5)
|
||||
|
||||
lm_logits = self.lm_head(sequence_output)
|
||||
present_self, _ = PastKeyValuesHelper.group_by_self_or_cross(present_key_values)
|
||||
|
||||
# Do not return present_cross since they are identical to corresponding past_cross input
|
||||
return lm_logits, present_self
|
||||
|
||||
|
||||
class T5DecoderInputs:
|
||||
|
||||
def __init__(self, decoder_input_ids, encoder_attention_mask, encoder_hidden_states, past_key_values=None):
|
||||
self.decoder_input_ids: torch.LongTensor = decoder_input_ids
|
||||
self.encoder_attention_mask: torch.LongTensor = encoder_attention_mask
|
||||
self.encoder_hidden_states: Union[torch.FloatTensor, torch.HalfTensor] = encoder_hidden_states
|
||||
self.past_key_values: Union[List[torch.FloatTensor], List[torch.HalfTensor], None] = past_key_values
|
||||
|
||||
@staticmethod
|
||||
def create_dummy(config: T5Config, batch_size: int, encode_sequence_length: int, past_decode_sequence_length: int,
|
||||
device: torch.device): # -> T5DecoderInputs:
|
||||
""" Create dummy inputs for T5Decoder.
|
||||
|
||||
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
|
||||
|
||||
Returns:
|
||||
T5DecoderInputs: dummy inputs for decoder
|
||||
"""
|
||||
hidden_size: int = config.d_model
|
||||
num_attention_heads: int = config.num_heads
|
||||
num_layers: int = config.num_layers
|
||||
vocab_size: int = config.vocab_size
|
||||
|
||||
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.int64,
|
||||
device=device)
|
||||
|
||||
encoder_inputs = T5EncoderInputs.create_dummy(batch_size, encode_sequence_length, vocab_size, device)
|
||||
|
||||
encoder_hidden_state = torch.rand(batch_size,
|
||||
encode_sequence_length,
|
||||
hidden_size,
|
||||
dtype=torch.float32,
|
||||
device=device)
|
||||
|
||||
if past_decode_sequence_length > 0:
|
||||
self_attention_past_shape = [
|
||||
batch_size, num_attention_heads, past_decode_sequence_length,
|
||||
int(hidden_size / num_attention_heads)
|
||||
]
|
||||
cross_attention_past_shape = [
|
||||
batch_size, num_attention_heads, encode_sequence_length,
|
||||
int(hidden_size / num_attention_heads)
|
||||
]
|
||||
|
||||
past = []
|
||||
for _ in range(2 * num_layers):
|
||||
past.append(torch.rand(self_attention_past_shape, dtype=torch.float32, device=device))
|
||||
|
||||
for _ in range(2 * num_layers):
|
||||
past.append(torch.rand(cross_attention_past_shape, dtype=torch.float32, device=device))
|
||||
else:
|
||||
past = None
|
||||
|
||||
return T5DecoderInputs(decoder_input_ids, encoder_inputs.attention_mask, encoder_hidden_state, past)
|
||||
|
||||
def to_list(self) -> List:
|
||||
input_list = [self.decoder_input_ids, self.encoder_attention_mask, self.encoder_hidden_states]
|
||||
if self.past_key_values:
|
||||
input_list.extend(self.past_key_values)
|
||||
return input_list
|
||||
|
||||
|
||||
class T5DecoderHelper:
|
||||
|
||||
@staticmethod
|
||||
def export_onnx(decoder: Union[T5Decoder, T5DecoderInit],
|
||||
device: torch.device,
|
||||
onnx_model_path: str,
|
||||
verbose: bool = True,
|
||||
use_external_data_format: bool = False):
|
||||
"""Export decoder to ONNX
|
||||
|
||||
Args:
|
||||
decoder (Union[T5Decoder, T5DecoderNoPastState]): 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.
|
||||
"""
|
||||
assert isinstance(decoder, (T5Decoder, T5DecoderInit))
|
||||
|
||||
inputs = T5DecoderInputs.create_dummy(decoder.config,
|
||||
batch_size=2,
|
||||
encode_sequence_length=3,
|
||||
past_decode_sequence_length=5 if isinstance(decoder, T5Decoder) else 0,
|
||||
device=device)
|
||||
input_list = inputs.to_list()
|
||||
with torch.no_grad():
|
||||
outputs = decoder(*input_list)
|
||||
|
||||
past_names = PastKeyValuesHelper.get_past_names(decoder.config.num_layers, present=False)
|
||||
present_names = PastKeyValuesHelper.get_past_names(decoder.config.num_layers, present=True)
|
||||
present_self_names = present_names[:2 * decoder.config.num_layers]
|
||||
|
||||
input_past_names = past_names if isinstance(decoder, T5Decoder) else []
|
||||
output_present_names = present_self_names if isinstance(decoder, T5Decoder) 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)
|
||||
# encoder_hidden_states: (batch_size, encode_sequence_length, hidden_size)
|
||||
# past_self_*: (batch_size, num_heads, past_decode_sequence_length, hidden_size/num_heads)
|
||||
# past_cross_*: (batch_size, num_heads, encode_sequence_length, hidden_size/num_heads)
|
||||
|
||||
# Shape of output tensors:
|
||||
# logits: (batch_size, sequence_length, vocab_size)
|
||||
# past_self_*: (batch_size, num_heads, past_decode_sequence_length + sequence_length, hidden_size/num_heads)
|
||||
# past_cross_*: (batch_size, num_heads, encode_sequence_length, hidden_size/num_heads)
|
||||
|
||||
input_names = ["input_ids"]
|
||||
input_names.append("encoder_attention_mask")
|
||||
input_names.append("encoder_hidden_states")
|
||||
input_names.extend(input_past_names)
|
||||
|
||||
dynamic_axes = {
|
||||
'input_ids': {
|
||||
0: 'batch_size',
|
||||
#1: 'sequence_length'
|
||||
},
|
||||
'encoder_attention_mask': {
|
||||
0: 'batch_size',
|
||||
1: 'encode_sequence_length'
|
||||
},
|
||||
'encoder_hidden_states': {
|
||||
0: 'batch_size',
|
||||
1: 'encode_sequence_length'
|
||||
},
|
||||
"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, T5Decoder):
|
||||
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)
|
||||
torch.onnx.export(decoder,
|
||||
args=tuple(input_list),
|
||||
f=onnx_model_path,
|
||||
export_params=True,
|
||||
input_names=input_names,
|
||||
output_names=output_names,
|
||||
example_outputs=outputs,
|
||||
dynamic_axes=dynamic_axes,
|
||||
opset_version=12,
|
||||
do_constant_folding=True,
|
||||
use_external_data_format=use_external_data_format,
|
||||
verbose=verbose)
|
||||
|
||||
@staticmethod
|
||||
def onnxruntime_inference(ort_session, inputs: T5DecoderInputs):
|
||||
""" Run inference of ONNX model.
|
||||
"""
|
||||
logger.debug(f"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()),
|
||||
'encoder_hidden_states': numpy.ascontiguousarray(inputs.encoder_hidden_states.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[T5Decoder, T5DecoderInit],
|
||||
ort_session: InferenceSession,
|
||||
device: torch.device,
|
||||
max_cases=4):
|
||||
""" Compare the result from PyTorch and OnnxRuntime to verify the ONNX model is good.
|
||||
"""
|
||||
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, T5DecoderInit):
|
||||
past_decode_sequence_length = 0
|
||||
|
||||
inputs = T5DecoderInputs.create_dummy(model.config,
|
||||
batch_size,
|
||||
encode_sequence_length,
|
||||
past_decode_sequence_length,
|
||||
device=device)
|
||||
input_list = inputs.to_list()
|
||||
|
||||
# Run inference of PyTorch model
|
||||
with torch.no_grad():
|
||||
torch_outputs = model(*input_list)
|
||||
|
||||
ort_outputs = T5DecoderHelper.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, T5DecoderInit):
|
||||
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}, past_decode_sequence_length={past_decode_sequence_length}, max_diff={max_diff_all}"
|
||||
)
|
||||
|
||||
return max_diff_all
|
||||
150
onnxruntime/python/tools/transformers/models/t5/t5_encoder.py
Normal file
150
onnxruntime/python/tools/transformers/models/t5/t5_encoder.py
Normal file
|
|
@ -0,0 +1,150 @@
|
|||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
import random
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
import logging
|
||||
import numpy
|
||||
import torch
|
||||
from transformers import T5Config
|
||||
from onnxruntime import InferenceSession
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class T5Encoder(torch.nn.Module):
|
||||
""" T5 encoder outputs only the last hidden state"""
|
||||
|
||||
def __init__(self, encoder, config: T5Config):
|
||||
super().__init__()
|
||||
self.encoder = encoder
|
||||
self.config = config
|
||||
|
||||
def forward(self, input_ids, attention_mask):
|
||||
return self.encoder(input_ids, attention_mask)[0]
|
||||
|
||||
|
||||
class T5EncoderInputs:
|
||||
|
||||
def __init__(self, input_ids, attention_mask):
|
||||
self.input_ids: torch.LongTensor = input_ids
|
||||
self.attention_mask: torch.LongTensor = attention_mask
|
||||
|
||||
@staticmethod
|
||||
def create_dummy(batch_size: int, sequence_length: int, vocab_size: int,
|
||||
device: torch.device): # -> T5EncoderInputs
|
||||
""" Create dummy inputs for T5 encoder.
|
||||
|
||||
Args:
|
||||
batch_size (int): batch size
|
||||
sequence_length (int): sequence length
|
||||
vocab_size (int): vocaburary size
|
||||
device (torch.device): device of output tensors
|
||||
|
||||
Returns:
|
||||
T5EncoderInputs: dummy inputs for encoder
|
||||
"""
|
||||
input_ids = torch.randint(low=0,
|
||||
high=vocab_size - 1,
|
||||
size=(batch_size, sequence_length),
|
||||
dtype=torch.int64,
|
||||
device=device)
|
||||
|
||||
attention_mask = torch.ones([batch_size, sequence_length], dtype=torch.int64, device=device)
|
||||
if sequence_length >= 2:
|
||||
for i in range(batch_size):
|
||||
padding_position = random.randint(0, sequence_length - 1)
|
||||
attention_mask[i, :padding_position] = 0
|
||||
return T5EncoderInputs(input_ids, attention_mask)
|
||||
|
||||
def to_list(self) -> List:
|
||||
input_list = [v for v in [self.input_ids, self.attention_mask] if v is not None]
|
||||
return input_list
|
||||
|
||||
|
||||
class T5EncoderHelper:
|
||||
|
||||
@staticmethod
|
||||
def export_onnx(encoder: T5Encoder,
|
||||
device: torch.device,
|
||||
onnx_model_path: str,
|
||||
verbose: bool = True,
|
||||
use_external_data_format: bool = False):
|
||||
"""Export encoder to ONNX
|
||||
|
||||
Args:
|
||||
encoder (T5Encoder): 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 = T5EncoderInputs.create_dummy(batch_size=2,
|
||||
sequence_length=4,
|
||||
vocab_size=config.vocab_size,
|
||||
device=device)
|
||||
|
||||
with torch.no_grad():
|
||||
outputs = encoder(encoder_inputs.input_ids, encoder_inputs.attention_mask)
|
||||
|
||||
Path(onnx_model_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
torch.onnx.export(encoder,
|
||||
args=tuple(encoder_inputs.to_list()),
|
||||
f=onnx_model_path,
|
||||
export_params=True,
|
||||
input_names=['input_ids', 'attention_mask'],
|
||||
output_names=['hidden_states'],
|
||||
example_outputs=outputs,
|
||||
dynamic_axes={
|
||||
'input_ids': {
|
||||
0: 'batch_size',
|
||||
1: 'sequence_length'
|
||||
},
|
||||
'attention_mask': {
|
||||
0: 'batch_size',
|
||||
1: 'sequence_length'
|
||||
},
|
||||
'hidden_states': {
|
||||
0: 'batch_size',
|
||||
1: 'sequence_length'
|
||||
},
|
||||
},
|
||||
opset_version=12,
|
||||
do_constant_folding=True,
|
||||
use_external_data_format=use_external_data_format,
|
||||
verbose=verbose)
|
||||
|
||||
@staticmethod
|
||||
def onnxruntime_inference(ort_session, inputs: T5EncoderInputs):
|
||||
""" 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: T5Encoder, ort_session: InferenceSession, device: torch.device):
|
||||
""" Compare the result from PyTorch and OnnxRuntime to verify the ONNX model is good.
|
||||
"""
|
||||
inputs = T5EncoderInputs.create_dummy(batch_size=4,
|
||||
sequence_length=11,
|
||||
vocab_size=model.config.vocab_size,
|
||||
device=device)
|
||||
input_list = inputs.to_list()
|
||||
torch_outputs = model(*input_list)
|
||||
|
||||
ort_outputs = T5EncoderHelper.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,231 @@
|
|||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
from pathlib import Path
|
||||
from typing import List
|
||||
import logging
|
||||
import numpy
|
||||
import torch
|
||||
from transformers import T5Config
|
||||
from onnxruntime import InferenceSession
|
||||
from t5_encoder import T5Encoder, T5EncoderInputs
|
||||
from t5_decoder import T5DecoderInit
|
||||
from past_helper import PastKeyValuesHelper
|
||||
|
||||
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.Module,
|
||||
config: T5Config,
|
||||
decoder_start_token_id: int = None):
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.t5_encoder = T5Encoder(encoder, config)
|
||||
self.t5_decoder_init = T5DecoderInit(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.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)
|
||||
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 = decoder_input_ids
|
||||
|
||||
@staticmethod
|
||||
def create_dummy(config: T5Config, batch_size: int, encode_sequence_length: int, use_decoder_input_ids: int,
|
||||
device: torch.device): # -> T5EncoderDecoderInitInputs:
|
||||
encoder_inputs: T5EncoderInputs = T5EncoderInputs.create_dummy(batch_size, encode_sequence_length,
|
||||
config.vocab_size, device)
|
||||
decoder_input_ids = None
|
||||
if use_decoder_input_ids:
|
||||
decoder_input_ids = torch.ones(
|
||||
(batch_size, 1), dtype=torch.long, 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:
|
||||
|
||||
@staticmethod
|
||||
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):
|
||||
"""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.
|
||||
"""
|
||||
assert isinstance(model, T5EncoderDecoderInit)
|
||||
|
||||
inputs = T5EncoderDecoderInitInputs.create_dummy(model.config,
|
||||
batch_size=2,
|
||||
encode_sequence_length=3,
|
||||
use_decoder_input_ids=use_decoder_input_ids,
|
||||
device=device)
|
||||
input_list = inputs.to_list()
|
||||
outputs = model(*input_list)
|
||||
|
||||
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, hidden_size/num_heads)
|
||||
# past_cross_*: (batch_size, num_heads, encode_sequence_length, hidden_size/num_heads)
|
||||
|
||||
# Shape of output tensors:
|
||||
# logits: (batch_size, sequence_length, vocab_size)
|
||||
# past_self_*: (batch_size, num_heads, past_decode_sequence_length + sequence_length, hidden_size/num_heads)
|
||||
# past_cross_*: (batch_size, num_heads, encode_sequence_length, hidden_size/num_heads)
|
||||
|
||||
input_names = ["encoder_input_ids", "encoder_attention_mask"]
|
||||
|
||||
# ONNX exporter might mark dimension like 'Transposepresent_value_self_1_dim_2'. Use more friendly string here.
|
||||
sequence_length = '1'
|
||||
num_heads = str(model.config.num_heads)
|
||||
hidden_size = str(model.config.d_model)
|
||||
head_size = str(model.config.d_model // model.config.num_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}
|
||||
|
||||
Path(onnx_model_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
torch.onnx.export(model,
|
||||
args=tuple(input_list),
|
||||
f=onnx_model_path,
|
||||
export_params=True,
|
||||
input_names=input_names,
|
||||
output_names=output_names,
|
||||
example_outputs=outputs,
|
||||
dynamic_axes=dynamic_axes,
|
||||
opset_version=12,
|
||||
do_constant_folding=True,
|
||||
use_external_data_format=use_external_data_format,
|
||||
verbose=verbose)
|
||||
|
||||
@staticmethod
|
||||
def onnxruntime_inference(ort_session, inputs: T5EncoderDecoderInitInputs):
|
||||
""" Run inference of ONNX model.
|
||||
"""
|
||||
logger.debug(f"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: T5EncoderDecoderInit, ort_session: InferenceSession, device: torch.device, max_cases=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)
|
||||
|
||||
ort_outputs = T5EncoderDecoderInitHelper.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)
|
||||
133
onnxruntime/python/tools/transformers/models/t5/t5_helper.py
Normal file
133
onnxruntime/python/tools/transformers/models/t5/t5_helper.py
Normal file
|
|
@ -0,0 +1,133 @@
|
|||
# -------------------------------------------------------------------------
|
||||
# Copyright (c) Microsoft Corporation. All rights reserved.
|
||||
# Licensed under the MIT License. See License.txt in the project root for
|
||||
# license information.
|
||||
# --------------------------------------------------------------------------
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Union, Dict
|
||||
import logging
|
||||
import torch
|
||||
from transformers import T5ForConditionalGeneration
|
||||
from onnxruntime import InferenceSession
|
||||
from t5_encoder import T5Encoder, T5EncoderHelper
|
||||
from t5_decoder import T5DecoderInit, T5Decoder, T5DecoderHelper
|
||||
from t5_encoder_decoder_init import T5EncoderDecoderInit, T5EncoderDecoderInitHelper
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
PRETRAINED_T5_MODELS = ["t5-small", "t5-base", "t5-large", "t5-3B", "t5-11B"]
|
||||
|
||||
|
||||
class T5Helper:
|
||||
|
||||
@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
|
||||
|
||||
dir = os.path.join(output_dir, model_name) if new_folder else output_dir
|
||||
return os.path.join(dir, 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 = T5ForConditionalGeneration.from_pretrained(model_name_or_path, cache_dir=cache_dir)
|
||||
|
||||
decoder = T5Decoder(model.decoder, model.lm_head, model.config)
|
||||
decoder.eval().to(device)
|
||||
|
||||
if merge_encoder_and_decoder_init:
|
||||
encoder_decoder_init = T5EncoderDecoderInit(model.encoder,
|
||||
model.decoder,
|
||||
model.lm_head,
|
||||
model.config,
|
||||
decoder_start_token_id=None)
|
||||
return {"encoder_decoder_init": encoder_decoder_init, "decoder": decoder}
|
||||
else:
|
||||
encoder = T5Encoder(model.encoder, model.config)
|
||||
encoder.eval().to(device)
|
||||
decoder_init = T5DecoderInit(model.decoder, model.lm_head, model.config)
|
||||
decoder_init.eval().to(device)
|
||||
return {"encoder": encoder, "decoder": decoder, "decoder_init": decoder_init}
|
||||
|
||||
@staticmethod
|
||||
def export_onnx(model: Union[T5Encoder, T5Decoder, T5DecoderInit, T5EncoderDecoderInit],
|
||||
device: torch.device,
|
||||
onnx_model_path: str,
|
||||
verbose: bool = True,
|
||||
use_external_data_format: bool = False,
|
||||
use_decoder_input_ids: bool = True):
|
||||
if isinstance(model, T5Encoder):
|
||||
T5EncoderHelper.export_onnx(model, device, onnx_model_path, verbose, use_external_data_format)
|
||||
elif isinstance(model, T5EncoderDecoderInit):
|
||||
T5EncoderDecoderInitHelper.export_onnx(model, device, onnx_model_path, use_decoder_input_ids, verbose,
|
||||
use_external_data_format)
|
||||
else:
|
||||
T5DecoderHelper.export_onnx(model, device, onnx_model_path, verbose, use_external_data_format)
|
||||
|
||||
@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):
|
||||
""" Optimize ONNX model with an option to convert it to use mixed precision.
|
||||
"""
|
||||
sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))
|
||||
from optimizer import optimize_model
|
||||
m = optimize_model(onnx_model_path,
|
||||
model_type='bert',
|
||||
num_heads=num_attention_heads,
|
||||
hidden_size=hidden_size,
|
||||
opt_level=0,
|
||||
optimization_options=None,
|
||||
use_gpu=False)
|
||||
if is_float16:
|
||||
m.convert_model_float32_to_float16(cast_input_output=False)
|
||||
|
||||
m.save_model_to_file(optimized_model_path, use_external_data_format)
|
||||
|
||||
@staticmethod
|
||||
def verify_onnx(model: Union[T5Encoder, T5Decoder, T5DecoderInit, T5EncoderDecoderInit],
|
||||
ort_session: InferenceSession, device: torch.device):
|
||||
""" Compare the result from PyTorch and OnnxRuntime to verify the ONNX model is good.
|
||||
"""
|
||||
if isinstance(model, T5Encoder):
|
||||
return T5EncoderHelper.verify_onnx(model, ort_session, device)
|
||||
elif isinstance(model, T5EncoderDecoderInit):
|
||||
return T5EncoderDecoderInitHelper.verify_onnx(model, ort_session, device)
|
||||
else:
|
||||
return T5DecoderHelper.verify_onnx(model, ort_session, device)
|
||||
Loading…
Reference in a new issue