microsoft/onnxruntime-extensions
Publicmirrored from https://github.com/microsoft/onnxruntime-extensionsAvailable
onnxruntime_extensions/pnp/_torchext.py
310lines · modecode
| 1 | import onnx |
| 2 | import torch |
| 3 | import numpy as np |
| 4 | from typing import Any |
| 5 | from onnx import helper |
| 6 | from onnx import onnx_pb as onnx_proto |
| 7 | from distutils.version import LooseVersion |
| 8 | from torch.onnx import register_custom_op_symbolic |
| 9 | |
| 10 | from ._utils import ONNXModelUtils |
| 11 | from ._base import CustomFunction, ProcessingTracedModule, is_processing_module |
| 12 | from ._onnx_ops import ox as _ox, schema as _schema |
| 13 | from ._onnx_ops import ONNXElementContainer, make_model_ex |
| 14 | from .._ortapi2 import OrtPyFunction, get_opset_version_from_ort |
| 15 | |
| 16 | |
| 17 | def _is_numpy_object(x): |
| 18 | return isinstance(x, (np.ndarray, np.generic)) |
| 19 | |
| 20 | |
| 21 | def _is_numpy_string_type(arr): |
| 22 | return arr.dtype.kind in {'U', 'S'} |
| 23 | |
| 24 | |
| 25 | def _is_string_type(x): |
| 26 | if isinstance(x, list): |
| 27 | return any(_is_string_type(e) for e in x) |
| 28 | elif isinstance(x, torch.Tensor): |
| 29 | return False |
| 30 | elif not _is_numpy_object(x): |
| 31 | x = np.array(x) |
| 32 | return _is_numpy_string_type(x) |
| 33 | |
| 34 | |
| 35 | def _to_onnx_type(dtype): |
| 36 | ty_dict = {torch.bool: onnx_proto.TensorProto.BOOL, |
| 37 | torch.float32: onnx_proto.TensorProto.FLOAT, |
| 38 | torch.float64: onnx_proto.TensorProto.DOUBLE, |
| 39 | torch.long: onnx_proto.TensorProto.INT64, |
| 40 | torch.int32: onnx_proto.TensorProto.INT32} |
| 41 | # ... |
| 42 | return ty_dict.get(dtype, onnx_proto.TensorProto.STRING) |
| 43 | |
| 44 | |
| 45 | class OnnxOpFunction(CustomFunction): |
| 46 | @classmethod |
| 47 | def get_next_id_name(cls, name_base): |
| 48 | name = 'cls' if name_base is None else name_base |
| 49 | _cid = getattr(cls, '_cid', 1) |
| 50 | cls._cid = _cid + 1 |
| 51 | return "{}_{}".format(name, _cid) |
| 52 | |
| 53 | @staticmethod |
| 54 | def jvp(ctx: Any, *grad_inputs: Any) -> Any: |
| 55 | pass |
| 56 | |
| 57 | @staticmethod |
| 58 | def backward(ctx: Any, *grad_outputs: Any) -> Any: |
| 59 | return grad_outputs |
| 60 | |
| 61 | @classmethod |
| 62 | def build_model(cls, opset_version, *args): |
| 63 | # build the one node graph |
| 64 | if isinstance(args[0], list): |
| 65 | args = [np.asarray(_i) for _i in args] |
| 66 | ec = ONNXElementContainer(get_opset_version_from_ort() if opset_version is None else opset_version) |
| 67 | attrs = cls.attrs |
| 68 | vi_inputs = [helper.make_tensor_value_info( |
| 69 | 'it_' + str(id(_arg)), _to_onnx_type(_arg.dtype), list(_arg.shape)) |
| 70 | for _arg in args] |
| 71 | inputs = [_vi.name for _vi in vi_inputs] |
| 72 | if hasattr(cls.opb_func, 'outputs') and len(cls.opb_func.outputs) > 0: |
| 73 | vi_outputs = [helper.make_tensor_value_info( |
| 74 | cls.get_next_id_name('ot'), *_schm) for _schm in cls.opb_func.outputs] |
| 75 | else: |
| 76 | vi_outputs = [helper.make_tensor_value_info( |
| 77 | cls.get_next_id_name('ot'), onnx_proto.TensorProto.FLOAT, [] |
| 78 | )] |
| 79 | outputs = [_vi.name for _vi in vi_outputs] |
| 80 | # build the node |
| 81 | opfunc = cls.opb_func |
| 82 | opfunc(inputs, outputs, ec, None, **attrs) |
| 83 | g = helper.make_graph(ec.nodes, cls.get_next_id_name('g'), vi_inputs, vi_outputs) |
| 84 | m = make_model_ex(g, ec.node_domain_version_pair_sets, ec.target_opset) |
| 85 | return m |
| 86 | |
| 87 | @classmethod |
| 88 | @torch.jit.unused |
| 89 | def _onnx_call(cls, ctx, *args) -> Any: |
| 90 | m = cls.build_model(None, *args) |
| 91 | try: |
| 92 | f = OrtPyFunction.from_model(m) |
| 93 | result = f(*list(_i.numpy() if isinstance(_i, torch.Tensor) else _i for _i in args)) |
| 94 | except Exception as e: |
| 95 | onnx.save_model(m, '_temp_debugging.onnx') |
| 96 | raise e |
| 97 | |
| 98 | results = result if isinstance(result, tuple) else [result] |
| 99 | return tuple([torch.from_numpy(_o) for _o in results]) if len(results) > 1 else torch.from_numpy(results[0]) |
| 100 | |
| 101 | @classmethod |
| 102 | def forward(cls, ctx: Any, *args: Any, **kwargs: Any) -> Any: |
| 103 | return cls._onnx_call(ctx, *args, **kwargs) |
| 104 | |
| 105 | @classmethod |
| 106 | def symbolic(cls, g, *args): |
| 107 | return g.op(cls.op_type, *args) |
| 108 | |
| 109 | |
| 110 | def create_op_function(op_type: str, func, **attrs): |
| 111 | if _ox.is_raw(func): |
| 112 | func = _schema(func.__func__) |
| 113 | cls = type(_ox.get_unique_operator_type_name(op_type), (OnnxOpFunction,), |
| 114 | dict( |
| 115 | op_type=op_type, |
| 116 | opb_func=func, |
| 117 | attrs=attrs |
| 118 | )) |
| 119 | return cls.apply # noqa |
| 120 | |
| 121 | |
| 122 | onnx_pad = create_op_function('Pad', _ox.pad) |
| 123 | onnx_where = create_op_function('Where', _ox.where) |
| 124 | onnx_greater = create_op_function('Greater', _ox.greater) |
| 125 | |
| 126 | |
| 127 | class _OnnxModelFunction: |
| 128 | id_object_map = {} # cannot use the string directly since jit.script doesn't support the data type |
| 129 | id_function_map = {} |
| 130 | str_model_function_id = '_model_function_id' |
| 131 | str_model_id = '_model_id' |
| 132 | str_model_attached = '_model_attached' |
| 133 | |
| 134 | |
| 135 | @torch.jit.ignore |
| 136 | def _invoke_onnx_model(model_id: int, *args, **kwargs): |
| 137 | func = _OnnxModelFunction.id_function_map.get(model_id, None) |
| 138 | if not func: |
| 139 | model_or_path = _OnnxModelFunction.id_object_map.get(model_id) |
| 140 | if model_or_path is None: |
| 141 | raise ValueError("cannot find id={} registered!".format(model_id)) |
| 142 | func = OrtPyFunction.from_model(model_or_path) |
| 143 | _OnnxModelFunction.id_function_map[model_id] = func |
| 144 | results = func(*list(_i.numpy() if isinstance(_i, torch.Tensor) else _i for _i in args), **kwargs) |
| 145 | return tuple( |
| 146 | [torch.from_numpy(_o) for _o in results]) if isinstance(results, tuple) else torch.from_numpy(results) |
| 147 | |
| 148 | |
| 149 | @torch.jit.ignore |
| 150 | def invoke_onnx_model1(model_id: int, arg0): |
| 151 | return _invoke_onnx_model(model_id, arg0) |
| 152 | |
| 153 | |
| 154 | @torch.jit.ignore |
| 155 | def invoke_onnx_model2(model_id: int, arg0, arg1): |
| 156 | return _invoke_onnx_model(model_id, arg0, arg1) |
| 157 | |
| 158 | |
| 159 | @torch.jit.ignore |
| 160 | def invoke_onnx_model3(model_id: int, arg0, arg1, arg2): |
| 161 | return _invoke_onnx_model(model_id, arg0, arg1, arg2) |
| 162 | |
| 163 | |
| 164 | class _OnnxTracedFunction(CustomFunction): |
| 165 | @classmethod |
| 166 | def forward(cls, ctx: Any, *args: Any, **kwargs: Any) -> Any: |
| 167 | return _invoke_onnx_model(args[0].item(), *args[1:], **kwargs) |
| 168 | |
| 169 | @classmethod |
| 170 | def symbolic(cls, g, *args): |
| 171 | ret = g.op('ai.onnx.contrib::_ModelFunctionCall', *args) |
| 172 | model_id = torch.onnx.symbolic_helper._maybe_get_scalar(args[0]) # noqa |
| 173 | if not model_id: |
| 174 | return ret |
| 175 | |
| 176 | func = _OnnxModelFunction.id_function_map.get(model_id.item(), None) |
| 177 | if not func or len(func.outputs) <= 1: |
| 178 | return ret |
| 179 | |
| 180 | outputs = [ret] |
| 181 | for _ in range(len(func.outputs) - 1): |
| 182 | outputs.append(ret.node().addOutput()) |
| 183 | |
| 184 | return tuple(outputs) |
| 185 | |
| 186 | |
| 187 | def create_model_function(model_or_path): |
| 188 | _id = id(model_or_path) |
| 189 | assert _id != 0, "internal error: the id of a Python object is 0." |
| 190 | _OnnxModelFunction.id_object_map[_id] = model_or_path |
| 191 | return _id |
| 192 | |
| 193 | |
| 194 | def get_id_models(): |
| 195 | return _OnnxModelFunction.id_object_map |
| 196 | |
| 197 | |
| 198 | class OnnxTracedModelFunction: |
| 199 | def __init__(self, onnx_model): |
| 200 | self.func_id = create_model_function(onnx_model) |
| 201 | |
| 202 | def __call__(self, *args, **kwargs): |
| 203 | return _OnnxTracedFunction.apply(torch.tensor(self.func_id), *args, **kwargs) |
| 204 | |
| 205 | |
| 206 | class _OnnxModelModule(torch.nn.Module): |
| 207 | def __init__(self, mdl): |
| 208 | super(_OnnxModelModule, self).__init__() |
| 209 | self.function = OnnxTracedModelFunction(mdl) |
| 210 | |
| 211 | def forward(self, *args): |
| 212 | return self.function(*args) |
| 213 | |
| 214 | |
| 215 | def _symbolic_pythonop(g: torch._C.Graph, n: torch._C.Node, *args, **kwargs): |
| 216 | name = kwargs["name"] |
| 217 | if name.startswith(invoke_onnx_model1.__name__[:-1]): |
| 218 | # NB: if you want to get the value of the first argument, i.e. the model id, |
| 219 | # you can get it by torch.onnx.symbolic_helper._maybe_get_scalar(args[0]).item() |
| 220 | ret = g.op("ai.onnx.contrib::_ModelFunctionCall", *args) |
| 221 | else: |
| 222 | # Logs a warning and returns None |
| 223 | import warnings |
| 224 | return warnings.warn("prim::PythonOp", "unknown node kind: " + name) |
| 225 | # Copy type and shape from original node. |
| 226 | ret.setType(n.output().type()) |
| 227 | return ret |
| 228 | |
| 229 | |
| 230 | if LooseVersion(torch.__version__) >= LooseVersion("1.11"): |
| 231 | register_custom_op_symbolic("prim::PythonOp", _symbolic_pythonop, 1) |
| 232 | |
| 233 | |
| 234 | class SequentialProcessingModule(ProcessingTracedModule): |
| 235 | def __init__(self, *models): |
| 236 | super(SequentialProcessingModule, self).__init__() |
| 237 | self.model_list = torch.nn.ModuleList() |
| 238 | for mdl_ in models: |
| 239 | if isinstance(mdl_, onnx.ModelProto): |
| 240 | self.model_list.append(_OnnxModelModule(mdl_)) |
| 241 | elif is_processing_module(mdl_): |
| 242 | self.model_list.append(mdl_) |
| 243 | else: |
| 244 | assert callable(mdl_), "the model type is not recognizable." |
| 245 | self.model_list.append(ProcessingTracedModule(mdl_)) |
| 246 | |
| 247 | def forward(self, *args): |
| 248 | outputs = args |
| 249 | with torch.no_grad(): |
| 250 | for idx_, mdl_ in enumerate(self.model_list): |
| 251 | if not isinstance(outputs, tuple): |
| 252 | outputs = (outputs,) |
| 253 | outputs = mdl_(*outputs) |
| 254 | |
| 255 | return outputs |
| 256 | |
| 257 | def export(self, *args, **kwargs): |
| 258 | prefix_m = None |
| 259 | core_m = self |
| 260 | raw_input_flag = any(_is_string_type(x_) for x_ in args) |
| 261 | if raw_input_flag: |
| 262 | # NB: torch.onnx.export doesn't support exporting a module accepting string type input, |
| 263 | # So, in this case, the module will be separated into two parts to use the customized export. |
| 264 | m0 = self.model_list[0] |
| 265 | new_args = m0(*args) |
| 266 | if not isinstance(new_args, tuple): |
| 267 | new_args = (new_args, ) |
| 268 | prefix_m = m0.export(*args, **kwargs) |
| 269 | args = new_args |
| 270 | core_m = SequentialProcessingModule(*self.model_list[1:]) |
| 271 | if prefix_m is None: |
| 272 | return super().export(*args, **kwargs) |
| 273 | else: |
| 274 | oxml = core_m.export(*args, **kwargs) |
| 275 | model = ONNXModelUtils.join_models(prefix_m, oxml) |
| 276 | |
| 277 | # Rename the input/output node names if the user has provided any substitutions! |
| 278 | # Ref: https://github.com/onnx/onnx/issues/2052 |
| 279 | # Known issue: This logic doesn't deal with subgraphs. |
| 280 | if (('input_names' in kwargs) or ('output_names' in kwargs)) and \ |
| 281 | (kwargs['input_names'] or kwargs['output_names']): |
| 282 | swaps = {} |
| 283 | if 'input_names' in kwargs and kwargs['input_names']: |
| 284 | assert len(model.graph.input) == len(kwargs['input_names']), \ |
| 285 | "Expecting {} input names but got {}".format( |
| 286 | len(model.graph.input), len(kwargs['input_names'])) |
| 287 | for n, new_name in zip(model.graph.input, kwargs['input_names']): |
| 288 | swaps[n.name] = new_name |
| 289 | n.name = new_name |
| 290 | |
| 291 | if 'output_names' in kwargs and kwargs['output_names']: |
| 292 | assert len(model.graph.output) == len(kwargs['output_names']), \ |
| 293 | "Expecting {} output names but got {}".format( |
| 294 | len(model.graph.output), len(kwargs['output_names'])) |
| 295 | for n, new_name in zip(model.graph.output, kwargs['output_names']): |
| 296 | swaps[n.name] = new_name |
| 297 | n.name = new_name |
| 298 | |
| 299 | if swaps: |
| 300 | for n in model.graph.node: |
| 301 | for j in range(len(n.input)): |
| 302 | n.input[j] = swaps.get(n.input[j], n.input[j]) |
| 303 | |
| 304 | for j in range(len(n.output)): |
| 305 | n.output[j] = swaps.get(n.output[j], n.output[j]) |
| 306 | |
| 307 | for n in model.graph.initializer: |
| 308 | n.name = swaps.get(n.name, n.name) |
| 309 | |
| 310 | return model |
| 311 | |