From a5b2291e0fe7c7d42f30154ccb20d6cde1380c3c Mon Sep 17 00:00:00 2001 From: trajep Date: Tue, 5 Dec 2023 04:26:50 +0800 Subject: [PATCH] [Transformer Optimization]Return model directly for unknown model type (#18642) This pull request is used to improves the handling of unsupported model types in the optimization process. --- onnxruntime/python/tools/transformers/optimizer.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/onnxruntime/python/tools/transformers/optimizer.py b/onnxruntime/python/tools/transformers/optimizer.py index 6842a97fe0..ba61f4f6e4 100644 --- a/onnxruntime/python/tools/transformers/optimizer.py +++ b/onnxruntime/python/tools/transformers/optimizer.py @@ -209,6 +209,10 @@ def optimize_by_fusion( if model_type not in ["bert", "swin", "unet", "vae", "clip"] and (num_heads == 0 or hidden_size == 0): logger.warning(f"Please specify parameters of num_heads and hidden_size for model_type {model_type}") + if model_type not in MODEL_TYPES: + logger.warning(f"Unsupported model type: {model_type} for graph fusion, directly return model.") + return OnnxModel(model) + (optimizer_class, producer, _) = MODEL_TYPES[model_type] if model.producer_name and producer != model.producer_name: @@ -290,6 +294,10 @@ def optimize_model( """ assert opt_level is None or opt_level in [0, 1, 2, 99] + if model_type not in MODEL_TYPES: + logger.warning(f"Unsupported model type: {model_type} for optimization, directly return model.") + return OnnxModel(load_model(input)) + (optimizer_class, _producer, default_opt_level) = MODEL_TYPES[model_type] if opt_level is None: