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