microsoft/onnxruntime-extensions
Publicmirrored from https://github.com/microsoft/onnxruntime-extensionsAvailable
onnxruntime_extensions/onnxprocess/_tensor.py
624lines · modecode
| 1 | import torch |
| 2 | import builtins |
| 3 | import functools |
| 4 | import numpy as np |
| 5 | from onnx import onnx_pb as onnx_proto |
| 6 | from typing import List, Tuple, Optional, Union, Any, ContextManager, overload, Iterator, NamedTuple |
| 7 | from torch.types import _int, _float, _bool, Number, _dtype, _device, _qscheme, _size, _layout # noqa |
| 8 | from torch import strided, memory_format, contiguous_format, StringType # noqa |
| 9 | |
| 10 | from ._onnx_ops import ox as _ox |
| 11 | from ..eager_op import EagerOp |
| 12 | |
| 13 | |
| 14 | class _EagerTensor: |
| 15 | def __init__(self, _t, name=None, sess=None, raw_data: Any = None): |
| 16 | self._t = _t if isinstance(_t, torch.Tensor) else torch.tensor(_t) |
| 17 | if isinstance(name, (tuple, list)): |
| 18 | assert len(name) == 1, "Multiple names for one tensor!" |
| 19 | name = name[0] |
| 20 | self.name = '' if name is None else name |
| 21 | self.raw_data = raw_data |
| 22 | self.symbolic_shape = [] |
| 23 | |
| 24 | def __repr__(self): |
| 25 | if self.raw_data is not None: |
| 26 | return "name: {}, \"{}\"".format(self.name, str(self.raw_data)) |
| 27 | else: |
| 28 | return "name: {}, {}, dtype={}".format(self.name, repr(self._t), str(self._t.dtype)) |
| 29 | |
| 30 | _all_ops = {} |
| 31 | |
| 32 | @property |
| 33 | def value(self) -> Union[torch.Tensor, Any]: |
| 34 | return self.raw_data if self.raw_data else self._t |
| 35 | |
| 36 | @property |
| 37 | def t(self): |
| 38 | return self._t |
| 39 | |
| 40 | @property |
| 41 | def onnx_type(self): |
| 42 | return self.to_onnx_type(self._t.dtype) |
| 43 | |
| 44 | @classmethod |
| 45 | def is_numeric(cls, np_arr): |
| 46 | return np_arr.dtype.kind in set('buifc') |
| 47 | |
| 48 | @classmethod |
| 49 | def set_active_session(cls, sess): |
| 50 | """ |
| 51 | set the active operator tracing log session. if sess is None, the active session will be removed |
| 52 | :param sess: |
| 53 | :return: |
| 54 | """ |
| 55 | if not hasattr(cls, '_active_session'): |
| 56 | cls._active_session = sess |
| 57 | if sess is None: |
| 58 | raise RuntimeError("unset the active session twice!") |
| 59 | else: |
| 60 | if sess is not None: |
| 61 | raise RuntimeError("The active session already assigned!") |
| 62 | delattr(cls, '_active_session') |
| 63 | |
| 64 | @classmethod |
| 65 | def get_trace_session(cls): |
| 66 | if not hasattr(cls, '_active_session'): |
| 67 | raise RuntimeError("the tracing not started yet!") |
| 68 | return cls._active_session # noqa |
| 69 | |
| 70 | @classmethod |
| 71 | def get_container(cls): |
| 72 | return cls.get_trace_session().container |
| 73 | |
| 74 | @classmethod |
| 75 | def from_onnx(cls, raw_val, ort_sess, name): |
| 76 | raw_data = None |
| 77 | if cls.is_numeric(raw_val): |
| 78 | val = torch.from_numpy(raw_val) |
| 79 | else: |
| 80 | # only keep the shape and the value was stored by it-self. |
| 81 | val = torch.empty(*raw_val.shape, dtype=torch.uint8) |
| 82 | raw_data = raw_val |
| 83 | t = cls(val, name, ort_sess, raw_data) |
| 84 | return t |
| 85 | |
| 86 | @classmethod |
| 87 | def from_torch(cls, _t, name): |
| 88 | t_name = name if name is not None else "id_{}".format(id(_t)) |
| 89 | ts = cls(_t, t_name) |
| 90 | return ts |
| 91 | |
| 92 | @classmethod |
| 93 | # torch.tensor prototype |
| 94 | def mytensor(cls, data: Any, dtype: Optional[_dtype] = None, device: Union[_device, str, None] = None, requires_grad: _bool = False): # noqa |
| 95 | y = torch.tensor(data, dtype=dtype, device=device, requires_grad=requires_grad) |
| 96 | val = _ox.make_tensor(cls.to_onnx_type(y.dtype), list(y.size()), |
| 97 | [data] if isinstance(data, (int, float, str, bool)) else data) |
| 98 | s = _ox.constant([], [_ox.get_unique_tensor_name('const')], cls.get_container(), None, value=val) |
| 99 | return cls.from_torch(y, s) |
| 100 | |
| 101 | def numpy(self): |
| 102 | return self._t.numpy() if self.raw_data is None else self.raw_data |
| 103 | |
| 104 | def item(self): |
| 105 | return self.numpy().item() |
| 106 | |
| 107 | def get_shape(self): |
| 108 | return self.t.size() if len(self.symbolic_shape) == 0 else self.symbolic_shape |
| 109 | |
| 110 | def _to_binary_tensor_args(self, other): |
| 111 | # convert self, other to [self, other], but if either is a number, convert that to a constant |
| 112 | x, y = self, other |
| 113 | if isinstance(y, (int, float, bool, np.ndarray)): |
| 114 | y = self.mytensor(y) |
| 115 | elif isinstance(x, (int, float, bool, np.ndarray)): |
| 116 | x = self.mytensor(x) |
| 117 | return x, y |
| 118 | |
| 119 | _dup_id = 0 |
| 120 | |
| 121 | def __copy__(self): |
| 122 | new_t = _EagerTensor.from_torch(self.t, self.name + '_{}'.format(_EagerTensor._dup_id)) |
| 123 | self._dup_id += 1 |
| 124 | new_t.raw_data = self.raw_data |
| 125 | return new_t |
| 126 | |
| 127 | def __add__(self, other): |
| 128 | x0, x1 = self._to_binary_tensor_args(other) |
| 129 | y = torch.add(x0._t, x1._t) |
| 130 | s = _ox.add(*_EagerTensor.ox_args([x0, x1])) |
| 131 | return self.from_torch(y, s) |
| 132 | |
| 133 | def __sub__(self, other): |
| 134 | x0, x1 = self._to_binary_tensor_args(other) |
| 135 | y = torch.sub(x0._t, x1._t) |
| 136 | s = _ox.sub(*_EagerTensor.ox_args([x0, x1])) |
| 137 | return self.from_torch(y, s) |
| 138 | |
| 139 | def __mul__(self, other): |
| 140 | x0, x1 = self._to_binary_tensor_args(other) |
| 141 | y = torch.mul(x0._t, x1._t) |
| 142 | s = _ox.mul(*_EagerTensor.ox_args([x0, x1])) |
| 143 | return self.from_torch(y, s) |
| 144 | |
| 145 | def __div__(self, other): |
| 146 | x0, x1 = self._to_binary_tensor_args(other) |
| 147 | y = torch.div(x0._t, x1._t) |
| 148 | s = _ox.div(*_EagerTensor.ox_args([x0, x1])) |
| 149 | return self.from_torch(y, s) |
| 150 | |
| 151 | def __pow__(self, other): |
| 152 | x0, x1 = self._to_binary_tensor_args(other) |
| 153 | y = torch.pow(x0._t, x1._t) |
| 154 | s = _ox.pow(*_EagerTensor.ox_args([x0, x1])) |
| 155 | return self.from_torch(y, s) |
| 156 | |
| 157 | def __matmul__(self, other): |
| 158 | x0, x1 = self._to_binary_tensor_args(other) |
| 159 | y = torch.matmul(x0._t, x1._t) |
| 160 | s = _ox.matmul(*_EagerTensor.ox_args([x0, x1])) |
| 161 | return self.from_torch(y, s) |
| 162 | |
| 163 | def __lt__(self, other): |
| 164 | x0, x1 = self._to_binary_tensor_args(other) |
| 165 | y = torch.less(x0._t, x1._t) |
| 166 | s = _ox.less(*_EagerTensor.ox_args([x0, x1])) |
| 167 | return self.from_torch(y, s) |
| 168 | |
| 169 | def __le__(self, other): |
| 170 | x0, x1 = self._to_binary_tensor_args(other) |
| 171 | y = torch.less_equal(x0._t, x1._t) |
| 172 | s = _ox.less_equal(*_EagerTensor.ox_args([x0, x1])) |
| 173 | return self.from_torch(y, s) |
| 174 | |
| 175 | def __eq__(self, other): |
| 176 | x0, x1 = self._to_binary_tensor_args(other) |
| 177 | y = torch.equal(x0._t, x1._t) |
| 178 | s = _ox.equal(*_EagerTensor.ox_args([x0, x1])) |
| 179 | return self.from_torch(y, s) |
| 180 | |
| 181 | def __ne__(self, other): |
| 182 | x0, x1 = self._to_binary_tensor_args(other) |
| 183 | y = torch.not_equal(x0._t, x1._t) |
| 184 | s = _ox.not_equal(*_EagerTensor.ox_args([x0, x1])) |
| 185 | return self.from_torch(y, s) |
| 186 | |
| 187 | def __gt__(self, other): |
| 188 | x0, x1 = self._to_binary_tensor_args(other) |
| 189 | y = torch.greater(x0._t, x1._t) |
| 190 | s = _ox.greater(*_EagerTensor.ox_args([x0, x1])) |
| 191 | return self.from_torch(y, s) |
| 192 | |
| 193 | def __ge__(self, other): |
| 194 | x0, x1 = self._to_binary_tensor_args(other) |
| 195 | y = torch.greater_equal(x0._t, x1._t) |
| 196 | s = _ox.greater_equal(*_EagerTensor.ox_args([x0, x1])) |
| 197 | return self.from_torch(y, s) |
| 198 | |
| 199 | def __invert__(self): |
| 200 | if self.t.dtype is torch.bool: |
| 201 | y = torch.logical_not(self.t) |
| 202 | s = _ox.not_op(*self.my_args()) |
| 203 | return self.from_torch(y, s) |
| 204 | else: |
| 205 | raise NotImplementedError("no numeric tensor inverse supported yet.") |
| 206 | |
| 207 | def __neg__(self): |
| 208 | y = torch.neg([self.t]) |
| 209 | s = _ox.neg(*self.my_args()) |
| 210 | return self.from_torch(y, s) |
| 211 | |
| 212 | def __not__(self): |
| 213 | y = torch.logical_not(self.t) |
| 214 | s = _ox.not_op(*self.my_args()) |
| 215 | return self.from_torch(y, s) |
| 216 | |
| 217 | def __or__(self, other): |
| 218 | x0, x1 = self._to_binary_tensor_args(other) |
| 219 | y = torch.logical_or(x0._t, x1._t) |
| 220 | s = _ox.or_op(*_EagerTensor.ox_args([x0, x1])) |
| 221 | return self.from_torch(y, s) |
| 222 | |
| 223 | def __getitem__(self, indices): |
| 224 | y = self.value.__getitem__(indices) |
| 225 | |
| 226 | # normalize indices to tuples of slices |
| 227 | # Formats encountered: |
| 228 | # - a single int |
| 229 | # - a tuple of (int or slice) |
| 230 | if not isinstance(indices, (tuple, list)): # single item: make it a tuple |
| 231 | indices = (indices,) |
| 232 | squeeze = [axis for axis, index in enumerate(indices) if |
| 233 | isinstance(index, int)] # which axes had a single index? |
| 234 | indices = tuple( |
| 235 | index if isinstance(index, slice) else slice(index, index + 1 if index != -1 else None, 1) for index in |
| 236 | indices) # make all tuple items of type Slice |
| 237 | bs, es, ss, ds = [], [], [], [] |
| 238 | INT_MAX = 2 ** 63 - 1 |
| 239 | for axis, index in enumerate(indices): |
| 240 | if not isinstance(index, slice): |
| 241 | raise ValueError("Index expected") |
| 242 | if index.start is None and index.stop is None: # [:] can be skipped |
| 243 | continue |
| 244 | b, e, s = index.start, index.stop, index.step |
| 245 | bs.append(b if b is not None else 0) |
| 246 | es.append(e if e is not None else INT_MAX) |
| 247 | ss.append(s if s is not None else 1) |
| 248 | ds.append(axis) |
| 249 | s = _ox.slice(*self.my_args(), starts=bs, ends=es, axes=ds, steps=ss) |
| 250 | if squeeze: # single index means we must drop the axis |
| 251 | s = _ox.squeeze(*self.ox_name_args(s), axes=squeeze) |
| 252 | |
| 253 | return self.from_torch(y, s) |
| 254 | |
| 255 | def __getattribute__(self, attr): |
| 256 | """ |
| 257 | A little hack that allows to call unary operators in a chaining fashion, |
| 258 | e.g. x.shape() instead of ox.shape(x). |
| 259 | """ |
| 260 | if attr in _EagerTensor._all_ops: |
| 261 | f = _EagerTensor._all_ops[attr] |
| 262 | return functools.partial(f, self) |
| 263 | else: |
| 264 | return object.__getattribute__(self, attr) |
| 265 | |
| 266 | @classmethod |
| 267 | def ox_name_args(cls, input_names, output_names=None): |
| 268 | """ |
| 269 | generate the arguments for ONNX model builder. |
| 270 | :param input_names: input name list |
| 271 | :param output_names: output name list, can be None, or [None]*output_n |
| 272 | :return: input_names, output_names, container, operator_name |
| 273 | """ |
| 274 | container = cls.get_trace_session().container |
| 275 | if output_names is None: |
| 276 | output_names = [None] # by default, there is only one output |
| 277 | |
| 278 | output_names = [_ox.get_unique_tensor_name(str(n_)) |
| 279 | if output_names[n_] is None else |
| 280 | output_names[n_] for n_ in range(len(output_names))] |
| 281 | operator_name = None |
| 282 | return input_names, output_names, container, operator_name |
| 283 | |
| 284 | @classmethod |
| 285 | def ort_verify(cls, ts_from, ts_to): |
| 286 | result, model = cls.get_trace_session().runops(ts_from, ts_to) |
| 287 | for idx in range(len(ts_to)): |
| 288 | if not np.allclose(ts_to[idx].numpy(), result[idx]): |
| 289 | # ONNX cannot be import globally, which is conflict with torch.onnx |
| 290 | import onnx # noqa |
| 291 | onnx.save_model(model, 'mt_debmodel.onnx') |
| 292 | raise RuntimeError("ONNXRuntime Result is not same pytorch!") |
| 293 | |
| 294 | def create_and_verify(self, value, name, additional_inputs=None): |
| 295 | ts_y = self.from_torch(value, name) |
| 296 | inputs = [self] + ([] if additional_inputs is None else additional_inputs) |
| 297 | self.ort_verify(inputs, [ts_y]) |
| 298 | return ts_y |
| 299 | |
| 300 | @classmethod |
| 301 | def ox_args(cls, tensors, output_names=None): |
| 302 | input_names = [ts_ if isinstance(ts_, str) else ts_.name for ts_ in tensors] |
| 303 | return cls.ox_name_args(input_names, output_names) |
| 304 | |
| 305 | def my_args(self): |
| 306 | return self.ox_args([self]) |
| 307 | |
| 308 | @staticmethod |
| 309 | def normalize_seq(list_or_tuple): |
| 310 | return [x.value.item() if isinstance(x, _EagerTensor) else x for x in list_or_tuple] |
| 311 | |
| 312 | @staticmethod |
| 313 | def to_onnx_type(torch_type): |
| 314 | ty_dict = {torch.bool: onnx_proto.TensorProto.BOOL, |
| 315 | torch.float32: onnx_proto.TensorProto.FLOAT, |
| 316 | torch.long: onnx_proto.TensorProto.INT64, |
| 317 | torch.int32: onnx_proto.TensorProto.INT32} |
| 318 | # ... |
| 319 | return ty_dict.get(torch_type, onnx_proto.TensorProto.STRING) |
| 320 | |
| 321 | def long(self): |
| 322 | y = self._t.long() |
| 323 | s = _ox.cast(*self.my_args(), to=onnx_proto.TensorProto.INT64) |
| 324 | return self.create_and_verify(y, s[0]) |
| 325 | |
| 326 | def cumsum(self, dim: _int, *, dtype: Optional[_dtype] = None): # noqa |
| 327 | y = self._t.cumsum(dim, dtype=dtype) |
| 328 | s = _ox.cumsum(*self.my_args(), axis=dim) |
| 329 | return self.create_and_verify(y, s[0]) |
| 330 | |
| 331 | def size(self): |
| 332 | y = self._t.size() |
| 333 | s = _ox.shape(*self.my_args()) |
| 334 | return self.create_and_verify(y, s[0]) |
| 335 | |
| 336 | def type(self, dtype: Union[str, _dtype], non_blocking: _bool=False): |
| 337 | y = self._t.type(dtype, non_blocking) |
| 338 | s = _ox.cast(*self.my_args(), to=self.to_onnx_type(dtype)) |
| 339 | return self.create_and_verify(y, s) |
| 340 | |
| 341 | def to(self, device): |
| 342 | y = self._t.to(device) |
| 343 | s = _ox.identity(*self.my_args()) |
| 344 | return self.create_and_verify(y, s[0]) |
| 345 | |
| 346 | def cpu(self): |
| 347 | y = self._t.cpu() |
| 348 | s = _ox.identity(*self.my_args()) |
| 349 | return self.create_and_verify(y, s[0]) |
| 350 | |
| 351 | def detach(self): |
| 352 | y = self._t.detach() |
| 353 | s = _ox.identity(*self.my_args()) |
| 354 | return self.create_and_verify(y, s[0]) |
| 355 | |
| 356 | def clone(self): |
| 357 | y = self._t.clone() |
| 358 | s = _ox.identity(*self.my_args()) |
| 359 | return self.create_and_verify(y, s[0]) |
| 360 | |
| 361 | def masked_fill(self, mask, value): |
| 362 | y = self._t.masked_fill(mask.value, value) |
| 363 | if not isinstance(value, _EagerTensor): |
| 364 | value = _EagerTensor.mytensor(value) |
| 365 | s = _ox.where(*_EagerTensor.ox_args([mask, value, self])) |
| 366 | return self.create_and_verify(y, s[0], additional_inputs=[mask, value]) |
| 367 | |
| 368 | def unsqueeze(self, dim: _int): |
| 369 | y = self._t.unsqueeze(dim) |
| 370 | s = _ox.unsqueeze(*self.my_args(), [dim]) |
| 371 | return self.create_and_verify(y, s[0]) |
| 372 | |
| 373 | def squeeze(self, dim: _int): |
| 374 | y = self._t.squeeze(dim) |
| 375 | s = _ox.squeeze(*self.my_args(), [dim]) |
| 376 | return self.create_and_verify(y, s[0]) |
| 377 | |
| 378 | |
| 379 | def _create_ox_sequence(*size): |
| 380 | container = _EagerTensor.get_container() |
| 381 | con_x = [] |
| 382 | if builtins.any(isinstance(n_, _EagerTensor) for n_ in size): |
| 383 | for x in size: |
| 384 | if isinstance(x, _EagerTensor): |
| 385 | x_h = _ox.unsqueeze(*_EagerTensor.ox_args([x]))[0] |
| 386 | else: |
| 387 | x_c = _ox.make_tensor(onnx_proto.TensorProto.INT64, [1], [x]) |
| 388 | x_h = _ox.constant([], [_ox.get_unique_tensor_name('const')], container, None, value=x_c)[0] |
| 389 | con_x.append(x_h) |
| 390 | return _ox.concat(con_x, [_ox.get_unique_tensor_name('concat')], container, None) |
| 391 | else: |
| 392 | ts_size = _ox.make_tensor(onnx_proto.TensorProto.INT64, [len(size)], size) |
| 393 | return _ox.constant([], [_ox.get_unique_tensor_name('const')], container, None, value=ts_size) |
| 394 | |
| 395 | |
| 396 | def _create_ox_sequence_constant(*size, init_value=None, onnx_type=None): |
| 397 | if onnx_type is None: |
| 398 | onnx_type = onnx_proto.TensorProto.FLOAT |
| 399 | names = _create_ox_sequence(*size) |
| 400 | ts_val = _ox.make_tensor(onnx_type, [1], [init_value]) |
| 401 | |
| 402 | container = _EagerTensor.get_container() |
| 403 | s = _ox.constant_of_shape(names, [_ox.get_unique_tensor_name('cos')], container, None, value=ts_val) |
| 404 | return s[0] |
| 405 | |
| 406 | |
| 407 | def empty(*size: Union[_int, _EagerTensor], memory_format: Optional[memory_format] = None, out: Optional[_EagerTensor] = None, |
| 408 | dtype: _dtype = None, layout: _layout = strided, device: Union[_device, str, None] = None, |
| 409 | requires_grad: _bool = False) -> _EagerTensor: # noqa |
| 410 | |
| 411 | if len(size) == 1 and isinstance(size[0], list): |
| 412 | size = size[0] |
| 413 | n_size = _EagerTensor.normalize_seq(size) |
| 414 | y = torch.empty(*n_size, memory_format=memory_format, out=out, |
| 415 | dtype=dtype, layout=layout, device=device, requires_grad=requires_grad) |
| 416 | s = _create_ox_sequence_constant(*size, init_value=0., onnx_type=_EagerTensor.to_onnx_type(y.dtype)) |
| 417 | return _EagerTensor.from_torch(y, s) |
| 418 | |
| 419 | |
| 420 | def zeros(*size: Union[_int, _EagerTensor], out: Optional[_EagerTensor] = None, dtype: _dtype = None, layout: _layout = strided, |
| 421 | device: Union[_device, str, None] = None, requires_grad: _bool = False) -> _EagerTensor: # noqa |
| 422 | |
| 423 | if len(size) == 1 and isinstance(size[0], list): |
| 424 | size = size[0] |
| 425 | n_size = _EagerTensor.normalize_seq(size) |
| 426 | y = torch.zeros(*n_size, out=out, dtype=dtype, |
| 427 | layout=layout, device=device, requires_grad=requires_grad) |
| 428 | s = _create_ox_sequence_constant(*size, init_value=0, onnx_type=_EagerTensor.to_onnx_type(y.dtype)) |
| 429 | return _EagerTensor.from_torch(y, s) |
| 430 | |
| 431 | |
| 432 | def ones(*size: Union[_int, _EagerTensor], out: Optional[_EagerTensor] = None, dtype: _dtype = None, layout: _layout = strided, |
| 433 | device: Union[_device, str, None] = None, requires_grad: _bool = False) -> _EagerTensor: # noqa |
| 434 | |
| 435 | if len(size) == 1 and isinstance(size[0], list): |
| 436 | size = size[0] |
| 437 | n_size = _EagerTensor.normalize_seq(size) |
| 438 | y = torch.ones(*n_size, out=out, dtype=dtype, |
| 439 | layout=layout, device=device, requires_grad=requires_grad) |
| 440 | s = _create_ox_sequence_constant(*size, init_value=1, onnx_type=_EagerTensor.to_onnx_type(y.dtype)) |
| 441 | return _EagerTensor.from_torch(y, s) |
| 442 | |
| 443 | |
| 444 | def repeat(input_ts: _EagerTensor, *repeats: Union[_int, _EagerTensor]) -> _EagerTensor: # noqa |
| 445 | |
| 446 | if len(repeats) == 1 and isinstance(repeats[0], list): |
| 447 | repeats = repeats[0] |
| 448 | n_size = _EagerTensor.normalize_seq(repeats) |
| 449 | y = input_ts.t.repeat(*n_size) |
| 450 | seq = _create_ox_sequence(*repeats) |
| 451 | s = _ox.tile(*input_ts.my_args(), repeats=seq[0]) |
| 452 | return _EagerTensor.from_torch(y, s[0]) |
| 453 | |
| 454 | |
| 455 | def argmax(input_ts: _EagerTensor, dim: Optional[_int] = None, keepdim: _bool = False) -> _EagerTensor: # noqa |
| 456 | y = torch.argmax(input_ts.value, dim, keepdim) |
| 457 | s = _ox.argmax(*input_ts.my_args(), axis=dim, keepdims=keepdim) |
| 458 | return _EagerTensor.from_torch(y, s) |
| 459 | |
| 460 | |
| 461 | def softmax(input_ts: _EagerTensor, dim: _int, dtype: Optional[_dtype]=None) -> _EagerTensor: |
| 462 | y = torch.softmax(input_ts.value, dim, dtype) |
| 463 | s = _ox.softmax(*input_ts.my_args(), axis=dim) |
| 464 | return _EagerTensor.from_torch(y, s) |
| 465 | |
| 466 | |
| 467 | def cat(tensors: Union[Tuple[_EagerTensor, ...], List[_EagerTensor]], |
| 468 | dim, *, out: Optional[_EagerTensor] = None) -> _EagerTensor: # noqa |
| 469 | res = torch.cat([t_.value for t_ in tensors], dim, out=out) |
| 470 | oname = _ox.concat(*_EagerTensor.ox_args(tensors), dim) |
| 471 | y = _EagerTensor.from_torch(res, oname[0]) |
| 472 | _EagerTensor.ort_verify(tensors, [y]) |
| 473 | return y |
| 474 | |
| 475 | |
| 476 | def all(input_ts: _EagerTensor, out: Optional[_EagerTensor]=None) -> _EagerTensor: # noqa |
| 477 | container = _EagerTensor.get_container() |
| 478 | y = torch.all(input_ts.value) |
| 479 | s_casted = _ox.cast(*input_ts.my_args(), to=onnx_proto.TensorProto.INT64) |
| 480 | s_redm = _ox.reducemin(s_casted, [_ox.get_unique_tensor_name('reducemin')], container, None, axes=[-1]) |
| 481 | s0 = _ox.constant([], [_ox.get_unique_tensor_name('const')], |
| 482 | container, None, value=_ox.make_tensor(onnx_proto.TensorProto.INT64, [1], [0])) |
| 483 | s = _ox.greater(s_redm + s0, [_ox.get_unique_tensor_name('greater')], container, None) |
| 484 | return input_ts.create_and_verify(y, s[0]) |
| 485 | |
| 486 | |
| 487 | def any(input_ts: _EagerTensor, out: Optional[_EagerTensor]=None) -> _EagerTensor: # noqa |
| 488 | container = _EagerTensor.get_container() |
| 489 | y = torch.any(input_ts.value) |
| 490 | s_casted = _ox.cast(*input_ts.my_args(), to=onnx_proto.TensorProto.INT64) |
| 491 | s_redm = _ox.reducesum(s_casted, [_ox.get_unique_tensor_name('reducesum')], container, None, axes=[-1]) |
| 492 | s0 = _ox.constant([], [_ox.get_unique_tensor_name('const')], |
| 493 | container, None, value=_ox.make_tensor(onnx_proto.TensorProto.INT64, [1], [0])) |
| 494 | s = _ox.greater(s_redm + s0, [_ox.get_unique_tensor_name('greater')], container, None) |
| 495 | return input_ts.create_and_verify(y, s[0]) |
| 496 | |
| 497 | |
| 498 | def reshape(input_ts: _EagerTensor, shape: _size): |
| 499 | y = input_ts.t.reshape(shape) |
| 500 | s = _ox.reshape(*input_ts.my_args(), desired_shape=shape) |
| 501 | return input_ts.create_and_verify(y, s[0]) |
| 502 | |
| 503 | |
| 504 | def transpose(input_ts: _EagerTensor, dim0: _int, dim1: _int): |
| 505 | y = input_ts.t.transpose(dim0, dim1) |
| 506 | axes = list(range(y.dim())) |
| 507 | axes[dim0], axes[dim1] = axes[dim1], axes[dim0] |
| 508 | s = _ox.transpose(*input_ts.my_args(), perm=axes) |
| 509 | return input_ts.create_and_verify(y, s[0]) |
| 510 | |
| 511 | |
| 512 | class _LoopIterator: |
| 513 | def __init__(self, ctx): |
| 514 | self.context = ctx |
| 515 | |
| 516 | def __iter__(self): |
| 517 | return self |
| 518 | |
| 519 | def __next__(self): |
| 520 | if self.context.is_stopped(): |
| 521 | _EagerTensor.get_trace_session().pop_container() |
| 522 | raise StopIteration |
| 523 | return self.context.current() |
| 524 | |
| 525 | |
| 526 | class _ControlFlowContext: |
| 527 | def __init__(self): |
| 528 | self.condition_i = None |
| 529 | self.condition = None |
| 530 | self.loop_count = None |
| 531 | self.iteration_num = None |
| 532 | self.states_i = [] |
| 533 | self.loop_states = [] |
| 534 | self.scan_outputs = [] |
| 535 | self.sub_graph = None |
| 536 | |
| 537 | def flow_output(self, cond, *outputs): |
| 538 | assert len(outputs) >= len(self.loop_states), "The loop body doesn't return enough objects" |
| 539 | if self.sub_graph is None: |
| 540 | trc = _EagerTensor.get_trace_session() |
| 541 | self.sub_graph = trc.build_graph(trc.container, |
| 542 | [self.iteration_num, self.condition] + self.loop_states, |
| 543 | [cond] + list(outputs)) |
| 544 | |
| 545 | self.condition = cond |
| 546 | c_state = len(self.loop_states) |
| 547 | self.loop_states = list(outputs[:c_state]) |
| 548 | if len(self.scan_outputs) == 0: |
| 549 | sc = [_EagerTensor(torch.unsqueeze(sci_.value, 0), 'sc_' + sci_.name) for sci_ in outputs[c_state:]] |
| 550 | self.scan_outputs = sc |
| 551 | else: |
| 552 | next_extra_vars = [] |
| 553 | for idx_, ext_ in enumerate(outputs[c_state:]): |
| 554 | et = self.scan_outputs[idx_] |
| 555 | next_extra_vars.append(_EagerTensor( |
| 556 | torch.cat([et.value, torch.unsqueeze(outputs[c_state + idx_].value, 0)]), name=et.name)) |
| 557 | self.scan_outputs = next_extra_vars |
| 558 | self.iteration_num.value.add_(1) |
| 559 | |
| 560 | def current(self): |
| 561 | return [self.iteration_num] + list(self.loop_states) |
| 562 | |
| 563 | def finalize(self): |
| 564 | # generate the outputs from the enclosing scope variables |
| 565 | full_outputs = [_EagerTensor(o_.value, 'lp_' + o_.name) for o_ in self.loop_states + self.scan_outputs] |
| 566 | _ox.loop(*_EagerTensor.ox_args( |
| 567 | [self.loop_count, self.condition_i] + list(self.states_i), |
| 568 | [ts_.name for ts_ in full_outputs]), body=self.sub_graph) |
| 569 | return tuple(full_outputs) |
| 570 | |
| 571 | def is_stopped(self): |
| 572 | return self.condition.item() is False or self.iteration_num.item() >= self.loop_count.item() |
| 573 | |
| 574 | def loop(self, loop_c, condition, *states): |
| 575 | self.condition = condition |
| 576 | self.condition_i = condition |
| 577 | self.states_i = states |
| 578 | _EagerTensor.get_trace_session().stack_container() |
| 579 | self.iteration_num = _EagerTensor.mytensor(0) |
| 580 | # clone the variables for the sub graph. |
| 581 | self.loop_states = [_EagerTensor(st_.value, st_.name) for st_ in states] |
| 582 | self.loop_count = loop_c |
| 583 | loop_b = _LoopIterator(self) |
| 584 | return iter(loop_b) |
| 585 | |
| 586 | |
| 587 | def control_flow(): |
| 588 | return _ControlFlowContext() |
| 589 | |
| 590 | |
| 591 | class _TracingEagerOp(EagerOp): |
| 592 | def __call__(self, *args, **kwargs): |
| 593 | np_args = [ts_.numpy() if isinstance(ts_, _EagerTensor) else ts_ for ts_ in args] |
| 594 | outseq = super().__call__(*np_args, **kwargs) |
| 595 | outseq = outseq if isinstance(outseq, (list, tuple)) else [outseq] |
| 596 | |
| 597 | outputs = [_EagerTensor.from_onnx(outseq[n_], self.ort_session, out_.name) |
| 598 | for n_, out_ in enumerate(self.ort_session.get_outputs())] |
| 599 | |
| 600 | y_names = [y.name for y in outputs] |
| 601 | _ox.model_call(*_EagerTensor.ox_args(args, output_names=y_names), oxml=self.onnx_model) |
| 602 | return tuple(outputs) if len(outputs) > 1 else outputs[0] |
| 603 | |
| 604 | |
| 605 | def op_from_customop(op_type, *args, **kwargs) -> _TracingEagerOp: |
| 606 | return _TracingEagerOp.from_customop(op_type, *args, **kwargs) |
| 607 | |
| 608 | |
| 609 | def op_from_model(path_or_model, *args, **kwargs) -> _TracingEagerOp: |
| 610 | return _TracingEagerOp.from_model(path_or_model, *args, **kwargs) |
| 611 | |
| 612 | |
| 613 | _EagerTensor._all_ops = {'argmax': argmax, |
| 614 | 'softmax': softmax, |
| 615 | 'reshape': reshape, |
| 616 | 'transpose': transpose, |
| 617 | 'repeat': repeat, |
| 618 | 'any': any, |
| 619 | 'all': all} |
| 620 | |
| 621 | tensor = _EagerTensor.mytensor |
| 622 | tensor_from_onnx = _EagerTensor.from_onnx |
| 623 | tensor_from_torch = _EagerTensor.from_torch |
| 624 | tensor_set_session = _EagerTensor.set_active_session |
| 625 | |