saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-28 20:11:22 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

Add --use_external_data_format in convert_to_onnx.py (#5393)

Browse source
This commit is contained in:
Tianlei Wu 2020-10-07 09:42:02 -07:00 committed by GitHub
parent 5947445457
commit 094384781e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
5 changed files with 67 additions and 26 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

11
onnxruntime/python/tools/transformers/benchmark_gpt2.py
View file

@ -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,

39
onnxruntime/python/tools/transformers/convert_to_onnx.py
View file

@ -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)

14
onnxruntime/python/tools/transformers/gpt2_helper.py
View file

@ -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"),

19
onnxruntime/python/tools/transformers/onnx_exporter.py
View file

@ -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

10
onnxruntime/python/tools/transformers/onnx_model.py
View file

@ -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):