# Copyright (c) Microsoft Corporation. All rights reserved. # Licensed under the MIT License. import onnx from onnx import version_converter from typing import List, Tuple, Union from .utils import ( IoMapEntry, IOEntryValuePreserver, create_custom_op_checker_context, sanitize_output_names, TENSOR_TYPE_TO_ONNX_TYPE, ) from .step import Step class PrePostProcessor: """ Class to handle running all the pre/post processing steps and updating the model. """ def __init__(self, inputs: List[onnx.ValueInfoProto] = None, onnx_opset: int = 16): """ Create a PrePostProcessor instance. Args: inputs: The inputs the model will use if pre-processing is added. onnx_opset: The ONNX opset to use. Minimum is 16. 18 or higher is strongly preferred if image resizing is involved due to its anti-aliasing ability. """ if onnx_opset < 16: raise ValueError("ONNX opset must be 16 or later.") self._onnx_opset = onnx_opset self._custom_op_checker_context = create_custom_op_checker_context(onnx_opset) self.pre_processors = [] self.post_processors = [] # Connections for each pre/post processor. 1:1 mapping with entries in pre_processors/post_processors self._pre_processor_connections = [] # type: List[List[IoMapEntry]] self._post_processor_connections = [] # type: List[List[IoMapEntry]] # explicitly join outputs from Steps in pre_processors to inputs of the original model # format is Step or step name, step_idx, name of graph input/output # # Pre-processing we connect Step output to original model: # - step_idx is for Step.output_names, and name is in graph.input # # Post-processing we connect the original model output to the Step input # - step_idx is for Step.input_names, and name is in graph.output self._pre_processing_joins = None # type: Union[None,List[Tuple[Union[Step, str], int, str]]] self._post_processing_joins = None # type: Union[None,List[Tuple[Union[Step, str], int, str]]] self._inputs = inputs if inputs else [] # preserve outputs from IOMapEntry, avoid it's consumed by the Follow-up steps. # we now can support a output value has more than one consumers with IOEntryValuePreserver. # IOEntryValuePreserver will preserve the output value and add it to the graph output # until consumer step is done. self.outputs_preserver = [] # type: List[IOEntryValuePreserver] def add_pre_processing(self, items: List[Union[Step, Tuple[Step, List[IoMapEntry]]]]): """ Add the pre-processing steps. The last step is automatically joined to the original model inputs. Options are: Add Step with default connection of outputs from the previous step (if available) to inputs of this step. Add tuple of Step and list of IoMapEntry instances for manual connections to previous steps. This will be used to override any automatic connections. If IoMapEntry.producer is None it is inferred to be the immediately previous Step. If IoMapEntry.producer is a step name it must match the name of a previous step. """ self.__add_processing(self.pre_processors, self._pre_processor_connections, items) def add_post_processing(self, items: List[Union[Step, Tuple[Step, List[IoMapEntry]]]]): """ Add the post-processing steps. The first step is automatically joined to the original model outputs. Options are: Add Step with default connection of outputs from the previous step (if available) to inputs of this step. Add tuple of Step and list of IoMapEntry instances for connections to previous steps. This will be used to override any automatic connections. If IoMapEntry.producer is None it is inferred to be the immediately previous Step. If IoMapEntry.producer is a step name it must match the name of a previous step. """ self.__add_processing(self.post_processors, self._post_processor_connections, items) def _add_connection(self, consumer: Step, entry: IoMapEntry): producer = self.__producer_from_step_or_str(entry.producer) # Black does annoying things with the multi-line 'if' conditions making the code far less readable # fmt: off if not ((producer in self.pre_processors or producer in self.post_processors) and (consumer in self.pre_processors or consumer in self.post_processors)): raise ValueError("Producer and Consumer processors must both be registered") if producer in self.pre_processors: if (consumer in self.pre_processors and self.pre_processors.index(producer) > self.pre_processors.index(consumer)): raise ValueError("Producer was registered after consumer and cannot be connected") elif producer in self.post_processors: if consumer not in self.post_processors: raise ValueError("Cannot connect pre-processor consumer with post-processor producer") elif self.post_processors.index(producer) > self.post_processors.index(consumer): raise ValueError("Producer was registered after consumer and cannot be connected") # fmt: on assert isinstance(producer, Step) consumer.connect(entry) def run(self, model: onnx.ModelProto): """ Update the model with the graph from each step in the pre and post processing pipelines. Args: model: model to add pre/post processing to. Returns: model with pre/post processing in it. """ # update the input model to the ONNX opset we're using. this is required as we implement the steps based on # the operator specs for this opset. model_opset = [ entry.version for entry in model.opset_import if entry.domain == "" or entry.domain == "ai.onnx" ][0] if model_opset > self._onnx_opset: # It will probably work if the user updates PRE_POST_PROCESSING_ONNX_OPSET to match the model # but there are no guarantees. # Would only break if ONNX operators used in the pre/post processing graphs have had spec changes. raise ValueError(f"Model opset is {model_opset} which is newer than the opset used by this script.") elif model_opset < self._onnx_opset: model = onnx.version_converter.convert_version(model, self._onnx_opset) def name_nodes(new_graph: onnx.GraphProto, prefix: str): # simple helper so all nodes are named. this makes it far easier to debug any issues. idx = 0 for n in new_graph.node: if not n.name: n.name = prefix + str(idx) idx += 1 def preserved_apply(processor: Step, *args): # Trying to activate the IOEntryValuePreserver and preserve outputs. # and deactivate the outputs when the current graph consumed them for preserver in self.outputs_preserver: if preserver.consumer == processor: preserver.is_active = False # IOEntryValuePreserver, preserve those outputs which has multiple consumers. # we explicitly add the output to the graph output. graph_outputs_to_maintain = [i.output for i in self.outputs_preserver if i.is_active] graph_for_step = processor.apply(*args, graph_outputs_to_maintain=graph_outputs_to_maintain) for preserver in self.outputs_preserver: if preserver.producer == processor: preserver.is_active = True preserver.output = processor.output_names[preserver.producer_idx] return graph_for_step def connect_and_run(graph: onnx.GraphProto, processor: Step, connections: List[IoMapEntry]): for connection in connections: assert connection.producer self._add_connection(processor, connection) return preserved_apply(processor, graph, self._custom_op_checker_context) # fix any invalid output names now if we're adding post-processing as the onnx parse_graph can't handle them if self.post_processors: sanitize_output_names(model.graph) graph = model.graph # add pre-processing if self.pre_processors: # create empty graph with pass through of the requested input name pre_process_graph = onnx.GraphProto() for i in self._inputs: pre_process_graph.input.append(i) pre_process_graph.output.append(i) # connect up the graph input names to the first pre-processing step based on order self.pre_processors[0]._connect_graph_inputs([vi.name for vi in self._inputs]) for idx, step in enumerate(self.pre_processors): pre_process_graph = connect_and_run(pre_process_graph, step, self._pre_processor_connections[idx]) # name all the nodes for easier debugging name_nodes(pre_process_graph, "pre_process_") if not self._pre_processing_joins: # default to 1:1 between outputs of last step with inputs of original model last_step = self.pre_processors[-1] num_entries = min(len(last_step.output_names), len(graph.input)) self._pre_processing_joins = [(last_step, i, graph.input[i].name) for i in range(0, num_entries)] # map the pre-processing outputs to graph inputs # we may need a natty way to get possible outputs after merge_graphs step_graph_outputs = [o.name for o in pre_process_graph.output] io_map = [] # type: List[Tuple[str, str]] for step, step_idx, graph_input in self._pre_processing_joins: io_map.append((step.output_names[step_idx], graph_input)) step_graph_outputs.remove((step.output_names[step_idx])) # add outputs from previous IoMapEntry producers to maintain them as graph outputs # until consumed by the final Step that requires them. step_graph_outputs += [ o.name for o in graph.output if o.name not in step_graph_outputs] external_outputs = [ i.output for i in self.outputs_preserver if i.is_active and i.output not in step_graph_outputs] if external_outputs: step_graph_outputs.extend(external_outputs) graph = onnx.compose.merge_graphs(pre_process_graph, graph, io_map, outputs=step_graph_outputs) # add post-processing if self.post_processors: orig_model_outputs = [o.name for o in model.graph.output] graph_outputs = [o.name for o in graph.output] # this may have additional outputs from pre-processing # create default joins if needed if not self._post_processing_joins: # default to 1:1 between outputs of original model with inputs of first post-processing step first_step = self.post_processors[0] num_entries = min(len(first_step.input_names), len(orig_model_outputs)) self._post_processing_joins = [(first_step, i, orig_model_outputs[i]) for i in range(0, num_entries)] # update the input names for the steps to match the values produced by the model for step, step_idx, graph_output in self._post_processing_joins: assert graph_output in graph_outputs step.input_names[step_idx] = graph_output # create empty graph with the values that will be available to the post-processing post_process_graph = onnx.GraphProto() for o in graph.output: post_process_graph.input.append(o) post_process_graph.output.append(o) for idx, step in enumerate(self.post_processors): post_process_graph = connect_and_run(post_process_graph, step, self._post_processor_connections[idx]) name_nodes(post_process_graph, "post_process_") # io_map should be 1:1 with the post-processing graph given we updated the step input names to match io_map = [(o, o) for o in graph_outputs] graph = onnx.compose.merge_graphs(graph, post_process_graph, io_map) # Make the output names nicer by removing prefixing from naming that occurred when applying the steps graph = PrePostProcessor.__cleanup_graph_output_names(graph) opset_imports = [onnx.helper.make_operatorsetid(domain, opset) for domain, opset in self._custom_op_checker_context.opset_imports.items()] # find_min_ir_version_for doesn't support custom domains until ONNX 1.14 so extract the ONNX opset from the # imports and only pass that in. ir_version = onnx.helper.find_min_ir_version_for([entry for entry in opset_imports if entry.domain == "" or entry.domain == "ai.onnx"]) new_model = onnx.helper.make_model(graph, opset_imports=opset_imports, ir_version=ir_version) onnx.checker.check_model(new_model) return new_model def __add_processing( self, processors: List[Step], processor_connections: List[List[IoMapEntry]], items: List[Union[Step, Tuple[Step, List[IoMapEntry]]]], ): """ Add the pre/post processing steps and join with existing steps. Args: processors: List of processors to add items to. processor_connections: Populated with connections between each step. 1:1 with entries in processors. items: Items to add to processors. Can be: A Step instance. This will be implicitly joined to the immediately previous Step if one exists. A tuple of (Step instance, list of IoMapEntry) The IoMapEntry values are used to manually join an output from a producer Step to an input of the current Step. In each IoMapEntry, if a step name is provided the producer Step will be searched for in all predecessor steps. It is valid for a post-processor step to consume output from a pre-processor step. """ for item in items: step = None explicit_io_map_entries = None if isinstance(item, Step): step = item elif isinstance(item, tuple): step, explicit_io_map_entries = item else: raise ValueError("Unexpected type " + str(type(item))) # start with implicit joins and replace with explicitly provided ones # this allows the user to specify the minimum number of manual joins. io_map_entries = [None] * len(step.input_names) # type: List[Union[None,IoMapEntry]] prev_step = None if len(processors) == 0 else processors[-1] if prev_step: # default is connecting as many outputs from the previous step as possible for i in range(0, min(len(prev_step.output_names), len(step.input_names))): io_map_entries[i] = IoMapEntry(prev_step, i, i) # add explicit connections if explicit_io_map_entries: for entry in explicit_io_map_entries: if not entry.producer: producer = prev_step else: producer = self.__producer_from_step_or_str(entry.producer) # throws if not found io_map_entries[entry.consumer_idx] = IoMapEntry(producer, entry.producer_idx, entry.consumer_idx) self.outputs_preserver.append(IOEntryValuePreserver(producer, step, entry.producer_idx)) processors.append(step) processor_connections.append([entry for entry in io_map_entries if entry is not None]) def __producer_from_step_or_str(self, entry: Union[Step, str]): if isinstance(entry, Step): return entry if isinstance(entry, str): match = (next((s for s in self.pre_processors if s.name == entry), None) or next((s for s in self.post_processors if s.name == entry), None)) # fmt: skip if not match: raise ValueError(f"Step named {entry} was not found") return match @staticmethod def __cleanup_graph_output_names(graph: onnx.GraphProto): """ Hide the prefixing of names that happens when we merge the graphs from the pre/post processing steps. Not essential but makes the graph outputs look far nicer. """ # for each output create identity node to remove prefixing io_map = [] fixes = onnx.GraphProto() # manually handle naming clashes input_names = set([i.name for i in graph.input]) used_names = set(input_names) conflicts = 0 for o in graph.output: if not o.name.startswith(Step.prefix): continue # we will create a small graph to do the renames so the output of the original graph will be an input # to that 'fixer' graph io_map.append((o.name, o.name)) clean_name = o.name while clean_name.startswith(Step.prefix): # output from last step will have one prefixing stage that adds Step._prefix + '_' # e.g. '_ppp8_' next_underscore = clean_name.find("_", 1) if next_underscore > 0: # this check shouldn't be necessary as we always add the trailing '_' when prefixing... if len(clean_name) > next_underscore + 1: next_underscore += 1 clean_name = clean_name[next_underscore:] # handle things like super resolution where there's an 'image' input and 'image' output if clean_name in input_names: clean_name += "_out" orig_clean_name = clean_name while clean_name in used_names: conflicts += 1 clean_name = f"{orig_clean_name}{conflicts}" used_names.add(clean_name) renamer = onnx.helper.make_node("Identity", [o.name], [clean_name], f"Rename {o.name}") fixes.node.append(renamer) fixes.input.append(o) new_output = fixes.output.add() new_output.name = clean_name new_output.type.CopyFrom(o.type) # merge if we have any renaming to do if io_map: graph = onnx.compose.merge_graphs(graph, fixes, io_map) return graph