diff --git a/onnxruntime/python/tools/transformers/benchmark_gpt2.py b/onnxruntime/python/tools/transformers/benchmark_gpt2.py index 4ffcd49dc1..ff5fc523c7 100644 --- a/onnxruntime/python/tools/transformers/benchmark_gpt2.py +++ b/onnxruntime/python/tools/transformers/benchmark_gpt2.py @@ -150,18 +150,23 @@ def main(args): has_attention_mask=use_padding) if args.optimize_onnx or args.precision != Precision.FLOAT32: - onnx_model_path = onnx_model_paths[str(args.precision)] + onnx_model_path = onnx_model_paths[str(args.precision) if args.precision != Precision.INT8 else 'fp32'] Gpt2Helper.optimize_onnx(onnx_model_paths["raw"], onnx_model_path, args.precision == Precision.FLOAT16, model.config.num_attention_heads, model.config.hidden_size, use_external_data_format) if args.precision == Precision.INT8: logger.info("quantizing model...") - QuantizeHelper.quantize_onnx_model(onnx_model_path, onnx_model_path, use_external_data_format) + QuantizeHelper.quantize_onnx_model(onnx_model_path, onnx_model_paths["int8"], use_external_data_format) model = QuantizeHelper.quantize_torch_model(model) logger.info("finished quantizing model") + onnx_model_path = onnx_model_paths["int8"] if args.torchscript: - model = Gpt2Helper.torchscript(model, config, device, has_position_ids, has_attention_mask) + model = Gpt2Helper.torchscript(model, + config, + device, + has_position_ids=use_padding, + has_attention_mask=use_padding) session = create_onnxruntime_session(onnx_model_path, args.use_gpu, diff --git a/onnxruntime/python/tools/transformers/convert_to_onnx.py b/onnxruntime/python/tools/transformers/convert_to_onnx.py index ebfddcb86d..7b2267905d 100644 --- a/onnxruntime/python/tools/transformers/convert_to_onnx.py +++ b/onnxruntime/python/tools/transformers/convert_to_onnx.py @@ -95,6 +95,9 @@ def parse_arguments(): 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) + args = parser.parse_args() return args @@ -122,6 +125,9 @@ def main(): if args.precision == Precision.INT8: assert not args.use_gpu, "quantization only supports CPU" + if args.use_external_data_format: + assert not args.output.endswith('.onnx'), "output shall be a directory for --use_external_data_format" + model_class = MODEL_CLASSES[args.model_class][0] config = AutoConfig.from_pretrained(args.model_name_or_path, cache_dir=cache_dir) model = model_class.from_pretrained(args.model_name_or_path, config=config, cache_dir=cache_dir) @@ -129,15 +135,15 @@ def main(): device = torch.device("cuda:0" if args.use_gpu else "cpu") model.eval().to(device) - use_external_data_format = (config.n_layer > 24) #TODO: find a way to check model size > 2GB + if (not args.use_external_data_format) and (config.n_layer > 24): + logger.info(f"Try --use_external_data_format when model size > 2GB") + onnx_model_paths = Gpt2Helper.get_onnx_paths(output_dir, args.model_name_or_path, args.model_class, - new_folder=use_external_data_format) - raw_onnx_model = args.output if args.output.endswith('.onnx') else onnx_model_paths["raw"] - output_path = raw_onnx_model if ( - args.output.endswith('.onnx') or - (args.precision == Precision.FLOAT32 and not args.optimize_onnx)) else onnx_model_paths[str(args.precision)] + new_folder=args.use_external_data_format) + + raw_onnx_model = onnx_model_paths["raw"] logger.info(f"Exporting ONNX model to {raw_onnx_model}") use_padding = MODEL_CLASSES[args.model_class][2] @@ -145,20 +151,33 @@ def main(): device, raw_onnx_model, args.verbose, - use_external_data_format, + args.use_external_data_format, has_position_ids=use_padding, has_attention_mask=use_padding) if args.optimize_onnx or args.precision != Precision.FLOAT32: + output_path = onnx_model_paths[str(args.precision) if args.precision != Precision.INT8 else 'fp32'] + logger.info(f"Optimizing model to {output_path}") Gpt2Helper.optimize_onnx(raw_onnx_model, output_path, args.precision == Precision.FLOAT16, - model.config.num_attention_heads, model.config.hidden_size) + model.config.num_attention_heads, model.config.hidden_size, + args.use_external_data_format) + else: + output_path = raw_onnx_model if args.precision == Precision.INT8: logger.info("quantizing model...") - QuantizeHelper.quantize_onnx_model(output_path, output_path) + QuantizeHelper.quantize_onnx_model(output_path, onnx_model_paths['int8'], args.use_external_data_format) model = QuantizeHelper.quantize_torch_model(model) logger.info("finished quantizing model") + output_path = onnx_model_paths['int8'] + + if args.output.endswith('.onnx') and output_path != args.output and not args.use_external_data_format: + import shutil + shutil.move(output_path, args.output) + output_path = args.output + + logger.info(f"Output path: {output_path}") session = create_onnxruntime_session(output_path, args.use_gpu, enable_all_optimization=True, verbose=args.verbose) if session is not None: @@ -177,7 +196,7 @@ def main(): # 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): + for _, 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) diff --git a/onnxruntime/python/tools/transformers/gpt2_helper.py b/onnxruntime/python/tools/transformers/gpt2_helper.py index 77afbb5959..9fbd58a386 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 pathlib import Path from typing import List, Dict, Tuple, Union from transformers import GPT2Model, GPT2LMHeadModel, GPT2Config from benchmark_helper import Precision @@ -285,6 +286,8 @@ class Gpt2Helper: f"Shapes: input_ids={dummy_inputs.input_ids.shape} past={dummy_inputs.past[0].shape} output={outputs[0].shape} present={outputs[1][0].shape}" ) + Path(onnx_model_path).parent.mkdir(parents=True, exist_ok=True) + torch.onnx.export(model, args=tuple(input_list), f=onnx_model_path, @@ -567,7 +570,6 @@ class Gpt2Helper: model_name = model_name_or_path if not re.match('^[\w_-]+$', model_name_or_path): # It is not a name, shall be a path assert os.path.isdir(model_name_or_path) - from pathlib import Path model_name = Path(model_name_or_path).parts[-1] if model_class != 'GPT2LMHeadModel': @@ -577,9 +579,13 @@ class Gpt2Helper: model_name += "_past" if new_folder: - output_dir = os.path.join(output_dir, model_name) - if not os.path.exists(output_dir): - os.makedirs(output_dir) + # store each model to its own directory (for external data format). + return { + "raw": os.path.join(os.path.join(output_dir, model_name), model_name + ".onnx"), + "fp32": os.path.join(os.path.join(output_dir, model_name + "_fp32"), model_name + "_fp32.onnx"), + "fp16": os.path.join(os.path.join(output_dir, model_name + "_fp16"), model_name + "_fp16.onnx"), + "int8": os.path.join(os.path.join(output_dir, model_name + "_int8"), model_name + "_int8.onnx") + } return { "raw": os.path.join(output_dir, model_name + ".onnx"), diff --git a/onnxruntime/python/tools/transformers/onnx_exporter.py b/onnxruntime/python/tools/transformers/onnx_exporter.py index 7e6acc6432..145ff4b0fd 100644 --- a/onnxruntime/python/tools/transformers/onnx_exporter.py +++ b/onnxruntime/python/tools/transformers/onnx_exporter.py @@ -134,7 +134,6 @@ def get_onnx_file_path(onnx_dir: str, model_name: str, input_count: int, optimiz filename += f"_ort" directory = onnx_dir - # ONNXRuntime will not write external data so the raw and optimized models shall be in same directory. if use_external_data and not optimized_by_onnxruntime: directory = os.path.join(onnx_dir, filename) @@ -158,6 +157,7 @@ def add_filename_suffix(file_path: str, suffix: str) -> str: def optimize_onnx_model_by_ort(onnx_model_path, ort_model_path, use_gpu, overwrite, model_fusion_statistics): if overwrite or not os.path.exists(ort_model_path): + Path(ort_model_path).parent.mkdir(parents=True, exist_ok=True) from optimizer import optimize_by_onnxruntime, get_fusion_statistics # Use onnxruntime to optimize model, which will be saved to *_ort.onnx opt_model = optimize_by_onnxruntime(onnx_model_path, @@ -170,8 +170,11 @@ def optimize_onnx_model_by_ort(onnx_model_path, ort_model_path, use_gpu, overwri def optimize_onnx_model(onnx_model_path, optimized_model_path, model_type, num_attention_heads, hidden_size, use_gpu, - precision, use_raw_attention_mask, overwrite, model_fusion_statistics): + precision, use_raw_attention_mask, overwrite, model_fusion_statistics, + use_external_data_format): if overwrite or not os.path.exists(optimized_model_path): + Path(optimized_model_path).parent.mkdir(parents=True, exist_ok=True) + from optimizer import optimize_model from onnx_model_bert import BertOptimizationOptions optimization_options = BertOptimizationOptions(model_type) @@ -199,7 +202,7 @@ def optimize_onnx_model(onnx_model_path, optimized_model_path, model_type, num_a if Precision.FLOAT16 == precision: opt_model.convert_model_float32_to_float16() - opt_model.save_model_to_file(optimized_model_path) + opt_model.save_model_to_file(optimized_model_path, use_external_data_format) else: logger.info(f"Skip optimization since model existed: {optimized_model_path}") @@ -271,8 +274,6 @@ def load_pt_model_from_tf(model_name): config, model = tf2pt_pipeline(model_name) return config, model - - def validate_and_optimize_onnx(model_name, use_external_data_format, model_type, onnx_dir, input_names, use_gpu, precision, optimize_onnx, validate_onnx, use_raw_attention_mask, overwrite, config, model_fusion_statistics, onnx_model_path, example_inputs, example_outputs_flatten): @@ -286,7 +287,7 @@ def validate_and_optimize_onnx(model_name, use_external_data_format, model_type, False, use_external_data_format) optimize_onnx_model(onnx_model_path, optimized_model_path, model_type, config.num_attention_heads, config.hidden_size, use_gpu, precision, use_raw_attention_mask, overwrite, - model_fusion_statistics) + model_fusion_statistics, use_external_data_format) onnx_model_path = optimized_model_path if validate_onnx: @@ -295,7 +296,7 @@ def validate_and_optimize_onnx(model_name, use_external_data_format, model_type, if precision == Precision.INT8: logger.info(f"Quantizing model: {onnx_model_path}") - QuantizeHelper.quantize_onnx_model(onnx_model_path, onnx_model_path) + QuantizeHelper.quantize_onnx_model(onnx_model_path, onnx_model_path, use_external_data_format) logger.info(f"Finished quantizing model: {onnx_model_path}") else: # Use OnnxRuntime to optimize @@ -335,6 +336,7 @@ def export_onnx_model_from_pt(model_name, opset_version, use_external_data_forma if overwrite or not os.path.exists(onnx_model_path): logger.info("Exporting ONNX model to {}".format(onnx_model_path)) + Path(onnx_model_path).parent.mkdir(parents=True, exist_ok=True) dynamic_axes, output_names = build_dynamic_axes(example_inputs, example_outputs_flatten) @@ -392,6 +394,8 @@ def export_onnx_model_from_tf(model_name, opset_version, use_external_data_forma if overwrite or not os.path.exists(onnx_model_path): logger.info("Exporting ONNX model to {}".format(onnx_model_path)) + Path(onnx_model_path).parent.mkdir(parents=True, exist_ok=True) + import keras2onnx onnx_model = keras2onnx.convert_keras(model, model.name, target_opset=opset_version) keras2onnx.save_model(onnx_model, onnx_model_path) @@ -406,3 +410,4 @@ def export_onnx_model_from_tf(model_name, opset_version, use_external_data_forma example_inputs, example_outputs_flatten) return onnx_model_file, is_valid_onnx_model, vocab_size, max_input_size + diff --git a/onnxruntime/python/tools/transformers/onnx_model.py b/onnxruntime/python/tools/transformers/onnx_model.py index f7cb09e059..694a105942 100644 --- a/onnxruntime/python/tools/transformers/onnx_model.py +++ b/onnxruntime/python/tools/transformers/onnx_model.py @@ -5,8 +5,10 @@ from typing import List, Tuple import logging +import os import sys import argparse +from pathlib import Path import numpy as np from collections import deque from onnx import ModelProto, TensorProto, numpy_helper, helper, external_data_helper, save_model @@ -660,6 +662,8 @@ class OnnxModel: def save_model_to_file(self, output_path, use_external_data_format=False): logger.info(f"Output model to {output_path}") + Path(output_path).parent.mkdir(parents=True, exist_ok=True) + if output_path.endswith(".json"): # Output text for testing small model. assert isinstance(self.model, ModelProto) with open(output_path, "w") as out: @@ -667,10 +671,12 @@ class OnnxModel: else: # Save model to external data, which is needed for model size > 2GB if use_external_data_format: - from pathlib import Path + data_file = str(Path(output_path).name + ".data") + if os.path.isfile(data_file): + os.remove(data_file) external_data_helper.convert_model_to_external_data(self.model, all_tensors_to_one_file=True, - location=Path(output_path).name + ".data") + location=data_file) save_model(self.model, output_path) def get_graph_inputs_excluding_initializers(self):