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

Show quantization model size in benchmark of transformer (#4626)

Browse source 
* Show quantization model size in benchmark of transformer

* refine model size calculation
This commit is contained in:
Yufeng Li 2020-07-27 23:56:33 -07:00 committed by GitHub
parent 73c99f8269
commit a06cf6a3b3
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
2 changed files with 19 additions and 5 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

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

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

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

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