This commit is contained in:
Ye Wang 2021-02-26 13:58:41 -08:00 committed by GitHub
parent d5175795d2
commit b4b87ac7a0
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 11 additions and 14 deletions

View file

@ -9,16 +9,6 @@ MODEL_CLASSES = [
'AutoModel', 'AutoModelWithLMHead', 'AutoModelForSequenceClassification', 'AutoModelForQuestionAnswering'
]
# List of models that require external data saving for onnx export but do not require it when saving optimized onnx model
# Very few models in the huggingface list require it for both: albert-xxlarge-v1, albert-xxlarge-v2
# TODO: most of the models in the below exempt list having runtime issues when saving these optimized onnx models
# using external data format. Need to address the issue in the future
EXEMPT_MODELS = [
"gpt2-large", "gpt2-xl", "xlm-mlm-en-2048", "xlm-mlm-17-1280", "xlm-mlm-100-1280", "ctrl", "albert-xlarge-v1",
"albert-xlarge-v2", "t5-large", "t5-3b", "t5-11b", "xlm-roberta-large", "microsoft/DialoGPT-large",
"facebook/mbart-large-en-ro"
]
# List of pretrained models: https://huggingface.co/transformers/pretrained_models.html
# Pretrained model name to a tuple of input names, opset_version, use_external_data_format, optimization model type
MODELS = {
@ -105,4 +95,14 @@ MODELS = {
# Longformer (use benchmark_longformer.py instead)
#"allenai/longformer-base-4096": (["input_ids"], 12, False, "bert"),
#"allenai/longformer-large-4096": (["input_ids"], 12, False, "bert"),
# MBart
"facebook/mbart-large-cc25": (["input_ids"], 11, True, "bert"),
"facebook/mbart-large-en-ro": (["input_ids"], 11, True, "bert"),
# Layoutlm
"microsoft/layoutlm-base-uncased": (["input_ids"], 11, False, "bert"),
"microsoft/layoutlm-large-uncased": (["input_ids"], 11, False, "bert"),
# Squeezebert
"squeezebert/squeezebert-uncased": (["input_ids"], 11, False, "bert"),
"squeezebert/squeezebert-mnli": (["input_ids"], 11, False, "bert"),
"squeezebert/squeezebert-mnli-headless": (["input_ids"], 11, False, "bert"),
}

View file

@ -13,7 +13,7 @@ from transformers import AutoConfig, AutoTokenizer, AutoModel
from benchmark_helper import create_onnxruntime_session, Precision
from gpt2_helper import GPT2ModelNoPastState, PRETRAINED_GPT2_MODELS
from quantize_helper import QuantizeHelper
from huggingface_models import MODEL_CLASSES, EXEMPT_MODELS
from huggingface_models import MODEL_CLASSES
logger = logging.getLogger(__name__)
@ -203,9 +203,6 @@ def optimize_onnx_model(model_name, onnx_model_path, optimized_model_path, model
if Precision.FLOAT16 == precision:
opt_model.convert_model_float32_to_float16()
if model_name in EXEMPT_MODELS:
use_external_data_format = False
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}")