microsoft/onnxruntime-extensions
Publicmirrored from https://github.com/microsoft/onnxruntime-extensionsAvailable
onnxruntime_extensions/tools/pre_post_processing/pre_post_processor.py
350lines · modecode
| 1 | # Copyright (c) Microsoft Corporation. All rights reserved. |
| 2 | # Licensed under the MIT License. |
| 3 | |
| 4 | import onnx |
| 5 | |
| 6 | from onnx import version_converter |
| 7 | from typing import List, Tuple, Union |
| 8 | |
| 9 | from .utils import ( |
| 10 | IoMapEntry, |
| 11 | create_custom_op_checker_context, |
| 12 | sanitize_output_names, |
| 13 | TENSOR_TYPE_TO_ONNX_TYPE, |
| 14 | ) |
| 15 | from .step import Step |
| 16 | |
| 17 | |
| 18 | class PrePostProcessor: |
| 19 | """ |
| 20 | Class to handle running all the pre/post processing steps and updating the model. |
| 21 | """ |
| 22 | |
| 23 | def __init__(self, inputs: List[onnx.ValueInfoProto] = None, onnx_opset: int = 16): |
| 24 | """ |
| 25 | Create a PrePostProcessor instance. |
| 26 | |
| 27 | Args: |
| 28 | inputs: The inputs the model will use if pre-processing is added. |
| 29 | onnx_opset: The ONNX opset to use. |
| 30 | Minimum is 16. 18 or higher is strongly preferred if image resizing is involved due to its |
| 31 | anti-aliasing ability. |
| 32 | """ |
| 33 | |
| 34 | if onnx_opset < 16: |
| 35 | raise ValueError("ONNX opset must be 16 or later.") |
| 36 | |
| 37 | self._onnx_opset = onnx_opset |
| 38 | self._custom_op_checker_context = create_custom_op_checker_context(onnx_opset) |
| 39 | |
| 40 | self.pre_processors = [] |
| 41 | self.post_processors = [] |
| 42 | |
| 43 | # Connections for each pre/post processor. 1:1 mapping with entries in pre_processors/post_processors |
| 44 | self._pre_processor_connections = [] # type: List[List[IoMapEntry]] |
| 45 | self._post_processor_connections = [] # type: List[List[IoMapEntry]] |
| 46 | |
| 47 | # explicitly join outputs from Steps in pre_processors to inputs of the original model |
| 48 | # format is Step or step name, step_idx, name of graph input/output |
| 49 | # |
| 50 | # Pre-processing we connect Step output to original model: |
| 51 | # - step_idx is for Step.output_names, and name is in graph.input |
| 52 | # |
| 53 | # Post-processing we connect the original model output to the Step input |
| 54 | # - step_idx is for Step.input_names, and name is in graph.output |
| 55 | self._pre_processing_joins = None # type: Union[None,List[Tuple[Union[Step, str], int, str]]] |
| 56 | self._post_processing_joins = None # type: Union[None,List[Tuple[Union[Step, str], int, str]]] |
| 57 | |
| 58 | self._inputs = inputs if inputs else [] |
| 59 | |
| 60 | def add_pre_processing(self, items: List[Union[Step, Tuple[Step, List[IoMapEntry]]]]): |
| 61 | """ |
| 62 | Add the pre-processing steps. The last step is automatically joined to the original model inputs. |
| 63 | |
| 64 | Options are: |
| 65 | Add Step with default connection of outputs from the previous step (if available) to inputs of this step. |
| 66 | Add tuple of Step and list of IoMapEntry instances for manual connections to previous steps. This will be |
| 67 | used to override any automatic connections. |
| 68 | If IoMapEntry.producer is None it is inferred to be the immediately previous Step. |
| 69 | If IoMapEntry.producer is a step name it must match the name of a previous step. |
| 70 | """ |
| 71 | self.__add_processing(self.pre_processors, self._pre_processor_connections, items) |
| 72 | |
| 73 | def add_post_processing(self, items: List[Union[Step, Tuple[Step, List[IoMapEntry]]]]): |
| 74 | """ |
| 75 | Add the post-processing steps. The first step is automatically joined to the original model outputs. |
| 76 | |
| 77 | Options are: |
| 78 | Add Step with default connection of outputs from the previous step (if available) to inputs of this step. |
| 79 | Add tuple of Step and list of IoMapEntry instances for connections to previous steps. This will be |
| 80 | used to override any automatic connections. |
| 81 | If IoMapEntry.producer is None it is inferred to be the immediately previous Step. |
| 82 | If IoMapEntry.producer is a step name it must match the name of a previous step. |
| 83 | """ |
| 84 | self.__add_processing(self.post_processors, self._post_processor_connections, items) |
| 85 | |
| 86 | def _add_connection(self, consumer: Step, entry: IoMapEntry): |
| 87 | producer = self.__producer_from_step_or_str(entry.producer) |
| 88 | |
| 89 | # Black does annoying things with the multi-line 'if' conditions making the code far less readable |
| 90 | # fmt: off |
| 91 | if not ((producer in self.pre_processors or producer in self.post_processors) and |
| 92 | (consumer in self.pre_processors or consumer in self.post_processors)): |
| 93 | raise ValueError("Producer and Consumer processors must both be registered") |
| 94 | |
| 95 | if producer in self.pre_processors: |
| 96 | if (consumer in self.pre_processors and |
| 97 | self.pre_processors.index(producer) > self.pre_processors.index(consumer)): |
| 98 | raise ValueError("Producer was registered after consumer and cannot be connected") |
| 99 | elif producer in self.post_processors: |
| 100 | if consumer not in self.post_processors: |
| 101 | raise ValueError("Cannot connect pre-processor consumer with post-processor producer") |
| 102 | elif self.post_processors.index(producer) > self.post_processors.index(consumer): |
| 103 | raise ValueError("Producer was registered after consumer and cannot be connected") |
| 104 | # fmt: on |
| 105 | |
| 106 | assert isinstance(producer, Step) |
| 107 | consumer.connect(entry) |
| 108 | |
| 109 | |
| 110 | def run(self, model: onnx.ModelProto): |
| 111 | """ |
| 112 | Update the model with the graph from each step in the pre and post processing pipelines. |
| 113 | |
| 114 | Args: |
| 115 | model: model to add pre/post processing to. |
| 116 | |
| 117 | Returns: |
| 118 | model with pre/post processing in it. |
| 119 | """ |
| 120 | |
| 121 | # update the input model to the ONNX opset we're using. this is required as we implement the steps based on |
| 122 | # the operator specs for this opset. |
| 123 | model_opset = [ |
| 124 | entry.version for entry in model.opset_import if entry.domain == "" or entry.domain == "ai.onnx" |
| 125 | ][0] |
| 126 | |
| 127 | if model_opset > self._onnx_opset: |
| 128 | # It will probably work if the user updates PRE_POST_PROCESSING_ONNX_OPSET to match the model |
| 129 | # but there are no guarantees. |
| 130 | # Would only break if ONNX operators used in the pre/post processing graphs have had spec changes. |
| 131 | raise ValueError(f"Model opset is {model_opset} which is newer than the opset used by this script.") |
| 132 | elif model_opset < self._onnx_opset: |
| 133 | model = onnx.version_converter.convert_version(model, self._onnx_opset) |
| 134 | |
| 135 | def name_nodes(new_graph: onnx.GraphProto, prefix: str): |
| 136 | # simple helper so all nodes are named. this makes it far easier to debug any issues. |
| 137 | idx = 0 |
| 138 | for n in new_graph.node: |
| 139 | if not n.name: |
| 140 | n.name = prefix + str(idx) |
| 141 | idx += 1 |
| 142 | |
| 143 | def connect_and_run(graph: onnx.GraphProto, processor: Step, connections: List[IoMapEntry]): |
| 144 | for connection in connections: |
| 145 | assert connection.producer |
| 146 | self._add_connection(processor, connection) |
| 147 | |
| 148 | return processor.apply(graph, self._custom_op_checker_context) |
| 149 | |
| 150 | # fix any invalid output names now if we're adding post-processing as the onnx parse_graph can't handle them |
| 151 | if self.post_processors: |
| 152 | sanitize_output_names(model.graph) |
| 153 | |
| 154 | graph = model.graph |
| 155 | # add pre-processing |
| 156 | if self.pre_processors: |
| 157 | # create empty graph with pass through of the requested input name |
| 158 | pre_process_graph = onnx.GraphProto() |
| 159 | for i in self._inputs: |
| 160 | pre_process_graph.input.append(i) |
| 161 | pre_process_graph.output.append(i) |
| 162 | |
| 163 | for idx, step in enumerate(self.pre_processors): |
| 164 | pre_process_graph = connect_and_run(pre_process_graph, step, self._pre_processor_connections[idx]) |
| 165 | |
| 166 | # name all the nodes for easier debugging |
| 167 | name_nodes(pre_process_graph, "pre_process_") |
| 168 | |
| 169 | if not self._pre_processing_joins: |
| 170 | # default to 1:1 between outputs of last step with inputs of original model |
| 171 | last_step = self.pre_processors[-1] |
| 172 | num_entries = min(len(last_step.output_names), len(graph.input)) |
| 173 | self._pre_processing_joins = [(last_step, i, graph.input[i].name) for i in range(0, num_entries)] |
| 174 | |
| 175 | # map the pre-processing outputs to graph inputs |
| 176 | io_map = [] # type: List[Tuple[str, str]] |
| 177 | for step, step_idx, graph_input in self._pre_processing_joins: |
| 178 | io_map.append((step.output_names[step_idx], graph_input)) |
| 179 | |
| 180 | graph = onnx.compose.merge_graphs(pre_process_graph, graph, io_map) |
| 181 | |
| 182 | # add post-processing |
| 183 | if self.post_processors: |
| 184 | orig_model_outputs = [o.name for o in model.graph.output] |
| 185 | graph_outputs = [o.name for o in graph.output] # this may have additional outputs from pre-processing |
| 186 | |
| 187 | # create default joins if needed |
| 188 | if not self._post_processing_joins: |
| 189 | # default to 1:1 between outputs of original model with inputs of first post-processing step |
| 190 | first_step = self.post_processors[0] |
| 191 | num_entries = min(len(first_step.input_names), len(orig_model_outputs)) |
| 192 | self._post_processing_joins = [(first_step, i, orig_model_outputs[i]) for i in range(0, num_entries)] |
| 193 | |
| 194 | # update the input names for the steps to match the values produced by the model |
| 195 | for step, step_idx, graph_output in self._post_processing_joins: |
| 196 | assert graph_output in graph_outputs |
| 197 | step.input_names[step_idx] = graph_output |
| 198 | |
| 199 | # create empty graph with the values that will be available to the post-processing |
| 200 | post_process_graph = onnx.GraphProto() |
| 201 | for o in graph.output: |
| 202 | post_process_graph.input.append(o) |
| 203 | post_process_graph.output.append(o) |
| 204 | |
| 205 | for idx, step in enumerate(self.post_processors): |
| 206 | post_process_graph = connect_and_run(post_process_graph, step, self._post_processor_connections[idx]) |
| 207 | |
| 208 | name_nodes(post_process_graph, "post_process_") |
| 209 | |
| 210 | # io_map should be 1:1 with the post-processing graph given we updated the step input names to match |
| 211 | io_map = [(o, o) for o in graph_outputs] |
| 212 | graph = onnx.compose.merge_graphs(graph, post_process_graph, io_map) |
| 213 | |
| 214 | # Make the output names nicer by removing prefixing from naming that occurred when applying the steps |
| 215 | graph = PrePostProcessor.__cleanup_graph_output_names(graph) |
| 216 | |
| 217 | opset_imports = [onnx.helper.make_operatorsetid(domain, opset) |
| 218 | for domain, opset in self._custom_op_checker_context.opset_imports.items()] |
| 219 | new_model = onnx.helper.make_model(graph, opset_imports=opset_imports) |
| 220 | |
| 221 | onnx.checker.check_model(new_model) |
| 222 | |
| 223 | return new_model |
| 224 | |
| 225 | def __add_processing( |
| 226 | self, |
| 227 | processors: List[Step], |
| 228 | processor_connections: List[List[IoMapEntry]], |
| 229 | items: List[Union[Step, Tuple[Step, List[IoMapEntry]]]], |
| 230 | ): |
| 231 | """ |
| 232 | Add the pre/post processing steps and join with existing steps. |
| 233 | |
| 234 | Args: |
| 235 | processors: List of processors to add items to. |
| 236 | processor_connections: Populated with connections between each step. 1:1 with entries in processors. |
| 237 | items: Items to add to processors. |
| 238 | Can be: |
| 239 | A Step instance. This will be implicitly joined to the immediately previous Step if one exists. |
| 240 | A tuple of (Step instance, list of IoMapEntry) |
| 241 | The IoMapEntry values are used to manually join an output from a producer Step to an input |
| 242 | of the current Step. |
| 243 | In each IoMapEntry, if a step name is provided the producer Step will be searched for in all |
| 244 | predecessor steps. It is valid for a post-processor step to consume output from a |
| 245 | pre-processor step. |
| 246 | """ |
| 247 | |
| 248 | for item in items: |
| 249 | step = None |
| 250 | explicit_io_map_entries = None |
| 251 | |
| 252 | if isinstance(item, Step): |
| 253 | step = item |
| 254 | elif isinstance(item, tuple): |
| 255 | step, explicit_io_map_entries = item |
| 256 | else: |
| 257 | raise ValueError("Unexpected type " + str(type(item))) |
| 258 | |
| 259 | # start with implicit joins and replace with explicitly provided ones |
| 260 | # this allows the user to specify the minimum number of manual joins. |
| 261 | io_map_entries = [None] * len(step.input_names) # type: List[Union[None,IoMapEntry]] |
| 262 | prev_step = None if len(processors) == 0 else processors[-1] |
| 263 | if prev_step: |
| 264 | # default is connecting as many outputs from the previous step as possible |
| 265 | for i in range(0, min(len(prev_step.output_names), len(step.input_names))): |
| 266 | io_map_entries[i] = IoMapEntry(prev_step, i, i) |
| 267 | |
| 268 | # add explicit connections |
| 269 | if explicit_io_map_entries: |
| 270 | for entry in explicit_io_map_entries: |
| 271 | if not entry.producer: |
| 272 | producer = prev_step |
| 273 | else: |
| 274 | producer = self.__producer_from_step_or_str(entry.producer) # throws if not found |
| 275 | |
| 276 | io_map_entries[entry.consumer_idx] = IoMapEntry(producer, entry.producer_idx, entry.consumer_idx) |
| 277 | |
| 278 | processors.append(step) |
| 279 | processor_connections.append([entry for entry in io_map_entries if entry is not None]) |
| 280 | |
| 281 | def __producer_from_step_or_str(self, entry: Union[Step, str]): |
| 282 | if isinstance(entry, Step): |
| 283 | return entry |
| 284 | if isinstance(entry, str): |
| 285 | match = (next((s for s in self.pre_processors if s.name == entry), None) or |
| 286 | next((s for s in self.post_processors if s.name == entry), None)) # fmt: skip |
| 287 | |
| 288 | if not match: |
| 289 | raise ValueError(f"Step named {entry} was not found") |
| 290 | |
| 291 | return match |
| 292 | |
| 293 | @staticmethod |
| 294 | def __cleanup_graph_output_names(graph: onnx.GraphProto): |
| 295 | """ |
| 296 | Hide the prefixing of names that happens when we merge the graphs from the pre/post processing steps. |
| 297 | Not essential but makes the graph outputs look far nicer. |
| 298 | """ |
| 299 | |
| 300 | # for each output create identity node to remove prefixing |
| 301 | io_map = [] |
| 302 | fixes = onnx.GraphProto() |
| 303 | fixes.input.extend(graph.output) |
| 304 | |
| 305 | # manually handle naming clashes |
| 306 | input_names = set([i.name for i in graph.input]) |
| 307 | used_names = set(input_names) |
| 308 | conflicts = 0 |
| 309 | |
| 310 | for o in graph.output: |
| 311 | if not o.name.startswith(Step.prefix): |
| 312 | continue |
| 313 | |
| 314 | # we will create a small graph to do the renames so the output of the original graph will be an input |
| 315 | # to that 'fixer' graph |
| 316 | io_map.append((o.name, o.name)) |
| 317 | clean_name = o.name |
| 318 | while clean_name.startswith(Step.prefix): |
| 319 | # output from last step will have one prefixing stage that adds Step._prefix + '_' |
| 320 | # e.g. '_ppp8_<orig_name>' |
| 321 | next_underscore = clean_name.find("_", 1) |
| 322 | if next_underscore > 0: |
| 323 | # this check shouldn't be necessary as we always add the trailing '_' when prefixing... |
| 324 | if len(clean_name) > next_underscore + 1: |
| 325 | next_underscore += 1 |
| 326 | clean_name = clean_name[next_underscore:] |
| 327 | |
| 328 | # handle things like super resolution where there's an 'image' input and 'image' output |
| 329 | if clean_name in input_names: |
| 330 | clean_name += "_out" |
| 331 | |
| 332 | orig_clean_name = clean_name |
| 333 | while clean_name in used_names: |
| 334 | conflicts += 1 |
| 335 | clean_name = f"{orig_clean_name}{conflicts}" |
| 336 | |
| 337 | used_names.add(clean_name) |
| 338 | |
| 339 | renamer = onnx.helper.make_node("Identity", [o.name], [clean_name], f"Rename {o.name}") |
| 340 | fixes.node.append(renamer) |
| 341 | |
| 342 | new_output = fixes.output.add() |
| 343 | new_output.name = clean_name |
| 344 | new_output.type.CopyFrom(o.type) |
| 345 | |
| 346 | # merge if we have any renaming to do |
| 347 | if io_map: |
| 348 | graph = onnx.compose.merge_graphs(graph, fixes, io_map) |
| 349 | |
| 350 | return graph |
| 351 | |