def run_onnx_optimizer(onnx_model):
"""Run ONNX's optimization routines.
ONNX Optimizer was moved to an external library in
version 1.9. Attempt to use the optimizer in onnx if
it is available, fall back to the standalone
onnxoptimizer otherwise, and return the model
unoptimized if neither are available.
"""
try:
onnx_polish_model = onnx.utils.polish_model
except AttributeError:
pass
else:
return onnx_polish_model(onnx_model)
try:
# pylint: disable=import-outside-toplevel
import onnxoptimizer
except ImportError:
pass
else:
return onnxoptimizer.optimize(onnx_model)
return model
def optimize(model: onnx.ModelProto, skip_fuse_bn: bool,
skipped_optimizers: Optional[Sequence[str]]) -> onnx.ModelProto:
"""
:param model: The onnx model.
:return: The optimized onnx model.
Before simplifying, use this method to generate value_info, which is used in `forward_all`
After simplifying, use this method to fold constants generated in previous step into initializer,
and eliminate unused constants.
"""
onnx.checker.check_model(model)
onnx.helper.strip_doc_string(model)
optimizers_list = onnxoptimizer.get_fuse_and_elimination_passes()
if not skip_fuse_bn:
optimizers_list.append('fuse_bn_into_conv')
if skipped_optimizers is not None:
for opt in skipped_optimizers:
try:
optimizers_list.remove(opt)
except ValueError:
pass
model = onnxoptimizer.optimize(model, optimizers_list, fixed_point=True)
onnx.checker.check_model(model)
return model
def optimize_fp_model(model_path):
check_model_extension(model_path)
model_proto = onnx.load(model_path)
# Convert static batch to dynamic batch
model_proto = convert_model_batch_to_dynamic(model_proto)
# Optimizer
model_proto = onnxoptimizer.optimize(model_proto,
['fuse_bn_into_conv', 'fuse_add_bias_into_conv', 'fuse_pad_into_conv',
'fuse_matmul_add_bias_into_gemm', 'fuse_transpose_into_gemm',
'eliminate_nop_flatten', 'eliminate_nop_pad', 'eliminate_nop_transpose',
'eliminate_unused_initializer', 'eliminate_duplicate_initializer'])
# Add Gemm name
add_fused_gemm_name(model_proto)
# Fuse bn into conv or gemm
fuse_bn(model_proto)
# Save optimized model
model_name = model_path.split(".")
model_path = model_name[0] + "_optimized.onnx"
onnx.save(model_proto, model_path)
return model_path
def remove_initializer_from_input():
args = get_args()
model = onnx.load(args.input)
if model.ir_version < 4:
print(
'Model with ir_version below 4 requires to include initilizer in graph input'
)
return
inputs = model.graph.input
name_to_input = {}
for input in inputs:
name_to_input[input.name] = input
for initializer in model.graph.initializer:
if initializer.name in name_to_input:
inputs.remove(name_to_input[initializer.name])
passes = [
"extract_constant_to_initializer",
"eliminate_unused_initializer"
]
optimized_model = onnxoptimizer.optimize(model, passes)
onnx.save(optimized_model, args.output)
def optimize(model: onnx.ModelProto, skip_fuse_bn: bool,
skipped_optimizers: Optional[Sequence[str]]) -> onnx.ModelProto:
"""Perform optimization on an ONNX model. Before simplifying, use this
method to generate value_info. After simplifying, use this method to fold
constants generated in previous step into initializer, and eliminate unused
constants.
Args:
model (onnx.ModelProto): The input ONNX model.
skip_fuse_bn (bool): Whether to skip fuse bn.
skipped_optimizers (Sequence[str]): List of optimizers to be skipped.
Returns:
onnx.ModelProto: The optimized model.
"""
# Due to a onnx bug, https://github.com/onnx/onnx/issues/2417,
# we need to add missing initializers into inputs
onnx.checker.check_model(model)
input_num = len(model.graph.input)
model = add_initializers_into_inputs(model)
onnx.helper.strip_doc_string(model)
onnx.checker.check_model(model)
optimizers_list = [
'eliminate_deadend', 'eliminate_nop_dropout', 'eliminate_nop_cast',
'eliminate_nop_monotone_argmax', 'eliminate_nop_pad',
'extract_constant_to_initializer', 'eliminate_unused_initializer',
'eliminate_nop_transpose', 'eliminate_identity',
'fuse_add_bias_into_conv', 'fuse_consecutive_concats',
'fuse_consecutive_log_softmax', 'fuse_consecutive_reduce_unsqueeze',
'fuse_consecutive_squeezes', 'fuse_consecutive_transposes',
'fuse_matmul_add_bias_into_gemm', 'fuse_pad_into_conv',
'fuse_transpose_into_gemm'
]
if not skip_fuse_bn:
optimizers_list.append('fuse_bn_into_conv')
if skipped_optimizers is not None:
for opt in skipped_optimizers:
try:
optimizers_list.remove(opt)
except ValueError:
pass
model = onnxoptimizer.optimize(model, optimizers_list, fixed_point=True)
if model.ir_version > 3:
del model.graph.input[input_num:]
onnx.checker.check_model(model)
return model
def transform(self, model: onnx.ModelProto) -> onnx.ModelProto:
for init in model.graph.initializer:
if init.name not in [inp.name for inp in model.graph.input]:
dims = numpy_helper.to_array(init).shape
model.graph.input.append(
make_tensor_value_info(name=init.name,
elem_type=init.data_type,
shape=dims))
model = optimizer.optimize(model,
passes=['extract_constant_to_initializer'])
# make scalar initializer's dim=0 in Add, Sub, Mul, Div
initializer = {init.name: init for init in model.graph.initializer}
input_vi = {inp.name: inp for inp in model.graph.input}
for node in model.graph.node:
if not any(node.op_type == op
for op in ['Add', 'Sub', 'Mul', 'Div']):
continue
for node_input in node.input:
if node_input not in initializer:
continue
init = initializer[node_input]
vi = input_vi[node_input]
arr = numpy_helper.to_array(init)
if arr.size == 1 and arr.shape == (1, ):
model.graph.initializer.remove(init)
model.graph.initializer.append(
make_tensor(name=init.name,
data_type=init.data_type,
dims=(),
vals=numpy_helper.to_array(init)))
model.graph.input.remove(vi)
model.graph.input.append(
make_tensor_value_info(
name=init.name,
elem_type=vi.type.tensor_type.elem_type,
shape=()))
model = make_model(model.graph)
model = utils.rebuild_model(model, model.graph.node)
return model
def export_onnx_model(model, inputs):
"""
Trace and export a model to onnx format.
Args:
model (nn.Module):
inputs (torch.Tensor): the model will be called by `model(*inputs)`
Returns:
an onnx model
"""
assert isinstance(model, torch.nn.Module)
# make sure all modules are in eval mode, onnx may change the training state
# of the module if the states are not consistent
def _check_eval(module):
assert not module.training
model.apply(_check_eval)
logger.info("Beginning ONNX file converting")
# Export the model to ONNX
with torch.no_grad():
with io.BytesIO() as f:
torch.onnx.export(
model,
inputs,
f,
operator_export_type=OperatorExportTypes.ONNX_ATEN_FALLBACK,
# verbose=True, # NOTE: uncomment this for debugging
# export_params=True,
)
onnx_model = onnx.load_from_string(f.getvalue())
logger.info("Completed convert of ONNX model")
# Apply ONNX's Optimization
logger.info("Beginning ONNX model path optimization")
all_passes = onnxoptimizer.get_available_passes()
passes = [
"extract_constant_to_initializer", "eliminate_unused_initializer",
"fuse_bn_into_conv"
]
assert all(p in all_passes for p in passes)
onnx_model = onnxoptimizer.optimize(onnx_model, passes)
logger.info("Completed ONNX model path optimization")
return onnx_model
def optimize(self, optimizations: List[str] = None, fixed_point: bool = False):
"""
Use ONNX optimizer to optimize the ONNX model. The optimizations supported can be seen by calling
'onnxoptimizer.get_available_passes()'
:param optimizations: List of possible optimizations. If None, all of the optimizations will be used. Defaulted
to None.
:param fixed_point: Optimize the weights using fixed point. Defaulted to False.
"""
# Set the ONNX optimizations list:
onnx_optimizations = onnxoptimizer.get_fuse_and_elimination_passes()
if optimizations is None:
# Set to all optimizations:
optimizations = onnx_optimizations
# Optimize the model:
self._model = onnxoptimizer.optimize(
self._model, passes=optimizations, fixed_point=fixed_point
)
def optimize_onnx(input, init=False, predict=False):
passes = [
'fuse_consecutive_transposes', 'eliminate_nop_transpose',
'fuse_transpose_into_gemm', 'lift_lexical_references'
]
if init:
passes.append('split_init')
if predict:
passes.append('split_predict')
try:
out = onnx.optimizer.optimize(input, passes)
except AttributeError:
warnings.warn(
"OptimizerWarning: optimizer module not found in ONNX version {}"
.format(onnx.__version__))
# ONNX does no ship onnx.optimizer since version 1.9+
import onnxoptimizer
out = onnxoptimizer.optimize(input, passes)
return out
def optimize(model: onnx.ModelProto, skip_fuse_bn: bool, skipped_optimizers: Optional[Sequence[str]]) -> onnx.ModelProto:
"""
:model参数: 待优化的ONXX模型.
:return: 优化之后的ONNX模型.
简化之前, 使用这个方法产生会在'forward_all'用到的ValueInfo
简化之后，使用这个方法去折叠前一步产生的常量到initializer中并且消除没被使用的常量
"""
onnx.checker.check_model(model)
onnx.helper.strip_doc_string(model)
optimizers_list = [
'eliminate_deadend',
'eliminate_nop_dropout',
'eliminate_nop_cast',
'eliminate_nop_monotone_argmax', 'eliminate_nop_pad',
'extract_constant_to_initializer', 'eliminate_unused_initializer',
'eliminate_nop_transpose',
'eliminate_nop_flatten', 'eliminate_identity',
'fuse_add_bias_into_conv',
'fuse_consecutive_concats',
'fuse_consecutive_log_softmax',
'fuse_consecutive_reduce_unsqueeze', 'fuse_consecutive_squeezes',
'fuse_consecutive_transposes', 'fuse_matmul_add_bias_into_gemm',
'fuse_pad_into_conv', 'fuse_transpose_into_gemm', 'eliminate_duplicate_initializer'
]
if not skip_fuse_bn:
optimizers_list.append('fuse_bn_into_conv')
if skipped_optimizers is not None:
for opt in skipped_optimizers:
try:
optimizers_list.remove(opt)
except ValueError:
pass
model = onnxoptimizer.optimize(model, optimizers_list,
fixed_point=True)
onnx.checker.check_model(model)
return model
def optimize(model: onnx.ModelProto, skip_fuse_bn: bool,
skipped_optimizers: Optional[Sequence[str]]) -> onnx.ModelProto:
"""
:param model: The onnx model.
:return: The optimized onnx model.
Before simplifying, use this method to generate value_info, which is used in `forward_all`
After simplifying, use this method to fold constants generated in previous step into initializer,
and eliminate unused constants.
"""
onnx.checker.check_model(model)
onnx.helper.strip_doc_string(model)
optimizers_list = [
'eliminate_deadend', 'eliminate_nop_dropout', 'eliminate_nop_cast',
'eliminate_nop_monotone_argmax', 'eliminate_nop_pad',
'extract_constant_to_initializer', 'eliminate_unused_initializer',
'eliminate_nop_transpose', 'eliminate_nop_flatten',
'eliminate_identity', 'fuse_add_bias_into_conv',
'fuse_consecutive_concats', 'fuse_consecutive_log_softmax',
'fuse_consecutive_reduce_unsqueeze', 'fuse_consecutive_squeezes',
'fuse_consecutive_transposes', 'fuse_matmul_add_bias_into_gemm',
'fuse_pad_into_conv', 'fuse_transpose_into_gemm',
'eliminate_duplicate_initializer'
]
if not skip_fuse_bn:
optimizers_list.append('fuse_bn_into_conv')
if skipped_optimizers is not None:
for opt in skipped_optimizers:
try:
optimizers_list.remove(opt)
except ValueError:
pass
model = onnxoptimizer.optimize(model, optimizers_list, fixed_point=True)
onnx.checker.check_model(model)
return model
def optimize(model: onnx.ModelProto, skip_fuse_bn: bool,
skipped_optimizers: Optional[Sequence[str]]) -> onnx.ModelProto:
"""
:model参数: 待优化的ONXX模型.
:return: 优化之后的ONNX模型.
简化之前, 使用这个方法产生会在'forward_all'用到的ValueInfo
简化之后，使用这个方法去折叠前一步产生的常量到initializer中并且消除没被使用的常量
"""
onnx.checker.check_model(model)
onnx.helper.strip_doc_string(model)
optimizers_list = onnxoptimizer.get_fuse_and_elimination_passes()
if not skip_fuse_bn:
optimizers_list.append('fuse_bn_into_conv')
if skipped_optimizers is not None:
for opt in skipped_optimizers:
try:
optimizers_list.remove(opt)
except ValueError:
pass
model = onnxoptimizer.optimize(model, optimizers_list, fixed_point=True)
onnx.checker.check_model(model)
return model
def convert(self):
inputs = []
parameters = []
onnx_nodes = []
outputs = []
unsupported_oprs = []
set_opset_version(self.opset_version)
def need_convert(opr):
is_const = [data.np_data is not None for data in opr.inp_tensors]
return not all(is_const) or len(opr.inp_tensors) == 0
def deduplication(inputs):
names = []
results = []
for i in inputs:
if i.name not in names:
results.append(i)
names.append(i.name)
return results
_, tensor_sources, _ = _add_input_tensors(self.net.graph_inputs)
inputs.extend(tensor_sources)
for opr in self.net.all_oprs:
if not need_convert(opr):
for tensor in opr.out_tensors:
if hasattr(tensor, "_var"):
tensor.np_data = get_symvar_value(tensor._var)
continue
converter_cls = MGE2ONNX.get(type(opr), None)
if converter_cls is None:
unsupported_oprs.append(opr)
continue
converter = converter_cls(opr, self.quantizer)
nodes, inps, params = converter.convert()
onnx_nodes.extend(nodes)
inputs.extend(inps)
parameters.extend(params)
inputs = deduplication(inputs)
parameters = deduplication(parameters)
unsupported_oprs = set(map(type, unsupported_oprs))
assert not unsupported_oprs, "Operators {} are not supported yet".format(
unsupported_oprs)
for output in self.net.graph_outputs:
def _get_onnx_dtype(output):
return mge2onnx_dtype_mapping[output.dtype]
out_tensor = onnx.helper.make_tensor_value_info(
output.name, _get_onnx_dtype(output), output.shape)
outputs.append(out_tensor)
onnx_graph = onnx.helper.make_graph(onnx_nodes,
self.graph_name,
inputs,
outputs,
initializer=parameters)
opset = onnx.helper.make_opsetid("", self.opset_version)
model = onnx.helper.make_model(
onnx_graph,
producer_name="MegEngine",
producer_version=mge.__version__,
opset_imports=[opset],
)
onnx.checker.check_model(model)
passes = [
"eliminate_deadend",
"extract_constant_to_initializer",
"eliminate_unused_initializer",
]
model = onnxoptimizer.optimize(model, passes=passes)
return model
def clean(graph: xpb2.GraphProto,
_optimize: bool = True,
_simplify: bool = True,
_remove_initializer: bool = True,
_producer: str = "sclblonnx",
_verbose: bool = True,
**kwargs):
""" clean cleans an ONNX graph using onnx tooling
This method will attempt to clean the supplied graph by
a. Removing initializers from input
b. Optimizing it using onnxoptimizer.optimize
c. Simplifying it using onnxsim.simplify
If one of these fails the method will print an error message and return the unaltered graph.
Args:
graph: An ONNX graph
_optimize: Boolean, default True. Optimize the model using onnxoptimizer.
_simplify: Boolean, default True. Simplify the model using simplify.
_remove_initializer: Boolean, default True. Remove initializers from input.
_producer: Optional string with producer name. Default 'sclblonnx' (used for internal conversion)
_verbose: Print user feedback; default True (note, errors are always printed).
**kwargs
Returns:
The cleaned ONNX graph, or the old graph if an error occurs.
"""
try:
if not 'opset_imports' in kwargs:
op = onnx.OperatorSetIdProto()
op.version = 12
mod = xhelp.make_model(graph,
producer_name=_producer,
opset_imports=[op],
**kwargs)
else:
mod = xhelp.make_model(graph, producer_name=_producer, **kwargs)
except Exception as e:
_print("Unable to create the model: " + str(e))
return graph
if _optimize:
try:
mod = onnxoptimizer.optimize(mod, glob.OPTIMIZER_PASSES, **kwargs)
except Exception as e:
_print("Unable to optimize your model: " + str(e))
return graph
if _simplify:
try:
mod, _ = simplify(mod, **kwargs)
except Exception as e:
_print("Unable to simplify your model: " + str(e))
return graph
# From: onnxruntime/tools/python/remove_initializer_from_input.py
graph = mod.graph
if _remove_initializer:
inputs = graph.input
name_to_input = {}
for input in inputs:
name_to_input[input.name] = input
for initializer in graph.initializer:
if initializer.name in name_to_input:
inputs.remove(name_to_input[initializer.name])
_print("The graph was successfully cleaned.", "MSG", (not _verbose))
return graph
def export_onnx(cls,
module: Module,
input_shape: Optional[Tuple[int, ...]] = None,
export_path: Optional[str] = None,
input_t: Optional[Union[Tensor, QuantTensor]] = None,
disable_warnings=True,
**kwargs):
if onnx is None or opt is None:
raise ModuleNotFoundError(
"Installation of onnx and onnxoptimizer is required.")
if input_shape is None and input_t is None:
raise RuntimeError(
"Export requires to pass in either input_shape or input_t")
if input_shape is not None and input_t is not None:
raise RuntimeError(
"Export accepts either an input shape or an input tensor, not both"
)
cls.solve_keep_initializers_as_inputs(kwargs)
cls.solve_enable_onnx_checker(kwargs)
cls.register_custom_fns()
with torch.no_grad():
with ExportContext(cls):
with warnings.catch_warnings():
if disable_warnings:
warnings.simplefilter("ignore")
training_state = module.training
module = module.eval()
module.apply(cls.set_export_handler)
if input_t is None:
input_t = torch.empty(input_shape, dtype=torch.float)
# do a forward pass with the dummy input to e.g. store input/output shapes
cls._cache_inp_out(module, input_t)
# Dequantize QuantTensor, if any and enabled
if isinstance(input_t, QuantTensor):
if cls.dequantize_tracing_input:
input_t = input_t.value
else:
input_t = (input_t, )
# enable export mode, this triggers collecting export values into handlers
module.apply(
lambda m: cls.set_export_mode(m, enabled=True))
# temporarily disable input caching to avoid collectives empty debug values
module.apply(
lambda m: _override_inp_caching_mode(m, enabled=False))
# perform export pass
with ExitStack() as stack:
for mgr in cls._trace_patches():
stack.enter_context(mgr)
if export_path is not None:
export_target = export_path
else:
model_bytes = BytesIO()
export_target = model_bytes
torch.onnx.export(module, input_t, export_target,
**kwargs)
# restore the model to previous properties
module.apply(lambda m: _restore_inp_caching_mode(m))
module.apply(
lambda m: cls.set_export_mode(m, enabled=False))
module.train(training_state)
# do some cleanup on the exported ONNX model
if export_path is not None:
model = onnx.load(export_path)
else:
model = onnx.ModelProto.FromString(
model_bytes.getvalue())
model = opt.optimize(model, cls.onnx_passes)
model = cls.apply_model_transforms(model)
if export_path is not None:
onnx.save(model, export_path)
return model
def export_onnx(cls,
module: Module,
input_shape: Optional[Tuple[int, ...]] = None,
export_path: Optional[str] = None,
input_t: Optional[Union[Tensor, QuantTensor]] = None,
**kwargs):
"""
* input_shape : tuple describing the shape of network input e.g. (1, 1, 28, 28)
* export_path : ONNX filename to export to
* input_t : if specified, do an initial forward pass with this value. this
may be necessary for QuantTensor caching.
* torch_onnx_kwargs : will be passed as kwargs to torch.onnx.export
"""
if onnx is None or opt is None:
raise ModuleNotFoundError(
"Installation of onnx and onnxoptimizer is required.")
if input_shape is None and input_t is None:
raise RuntimeError(
"Export requires to pass in either input_shape or input_t")
if input_shape is not None and input_t is not None:
raise RuntimeError(
"Export accepts either an input shape or an input tensor, not both"
)
cls.solve_keep_initializers_as_inputs(kwargs)
cls.solve_enable_onnx_checker(kwargs)
with torch.no_grad():
with ExportContext(cls):
training_state = module.training
module = module.eval()
module.apply(cls.set_export_handler)
if input_t is None:
input_t = torch.empty(input_shape, dtype=torch.float)
# do a forward pass with the dummy input to e.g. store input/output shapes
cls._cache_inp_out(module, input_t)
# Dequantize QuantTensor, if any
if isinstance(input_t, QuantTensor):
input_t = input_t.value
# enable export mode, this triggers collecting export values into handlers
module.apply(lambda m: cls.set_export_mode(m, enabled=True))
# temporarily disable input caching to avoid collectives empty debug values
module.apply(
lambda m: _override_inp_caching_mode(m, enabled=False))
# perform export pass
with ExitStack() as stack:
for mgr in cls._trace_patches():
stack.enter_context(mgr)
if export_path is not None:
torch.onnx.export(module, input_t, export_path,
**kwargs)
else:
model_bytes = BytesIO()
torch.onnx.export(module, input_t, model_bytes,
**kwargs)
# restore the model to previous properties
module.apply(lambda m: _restore_inp_caching_mode(m))
module.apply(lambda m: cls.set_export_mode(m, enabled=False))
module.train(training_state)
# do some cleanup on the exported ONNX model
if export_path is not None:
model = onnx.load(export_path)
else:
model = onnx.ModelProto.FromString(model_bytes.getvalue())
model = opt.optimize(model, cls.onnx_passes)
model = cls.apply_model_transforms(model)
if export_path is not None:
onnx.save(model, export_path)
return model
