diff --git a/onnxruntime/python/tools/transformers/benchmark_gpt2.py b/onnxruntime/python/tools/transformers/benchmark_gpt2.py index 0ca2742e24..d48b9f6170 100644 --- a/onnxruntime/python/tools/transformers/benchmark_gpt2.py +++ b/onnxruntime/python/tools/transformers/benchmark_gpt2.py @@ -28,11 +28,11 @@ def parse_arguments(): parser = argparse.ArgumentParser() parser.add_argument('-m', - '--model_name', + '--model_name_or_path', required=True, type=str, - choices=PRETRAINED_GPT2_MODELS, - help='Pretrained model selected in the list: ' + ', '.join(PRETRAINED_GPT2_MODELS)) + help='Model path, or pretrained model name selected in the list: ' + + ', '.join(PRETRAINED_GPT2_MODELS)) parser.add_argument('--model_class', required=False, @@ -84,8 +84,8 @@ def parse_arguments(): parser.add_argument('-b', '--batch_sizes', nargs='+', type=int, default=[1], help="batch size") - parser.add_argument('--past_sequence_lengths', - '-s', + parser.add_argument('-s', + '--past_sequence_lengths', nargs='+', type=int, default=[8, 16, 32, 64, 128, 256], @@ -126,10 +126,8 @@ def main(): model_class = MODEL_CLASSES[args.model_class][0] - config = AutoConfig.from_pretrained(args.model_name, torchscript=args.torchscript, cache_dir=cache_dir) - if hasattr(config, 'return_tuple'): - config.return_tuple = True - model = model_class.from_pretrained(args.model_name, config=config, cache_dir=cache_dir) + config = AutoConfig.from_pretrained(args.model_name_or_path, torchscript=args.torchscript, cache_dir=cache_dir) + model = model_class.from_pretrained(args.model_name_or_path, config=config, cache_dir=cache_dir) # This scirpt does not support float16 for PyTorch. #if args.float16: @@ -139,13 +137,20 @@ def main(): model.to(device) use_external_data_format = (config.n_layer > 24) #TODO: find a way to check model size > 2GB onnx_model_paths = Gpt2Helper.get_onnx_paths(output_dir, - args.model_name, + args.model_name_or_path, args.model_class, has_past=True, new_folder=use_external_data_format) onnx_model_path = onnx_model_paths["raw"] - Gpt2Helper.export_onnx(model, device, onnx_model_path, args.verbose, use_external_data_format) + use_padding = MODEL_CLASSES[args.model_class][2] + Gpt2Helper.export_onnx(model, + device, + onnx_model_path, + args.verbose, + use_external_data_format, + has_position_ids=use_padding, + has_attention_mask=use_padding) if args.optimize_onnx or args.precision != Precision.FLOAT32: onnx_model_path = onnx_model_paths[str(args.precision)] @@ -159,7 +164,7 @@ def main(): logger.info("finished quantizing model") if args.torchscript: - model = Gpt2Helper.torchscript(model, config, device) + model = Gpt2Helper.torchscript(model, config, device, has_position_ids, has_attention_mask) session = create_onnxruntime_session(onnx_model_path, args.use_gpu, @@ -189,10 +194,17 @@ def main(): for batch_size in args.batch_sizes: for past_sequence_length in args.past_sequence_lengths: logger.debug(f"Running test for batch_size={batch_size} past_sequence_length={past_sequence_length}...") - dummy_inputs = Gpt2Helper.get_dummy_inputs(batch_size, past_sequence_length, sequence_length, - config.num_attention_heads, config.hidden_size, - config.n_layer, config.vocab_size, device, - args.precision == Precision.FLOAT16) + dummy_inputs = Gpt2Helper.get_dummy_inputs(batch_size, + past_sequence_length, + sequence_length, + config.num_attention_heads, + config.hidden_size, + config.n_layer, + config.vocab_size, + device, + float16=(args.precision == Precision.FLOAT16), + has_position_ids=use_padding, + has_attention_mask=use_padding) output_shapes = Gpt2Helper.get_output_shapes(batch_size, past_sequence_length, sequence_length, config, args.model_class) @@ -232,7 +244,7 @@ def main(): ) row = { - "model_name": args.model_name, + "model_name": args.model_name_or_path, "model_class": args.model_class, "gpu": args.use_gpu, "precision": args.precision, diff --git a/onnxruntime/python/tools/transformers/convert_to_onnx.py b/onnxruntime/python/tools/transformers/convert_to_onnx.py index f4fdda358f..9e83584e53 100644 --- a/onnxruntime/python/tools/transformers/convert_to_onnx.py +++ b/onnxruntime/python/tools/transformers/convert_to_onnx.py @@ -23,7 +23,7 @@ import torch import numpy import json from transformers import AutoConfig -from gpt2_helper import Gpt2Helper, MODEL_CLASSES, DEFAULT_TOLERANCE +from gpt2_helper import Gpt2Helper, MODEL_CLASSES, DEFAULT_TOLERANCE, PRETRAINED_GPT2_MODELS from gpt2_tester import Gpt2Tester from quantize_helper import QuantizeHelper from benchmark_helper import create_onnxruntime_session, setup_logger, prepare_environment, Precision @@ -38,7 +38,7 @@ def parse_arguments(): '--model_name_or_path', required=True, type=str, - help='Model path, or pretrained model name in the list: ' + ', '.join(['gpt2', 'distilgpt2'])) + help='Model path, or pretrained model name in the list: ' + ', '.join(PRETRAINED_GPT2_MODELS)) parser.add_argument('--model_class', required=False, @@ -123,8 +123,6 @@ def main(): model_class = MODEL_CLASSES[args.model_class][0] config = AutoConfig.from_pretrained(args.model_name_or_path, cache_dir=cache_dir) - if hasattr(config, 'return_tuple'): - config.return_tuple = True model = model_class.from_pretrained(args.model_name_or_path, config=config, cache_dir=cache_dir) device = torch.device("cuda:0" if args.use_gpu else "cpu") @@ -141,7 +139,14 @@ def main(): (args.precision == Precision.FLOAT32 and not args.optimize_onnx)) else onnx_model_paths[str(args.precision)] logger.info(f"Exporting ONNX model to {raw_onnx_model}") - Gpt2Helper.export_onnx(model, device, raw_onnx_model, args.verbose, use_external_data_format) + use_padding = MODEL_CLASSES[args.model_class][2] + Gpt2Helper.export_onnx(model, + device, + raw_onnx_model, + args.verbose, + use_external_data_format, + has_position_ids=use_padding, + has_attention_mask=use_padding) if args.optimize_onnx or args.precision != Precision.FLOAT32: logger.info(f"Optimizing model to {output_path}") @@ -162,34 +167,45 @@ def main(): args.precision == Precision.FLOAT16, rtol=args.tolerance, atol=args.tolerance, - model_class=args.model_class) + model_class=args.model_class, + has_position_ids=use_padding, + has_attention_mask=use_padding) if args.input_test_file: test_inputs = [] + # Each line of test file is a JSON string like: + # {"input_ids": [[14698, 257, 1310, 13688, 319, 326]]} with open(args.input_test_file) as read_f: for i, line in enumerate(read_f): line = line.rstrip() data = json.loads(line) input_ids = torch.from_numpy(numpy.asarray(data["input_ids"], dtype=numpy.int64)).to(device) - if "attention_mask" in data: - numpy_float = numpy.float16 if args.precision == Precision.FLOAT16 else numpy.float32 - attention_mask = torch.from_numpy(numpy.asarray(data["attention_mask"], - dtype=numpy_float)).to(device) - else: - padding = -1 - attention_mask = (input_ids != padding - ).type(torch.float16 if args.precision == Precision.FLOAT16 else torch.float32) - input_ids.masked_fill_(input_ids == padding, 0) + if use_padding: + if "attention_mask" in data: + numpy_float = numpy.float16 if args.precision == Precision.FLOAT16 else numpy.float32 + attention_mask = torch.from_numpy(numpy.asarray(data["attention_mask"], + dtype=numpy_float)).to(device) + else: + padding = -1 + attention_mask = ( + input_ids != + padding).type(torch.float16 if args.precision == Precision.FLOAT16 else torch.float32) + input_ids.masked_fill_(input_ids == padding, 0) - if "position_ids" in data: - position_ids = torch.from_numpy(numpy.asarray(data["position_ids"], dtype=numpy.int64)).to(device) - else: - position_ids = (attention_mask.long().cumsum(-1) - 1) - position_ids.masked_fill_(position_ids < 0, 0) + if "position_ids" in data: + position_ids = torch.from_numpy(numpy.asarray(data["position_ids"], + dtype=numpy.int64)).to(device) + else: + position_ids = (attention_mask.long().cumsum(-1) - 1) + position_ids.masked_fill_(position_ids < 0, 0) + + inputs = {"input_ids": input_ids, "position_ids": position_ids, "attention_mask": attention_mask} + else: + inputs = {"input_ids": input_ids} - inputs = {"input_ids": input_ids, "position_ids": position_ids, "attention_mask": attention_mask} test_inputs.append(inputs) + Gpt2Tester.test_generation(session, model, device, diff --git a/onnxruntime/python/tools/transformers/gpt2_helper.py b/onnxruntime/python/tools/transformers/gpt2_helper.py index 67ba49263c..07bb9f274a 100644 --- a/onnxruntime/python/tools/transformers/gpt2_helper.py +++ b/onnxruntime/python/tools/transformers/gpt2_helper.py @@ -12,6 +12,7 @@ import random import numpy import time import re +from typing import List, Dict, Tuple, Union from transformers import GPT2Model, GPT2LMHeadModel, GPT2Config from benchmark_helper import Precision @@ -52,23 +53,64 @@ class MyGPT2LMHeadModel(GPT2LMHeadModel): return super().forward(input_ids, position_ids=position_ids, attention_mask=attention_mask, past=past) -# Maps model class name to a tuple of model class and name of first output -MODEL_CLASSES = {'GPT2LMHeadModel': (MyGPT2LMHeadModel, 'logits'), 'GPT2Model': (MyGPT2Model, 'last_state')} +class MyGPT2LMHeadModel_NoPadding(GPT2LMHeadModel): + """ Here we wrap a class for Onnx model conversion for GPT2LMHeadModel with past state and no padding. + When you always use batch_size=1 in inference, there is no padding in inputs. In such case, position_ids + and attention_mask need no be in inputs. + """ + def __init__(self, config): + super().__init__(config) + + def forward(self, input_ids, *past): + return super().forward(input_ids, past=past) + + +# Maps model class name to a tuple of model class, name of first output and use padding or not +MODEL_CLASSES = { + 'GPT2LMHeadModel': (MyGPT2LMHeadModel, 'logits', True), + 'GPT2LMHeadModel_NoPadding': (MyGPT2LMHeadModel_NoPadding, 'logits', False), + 'GPT2Model': (MyGPT2Model, 'last_state', True), +} + + +class Gpt2Inputs: + def __init__(self, input_ids, position_ids, attention_mask, past): + self.input_ids: torch.LongTensor = input_ids + self.position_ids: torch.LongTensor = position_ids + self.attention_mask: Union[torch.FloatTensor, torch.HalfTensor] = attention_mask + self.past: Union[List[torch.FloatTensor], List[torch.HalfTensor]] = past + + def to_list(self) -> List: + input_list = [v for v in [self.input_ids, self.position_ids, self.attention_mask] if v is not None] + if self.past: + input_list.extend(self.past) + + return input_list + + def to_tuple(self) -> Tuple: + return tuple(v for v in [self.input_ids, self.position_ids, self.attention_mask, self.past] if v is not None) + + def to_fp32(self): + attention_mask = self.attention_mask.to(dtype=torch.float32) if self.attention_mask is not None else None + past = [p.to(dtype=torch.float32) for p in self.past] + return Gpt2Inputs(self.input_ids, self.position_ids, attention_mask, past) class Gpt2Helper: """ A helper class for Gpt2 model conversion, inference and verification. """ @staticmethod - def get_dummy_inputs(batch_size, - past_sequence_length, - sequence_length, - num_attention_heads, - hidden_size, - num_layer, - vocab_size, - device, - float16=False): + def get_dummy_inputs(batch_size: int, + past_sequence_length: int, + sequence_length: int, + num_attention_heads: int, + hidden_size: int, + num_layer: int, + vocab_size: int, + device: torch.device, + float16: bool = False, + has_position_ids: bool = True, + has_attention_mask: bool = True) -> Gpt2Inputs: """ Create random inputs for GPT2 model. Returns torch tensors of input_ids, position_ids, attention_mask and a list of past state tensors. """ @@ -82,21 +124,29 @@ class Gpt2Helper: dtype=torch.int64, device=device) - total_sequence_length = past_sequence_length + sequence_length - attention_mask = torch.ones([batch_size, total_sequence_length], dtype=float_type, device=device) - if total_sequence_length >= 2: - padding_position = random.randint(0, total_sequence_length - 1) # test input with padding. - attention_mask[:, padding_position] = 0 + attention_mask = None + if has_attention_mask: + total_sequence_length = past_sequence_length + sequence_length + attention_mask = torch.ones([batch_size, total_sequence_length], dtype=float_type, device=device) + if total_sequence_length >= 2: + padding_position = random.randint(0, total_sequence_length - 1) # test input with padding. + attention_mask[:, padding_position] = 0 # Deduce position_ids from attention mask - position_ids = (attention_mask.long().cumsum(-1) - 1) - position_ids.masked_fill_(position_ids < 0, 0) - position_ids = position_ids[:, past_sequence_length:] + position_ids = None + if has_position_ids: + position_ids = (attention_mask.long().cumsum(-1) - 1) + position_ids.masked_fill_(position_ids < 0, 0) + position_ids = position_ids[:, past_sequence_length:] - return input_ids, position_ids, attention_mask, past + return Gpt2Inputs(input_ids, position_ids, attention_mask, past) @staticmethod - def get_output_shapes(batch_size, past_sequence_length, sequence_length, config, model_class="GPT2LMHeadModel"): + def get_output_shapes(batch_size: int, + past_sequence_length: int, + sequence_length: int, + config: GPT2Config, + model_class: str = "GPT2LMHeadModel") -> Dict[str, List[int]]: """ Returns a dictionary with output name as key, and shape as value. """ num_attention_heads = config.num_attention_heads @@ -106,9 +156,7 @@ class Gpt2Helper: output_name = MODEL_CLASSES[model_class][1] - last_state_shape = [ - batch_size, sequence_length, vocab_size if model_class == "GPT2LMHeadModel" else hidden_size - ] + last_state_shape = [batch_size, sequence_length, vocab_size if output_name == "logits" else hidden_size] present_state_shape = [ 2, batch_size, num_attention_heads, past_sequence_length + sequence_length, int(hidden_size / num_attention_heads) @@ -176,8 +224,14 @@ class Gpt2Helper: return is_all_close @staticmethod - def export_onnx(model, device, onnx_model_path, verbose=False, use_external_data_format=False): - """ Export GPT-2 model with past state to ONNX model + def export_onnx(model, + device, + onnx_model_path: str, + verbose: bool = False, + use_external_data_format: bool = False, + has_position_ids: bool = True, + has_attention_mask: bool = True): + """ Export GPT-2 model with past state to ONNX model. """ config: GPT2Config = model.config num_layer = config.n_layer @@ -189,11 +243,11 @@ class Gpt2Helper: num_layer=num_layer, vocab_size=config.vocab_size, device=device, - float16=False) + float16=False, + has_position_ids=has_position_ids, + has_attention_mask=has_attention_mask) + input_list = dummy_inputs.to_list() - dummy_input_ids, dummy_position_ids, dummy_attention_mask, dummy_past = dummy_inputs - - input_list = [dummy_input_ids, dummy_position_ids, dummy_attention_mask] + dummy_past with torch.no_grad(): outputs = model(*input_list) @@ -218,17 +272,23 @@ class Gpt2Helper: for name in present_names: dynamic_axes[name] = {1: 'batch_size', 3: 'total_seq_len'} - dynamic_axes['attention_mask'] = {0: 'batch_size', 1: 'total_seq_len'} - dynamic_axes['position_ids'] = {0: 'batch_size', 1: 'seq_len'} + input_names = ['input_ids'] + if has_position_ids: + dynamic_axes['position_ids'] = {0: 'batch_size', 1: 'seq_len'} + input_names.append('position_ids') + if has_attention_mask: + dynamic_axes['attention_mask'] = {0: 'batch_size', 1: 'total_seq_len'} + input_names.append('attention_mask') + input_names.extend(past_names) logger.info( - f"Shapes: input_ids={dummy_input_ids.shape} past={dummy_past[0].shape} output={outputs[0].shape} present={outputs[1][0].shape}" + f"Shapes: input_ids={dummy_inputs.input_ids.shape} past={dummy_inputs.past[0].shape} output={outputs[0].shape} present={outputs[1][0].shape}" ) torch.onnx.export(model, args=tuple(input_list), f=onnx_model_path, - input_names=['input_ids', 'position_ids', 'attention_mask'] + past_names, + input_names=input_names, output_names=output_names, example_outputs=outputs, dynamic_axes=dynamic_axes, @@ -259,18 +319,13 @@ class Gpt2Helper: m.save_model_to_file(optimized_model_path, use_external_data_format) @staticmethod - def pytorch_inference(model, inputs, total_runs=0): + def pytorch_inference(model, inputs: Gpt2Inputs, total_runs: int = 0): """ Run inference of PyTorch model, and returns average latency in ms when total_runs > 0 besides outputs. """ logger.debug(f"start pytorch_inference") - input_ids, position_ids, attention_mask, past = inputs - # Convert it back to fp32 as the PyTroch model cannot deal with half input. - if attention_mask is not None: - attention_mask = attention_mask.to(dtype=torch.float32) - past = [p.to(dtype=torch.float32) for p in past] - - input_list = [input_ids, position_ids, attention_mask] + past + # Convert it to fp32 as the PyTroch model cannot deal with half input. + input_list = inputs.to_fp32().to_list() with torch.no_grad(): outputs = model(*input_list) @@ -291,23 +346,22 @@ class Gpt2Helper: return outputs, average_latency @staticmethod - def onnxruntime_inference(ort_session, inputs, total_runs=0): + def onnxruntime_inference(ort_session, inputs: Gpt2Inputs, total_runs: int = 0): """ Run inference of ONNX model, and returns average latency in ms when total_runs > 0 besides outputs. """ logger.debug(f"start onnxruntime_inference") - input_ids, position_ids, attention_mask, past = inputs - ort_inputs = {'input_ids': numpy.ascontiguousarray(input_ids.cpu().numpy())} + ort_inputs = {'input_ids': numpy.ascontiguousarray(inputs.input_ids.cpu().numpy())} - if past is not None: - for i, past_i in enumerate(past): + if inputs.past is not None: + for i, past_i in enumerate(inputs.past): ort_inputs[f'past_{i}'] = numpy.ascontiguousarray(past_i.cpu().numpy()) - if attention_mask is not None: - ort_inputs['attention_mask'] = numpy.ascontiguousarray(attention_mask.cpu().numpy()) + if inputs.attention_mask is not None: + ort_inputs['attention_mask'] = numpy.ascontiguousarray(inputs.attention_mask.cpu().numpy()) - if position_ids is not None: - ort_inputs['position_ids'] = numpy.ascontiguousarray(position_ids.cpu().numpy()) + if inputs.position_ids is not None: + ort_inputs['position_ids'] = numpy.ascontiguousarray(inputs.position_ids.cpu().numpy()) ort_outputs = ort_session.run(None, ort_inputs) if total_runs == 0: @@ -384,20 +438,19 @@ class Gpt2Helper: @staticmethod def onnxruntime_inference_with_binded_io(ort_session, - inputs, - output_buffers, - output_shapes, - total_runs=0, - return_numpy=True, - include_copy_output_latency=False): + inputs: Gpt2Inputs, + output_buffers: Dict[str, torch.Tensor], + output_shapes : Dict[str, List[int]], + total_runs: int = 0, + return_numpy: bool = True, + include_copy_output_latency: bool = False): """ Inference with IO binding. Returns outputs, and optional latency when total_runs > 0. """ logger.debug(f"start onnxruntime_inference_with_binded_io") - input_ids, position_ids, attention_mask, past = inputs # Bind inputs and outputs to onnxruntime session - io_binding = Gpt2Helper.prepare_io_binding(ort_session, input_ids, position_ids, attention_mask, past, - output_buffers, output_shapes) + io_binding = Gpt2Helper.prepare_io_binding(ort_session, inputs.input_ids, inputs.position_ids, + inputs.attention_mask, inputs.past, output_buffers, output_shapes) # Run onnxruntime with io binding ort_session.run_with_iobinding(io_binding) @@ -433,7 +486,9 @@ class Gpt2Helper: atol=5e-4, total_test_cases=100, use_io_binding=True, - model_class="GPT2LMHeadModel"): + model_class="GPT2LMHeadModel", + has_position_ids=True, + has_attention_mask=True): """ Generate random inputs and compare the results of PyTorch and Onnx Runtime. """ @@ -463,7 +518,9 @@ class Gpt2Helper: f"Running parity test for batch_size={batch_size} past_sequence_length={past_sequence_length}...") dummy_inputs = Gpt2Helper.get_dummy_inputs(batch_size, past_sequence_length, sequence_length, config.num_attention_heads, config.hidden_size, config.n_layer, - config.vocab_size, device, is_float16) + config.vocab_size, device, is_float16, has_position_ids, + has_attention_mask) + outputs = Gpt2Helper.pytorch_inference(model, dummy_inputs) if use_io_binding: ort_outputs = Gpt2Helper.onnxruntime_inference(ort_session, dummy_inputs) @@ -482,20 +539,21 @@ class Gpt2Helper: return passed_test_cases == total_test_cases @staticmethod - def torchscript(model, config, device): + def torchscript(model, config, device, has_position_ids=True, has_attention_mask=True): """ JIT trace for TorchScript. """ - dummy_inputs = Gpt2Helper.get_dummy_inputs(batch_size=1, - past_sequence_length=1, - sequence_length=1, - num_attention_heads=config.num_attention_heads, - hidden_size=config.hidden_size, - num_layer=config.n_layer, - vocab_size=config.vocab_size, - device=device, - float16=False) - dummy_input_ids, dummy_position_ids, dummy_attention_mask, dummy_past = dummy_inputs - return torch.jit.trace(model, [dummy_input_ids, dummy_position_ids, dummy_attention_mask] + dummy_past) + input_list = Gpt2Helper.get_dummy_inputs(batch_size=1, + past_sequence_length=1, + sequence_length=1, + num_attention_heads=config.num_attention_heads, + hidden_size=config.hidden_size, + num_layer=config.n_layer, + vocab_size=config.vocab_size, + device=device, + float16=False, + has_position_ids=has_position_ids, + has_attention_mask=has_attention_mask).to_list() + return torch.jit.trace(model, input_list) @staticmethod def get_onnx_paths(output_dir, diff --git a/onnxruntime/python/tools/transformers/gpt2_tester.py b/onnxruntime/python/tools/transformers/gpt2_tester.py index bd1cbc30ff..e8c540c7a3 100644 --- a/onnxruntime/python/tools/transformers/gpt2_tester.py +++ b/onnxruntime/python/tools/transformers/gpt2_tester.py @@ -14,7 +14,7 @@ import time import timeit import math import statistics -from gpt2_helper import Gpt2Helper +from gpt2_helper import Gpt2Helper, Gpt2Inputs from benchmark_helper import Precision logger = logging.getLogger(__name__) @@ -133,6 +133,9 @@ class Gpt2Tester: self.position_ids = position_ids self.attention_mask = attention_mask + self.has_position_ids = position_ids is not None + self.has_attention_mask = attention_mask is not None + # Emtpy past state for first inference self.past = [] past_shape = [2, self.batch_size, num_attention_heads, 0, hidden_size // num_attention_heads] @@ -146,8 +149,8 @@ class Gpt2Tester: self.top_k = top_k self.top_k_required_order = top_k_required_order - def get_input_tuple(self): - return self.input_ids, self.position_ids, self.attention_mask, self.past + def get_inputs(self) -> Gpt2Inputs: + return Gpt2Inputs(self.input_ids, self.position_ids, self.attention_mask, self.past) def update(self, output, step, device): """ @@ -160,10 +163,16 @@ class Gpt2Tester: self.top_k_tokens = Gpt2Tester.predict_next_token(self.logits, self.top_k, self.top_k_required_order) self.input_ids = self.top_1_tokens.clone().detach().reshape([self.batch_size, 1]).to(device) - self.position_ids = torch.tensor([self.input_length + step - 1]).unsqueeze(0).repeat(self.batch_size, - 1).to(device) - self.attention_mask = torch.cat( - [self.attention_mask, torch.ones([self.batch_size, 1]).type_as(self.attention_mask)], 1).to(device) + + if self.has_position_ids: + self.position_ids = torch.tensor([self.input_length + step - 1]).unsqueeze(0).repeat(self.batch_size, + 1).to(device) + + if self.has_attention_mask: + self.attention_mask = torch.cat( + [self.attention_mask, + torch.ones([self.batch_size, 1]).type_as(self.attention_mask)], 1).to(device) + self.past = [] if isinstance(output[1], tuple): # past in torch output is tuple @@ -188,11 +197,13 @@ class Gpt2Tester: if not torch.all(self.input_ids == baseline.input_ids): print('Input_ids is different', self.input_ids, baseline.input_ids) - if not torch.all(self.position_ids == baseline.position_ids): - print('position_ids is different', self.position_ids, baseline.position_ids) + if self.has_position_ids: + if not torch.all(self.position_ids == baseline.position_ids): + print('position_ids is different', self.position_ids, baseline.position_ids) - if not torch.all(self.attention_mask == baseline.attention_mask): - print('attention_mask is different', self.attention_mask, baseline.attention_mask) + if self.has_attention_mask: + if not torch.all(self.attention_mask == baseline.attention_mask): + print('attention_mask is different', self.attention_mask, baseline.attention_mask) assert len(self.past) == len(baseline.past) @@ -301,8 +312,8 @@ class Gpt2Tester: if i % 10 == 0: print(f"{i}") input_ids = inputs["input_ids"] - position_ids = inputs["position_ids"] - attention_mask = inputs["attention_mask"] + position_ids = inputs["position_ids"] if "position_ids" in inputs else None + attention_mask = inputs["attention_mask"] if "attention_mask" in inputs else None onnx_runner = Gpt2Tester(input_ids, position_ids, attention_mask, n_head, n_embd, n_layer, device, is_float16, top_k, not top_k_no_order) @@ -322,12 +333,13 @@ class Gpt2Tester: past_seq_len = list(onnx_runner.past[0].size())[3] start_time = timeit.default_timer() - pytorch_output = Gpt2Helper.pytorch_inference(model, torch_runner.get_input_tuple()) + pytorch_output = Gpt2Helper.pytorch_inference(model, torch_runner.get_inputs()) torch_metric.add_latency(past_seq_len, timeit.default_timer() - start_time) torch_runner.update(pytorch_output, step, device) - input_tuple = onnx_runner.get_input_tuple() - onnx_output, avg_latency_ms = Gpt2Helper.onnxruntime_inference(session, input_tuple, total_runs=1) + onnx_output, avg_latency_ms = Gpt2Helper.onnxruntime_inference(session, + onnx_runner.get_inputs(), + total_runs=1) onnx_metric.add_latency(past_seq_len, avg_latency_ms / 1000.0) onnx_runner.update(onnx_output, step, device) @@ -338,10 +350,9 @@ class Gpt2Tester: model_class=model_class) Gpt2Helper.auto_increase_buffer_size(output_buffers, output_shapes) - input_tuple = onnx_io_runner.get_input_tuple() onnx_io_output, avg_latency_ms = Gpt2Helper.onnxruntime_inference_with_binded_io( session, - input_tuple, + onnx_io_runner.get_inputs(), output_buffers, output_shapes, total_runs=1,