diff --git a/onnxruntime/python/tools/transformers/onnx_exporter.py b/onnxruntime/python/tools/transformers/onnx_exporter.py index 61e2f93124..1e22f5f4ac 100644 --- a/onnxruntime/python/tools/transformers/onnx_exporter.py +++ b/onnxruntime/python/tools/transformers/onnx_exporter.py @@ -127,15 +127,17 @@ 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, - fp16, use_raw_attention_mask, overwrite, model_fusion_statistics): + precision, use_raw_attention_mask, overwrite, model_fusion_statistics): if overwrite or not os.path.exists(optimized_model_path): from optimizer import optimize_model from onnx_model_bert import BertOptimizationOptions optimization_options = BertOptimizationOptions(model_type) if use_raw_attention_mask: optimization_options.use_raw_attention_mask() - if fp16: + if Precision.FLOAT16 == precision: optimization_options.enable_gelu_approximation = True + if Precision.INT8 == precision: + optimization_options.enable_embed_layer_norm = False # Use script to optimize model. # Use opt_level <= 1 for models to be converted to fp16, because some fused op (like FusedGemm) has only fp32 and no fp16. @@ -150,7 +152,7 @@ def optimize_onnx_model(onnx_model_path, optimized_model_path, model_type, num_a only_onnxruntime=False) model_fusion_statistics[optimized_model_path] = opt_model.get_fused_operator_statistics() - if fp16: + if Precision.FLOAT16 == precision: opt_model.convert_model_float32_to_float16() opt_model.save_model_to_file(optimized_model_path) else: @@ -215,7 +217,7 @@ def export_onnx_model(model_name, opset_version, use_external_data_format, model optimized_model_path = get_onnx_file_path(onnx_dir, model_name, len(input_names), True, use_gpu, precision, 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 == Precision.FLOAT16, use_raw_attention_mask, + config.hidden_size, use_gpu, precision, use_raw_attention_mask, overwrite, model_fusion_statistics) onnx_model_path = optimized_model_path diff --git a/onnxruntime/python/tools/transformers/quantize_helper.py b/onnxruntime/python/tools/transformers/quantize_helper.py index 4d073f427e..3ba683ee1e 100644 --- a/onnxruntime/python/tools/transformers/quantize_helper.py +++ b/onnxruntime/python/tools/transformers/quantize_helper.py @@ -7,6 +7,7 @@ import logging import torch import onnx +import os from transformers.modeling_utils import Conv1D logger = logging.getLogger(__name__) @@ -34,6 +35,12 @@ def conv1d_to_linear(model): conv1d_to_linear(module) +def _get_size_of_pytorch_model(model): + torch.save(model.state_dict(), "temp.p") + size = os.path.getsize("temp.p")/(1024*1024) + os.remove('temp.p') + return size + class QuantizeHelper: @staticmethod def quantize_torch_model(model, dtype=torch.qint8): @@ -43,11 +50,15 @@ class QuantizeHelper: TODO: mix of in-place and return, but results are different ''' conv1d_to_linear(model) - return torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=dtype) + quantized_model = torch.quantization.quantize_dynamic(model, {torch.nn.Linear}, dtype=dtype) + logger.info(f'Size of full precision Torch model(MB):{_get_size_of_pytorch_model(model)}') + logger.info(f'Size of quantized Torch model(MB):{_get_size_of_pytorch_model(quantized_model)}') + return quantized_model @staticmethod def quantize_onnx_model(onnx_model_path, quantized_model_path): from onnxruntime.quantization import quantize, QuantizationMode + logger.info(f'Size of full precision ONNX model(MB):{os.path.getsize(onnx_model_path)/(1024*1024)}') onnx_opt_model = onnx.load(onnx_model_path) quantized_onnx_model = quantize(onnx_opt_model, quantization_mode=QuantizationMode.IntegerOps, @@ -55,3 +66,4 @@ class QuantizeHelper: force_fusions=True) onnx.save(quantized_onnx_model, quantized_model_path) logger.info(f"quantized model saved to:{quantized_model_path}") + logger.info(f'Size of quantized ONNX model(MB):{os.path.getsize(quantized_model_path)/(1024*1024)}')