onnxruntime_inference_collection.py 58 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071727374757677787980818283848586878889909192939495969798991001011021031041051061071081091101111121131141151161171181191201211221231241251261271281291301311321331341351361371381391401411421431441451461471481491501511521531541551561571581591601611621631641651661671681691701711721731741751761771781791801811821831841851861871881891901911921931941951961971981992002012022032042052062072082092102112122132142152162172182192202212222232242252262272282292302312322332342352362372382392402412422432442452462472482492502512522532542552562572582592602612622632642652662672682692702712722732742752762772782792802812822832842852862872882892902912922932942952962972982993003013023033043053063073083093103113123133143153163173183193203213223233243253263273283293303313323333343353363373383393403413423433443453463473483493503513523533543553563573583593603613623633643653663673683693703713723733743753763773783793803813823833843853863873883893903913923933943953963973983994004014024034044054064074084094104114124134144154164174184194204214224234244254264274284294304314324334344354364374384394404414424434444454464474484494504514524534544554564574584594604614624634644654664674684694704714724734744754764774784794804814824834844854864874884894904914924934944954964974984995005015025035045055065075085095105115125135145155165175185195205215225235245255265275285295305315325335345355365375385395405415425435445455465475485495505515525535545555565575585595605615625635645655665675685695705715725735745755765775785795805815825835845855865875885895905915925935945955965975985996006016026036046056066076086096106116126136146156166176186196206216226236246256266276286296306316326336346356366376386396406416426436446456466476486496506516526536546556566576586596606616626636646656666676686696706716726736746756766776786796806816826836846856866876886896906916926936946956966976986997007017027037047057067077087097107117127137147157167177187197207217227237247257267277287297307317327337347357367377387397407417427437447457467477487497507517527537547557567577587597607617627637647657667677687697707717727737747757767777787797807817827837847857867877887897907917927937947957967977987998008018028038048058068078088098108118128138148158168178188198208218228238248258268278288298308318328338348358368378388398408418428438448458468478488498508518528538548558568578588598608618628638648658668678688698708718728738748758768778788798808818828838848858868878888898908918928938948958968978988999009019029039049059069079089099109119129139149159169179189199209219229239249259269279289299309319329339349359369379389399409419429439449459469479489499509519529539549559569579589599609619629639649659669679689699709719729739749759769779789799809819829839849859869879889899909919929939949959969979989991000100110021003100410051006100710081009101010111012101310141015101610171018101910201021102210231024102510261027102810291030103110321033103410351036103710381039104010411042104310441045104610471048104910501051105210531054105510561057105810591060106110621063106410651066106710681069107010711072107310741075107610771078107910801081108210831084108510861087108810891090109110921093109410951096109710981099110011011102110311041105110611071108110911101111111211131114111511161117111811191120112111221123112411251126112711281129113011311132113311341135113611371138113911401141114211431144114511461147114811491150115111521153115411551156115711581159116011611162116311641165116611671168116911701171117211731174117511761177117811791180118111821183118411851186118711881189119011911192119311941195119611971198119912001201120212031204120512061207120812091210121112121213121412151216121712181219122012211222122312241225122612271228122912301231123212331234123512361237123812391240124112421243124412451246124712481249125012511252125312541255125612571258125912601261126212631264126512661267126812691270127112721273127412751276127712781279128012811282128312841285128612871288128912901291129212931294129512961297129812991300130113021303130413051306130713081309131013111312131313141315131613171318131913201321132213231324132513261327132813291330133113321333133413351336133713381339134013411342134313441345134613471348134913501351135213531354135513561357135813591360136113621363136413651366136713681369137013711372137313741375137613771378137913801381138213831384138513861387138813891390139113921393139413951396139713981399140014011402140314041405140614071408140914101411141214131414
  1. # -------------------------------------------------------------------------
  2. # Copyright (c) Microsoft Corporation. All rights reserved.
  3. # Licensed under the MIT License.
  4. # --------------------------------------------------------------------------
  5. from __future__ import annotations
  6. import collections
  7. import collections.abc
  8. import os
  9. import typing
  10. import warnings
  11. from collections.abc import Callable, Sequence
  12. from typing import Any
  13. from onnxruntime.capi import _pybind_state as C
  14. if typing.TYPE_CHECKING:
  15. import numpy as np
  16. import numpy.typing as npt
  17. import onnxruntime
  18. def get_ort_device_type(device_type: str) -> int:
  19. if device_type == "cuda":
  20. return C.OrtDevice.cuda()
  21. elif device_type == "cann":
  22. return C.OrtDevice.cann()
  23. elif device_type == "cpu":
  24. return C.OrtDevice.cpu()
  25. elif device_type == "dml":
  26. return C.OrtDevice.dml()
  27. elif device_type == "webgpu":
  28. return C.OrtDevice.webgpu()
  29. elif device_type == "gpu":
  30. return C.OrtDevice.gpu()
  31. elif device_type == "npu":
  32. return C.OrtDevice.npu()
  33. else:
  34. raise Exception("Unsupported device type: " + device_type)
  35. class AdapterFormat:
  36. """
  37. This class is used to create adapter files from python structures
  38. """
  39. def __init__(self, adapter=None) -> None:
  40. if adapter is None:
  41. self._adapter = C.AdapterFormat()
  42. else:
  43. self._adapter = adapter
  44. @staticmethod
  45. def read_adapter(file_path: os.PathLike) -> AdapterFormat:
  46. return AdapterFormat(C.AdapterFormat.read_adapter(file_path))
  47. def export_adapter(self, file_path: os.PathLike):
  48. """
  49. This function writes a file at the specified location
  50. in onnxrunitme adapter format containing Lora parameters.
  51. :param file_path: absolute path for the adapter
  52. """
  53. self._adapter.export_adapter(file_path)
  54. def get_format_version(self) -> int:
  55. return self._adapter.format_version
  56. def set_adapter_version(self, adapter_version: int) -> None:
  57. self._adapter.adapter_version = adapter_version
  58. def get_adapter_version(self) -> int:
  59. return self._adapter.adapter_version
  60. def set_model_version(self, model_version: int) -> None:
  61. self._adapter.model_version = model_version
  62. def get_model_version(self) -> int:
  63. return self._adapter.model_version
  64. def set_parameters(self, params: dict[str, OrtValue]) -> None:
  65. self._adapter.parameters = {k: v._ortvalue for k, v in params.items()}
  66. def get_parameters(self) -> dict[str, OrtValue]:
  67. return {k: OrtValue(v) for k, v in self._adapter.parameters.items()}
  68. def check_and_normalize_provider_args(
  69. providers: Sequence[str | tuple[str, dict[Any, Any]]] | None,
  70. provider_options: Sequence[dict[Any, Any]] | None,
  71. available_provider_names: Sequence[str],
  72. ):
  73. """
  74. Validates the 'providers' and 'provider_options' arguments and returns a
  75. normalized version.
  76. :param providers: Optional sequence of providers in order of decreasing
  77. precedence. Values can either be provider names or tuples of
  78. (provider name, options dict).
  79. :param provider_options: Optional sequence of options dicts corresponding
  80. to the providers listed in 'providers'.
  81. :param available_provider_names: The available provider names.
  82. :return: Tuple of (normalized 'providers' sequence, normalized
  83. 'provider_options' sequence).
  84. 'providers' can contain either names or names and options. When any options
  85. are given in 'providers', 'provider_options' should not be used.
  86. The normalized result is a tuple of:
  87. 1. Sequence of provider names in the same order as 'providers'.
  88. 2. Sequence of corresponding provider options dicts with string keys and
  89. values. Unspecified provider options yield empty dicts.
  90. """
  91. if providers is None:
  92. return [], []
  93. provider_name_to_options = collections.OrderedDict()
  94. def set_provider_options(name, options):
  95. if name not in available_provider_names:
  96. warnings.warn(
  97. "Specified provider '{}' is not in available provider names.Available providers: '{}'".format(
  98. name, ", ".join(available_provider_names)
  99. )
  100. )
  101. if name in provider_name_to_options:
  102. warnings.warn(f"Duplicate provider '{name}' encountered, ignoring.")
  103. return
  104. normalized_options = {str(key): str(value) for key, value in options.items()}
  105. provider_name_to_options[name] = normalized_options
  106. if not isinstance(providers, collections.abc.Sequence):
  107. raise ValueError("'providers' should be a sequence.")
  108. if provider_options is not None:
  109. if not isinstance(provider_options, collections.abc.Sequence):
  110. raise ValueError("'provider_options' should be a sequence.")
  111. if len(providers) != len(provider_options):
  112. raise ValueError("'providers' and 'provider_options' should be the same length if both are given.")
  113. if not all(isinstance(provider, str) for provider in providers):
  114. raise ValueError("Only string values for 'providers' are supported if 'provider_options' is given.")
  115. if not all(isinstance(options_for_provider, dict) for options_for_provider in provider_options):
  116. raise ValueError("'provider_options' values must be dicts.")
  117. for name, options in zip(providers, provider_options, strict=False):
  118. set_provider_options(name, options)
  119. else:
  120. for provider in providers:
  121. if isinstance(provider, str):
  122. set_provider_options(provider, {})
  123. elif (
  124. isinstance(provider, tuple)
  125. and len(provider) == 2
  126. and isinstance(provider[0], str)
  127. and isinstance(provider[1], dict)
  128. ):
  129. set_provider_options(provider[0], provider[1])
  130. else:
  131. raise ValueError("'providers' values must be either strings or (string, dict) tuples.")
  132. return list(provider_name_to_options.keys()), list(provider_name_to_options.values())
  133. class Session:
  134. """
  135. This is the main class used to run a model.
  136. """
  137. def __init__(self, enable_fallback: bool = True):
  138. # self._sess is managed by the derived class and relies on bindings from C.InferenceSession
  139. self._sess = None
  140. self._enable_fallback = enable_fallback
  141. def get_session_options(self) -> onnxruntime.SessionOptions:
  142. "Return the session options. See :class:`onnxruntime.SessionOptions`."
  143. return self._sess_options
  144. def get_inputs(self) -> Sequence[onnxruntime.NodeArg]:
  145. "Return the inputs metadata as a list of :class:`onnxruntime.NodeArg`."
  146. return self._inputs_meta
  147. def get_outputs(self) -> Sequence[onnxruntime.NodeArg]:
  148. "Return the outputs metadata as a list of :class:`onnxruntime.NodeArg`."
  149. return self._outputs_meta
  150. def get_overridable_initializers(self) -> Sequence[onnxruntime.NodeArg]:
  151. "Return the inputs (including initializers) metadata as a list of :class:`onnxruntime.NodeArg`."
  152. return self._overridable_initializers
  153. def get_modelmeta(self) -> onnxruntime.ModelMetadata:
  154. "Return the metadata. See :class:`onnxruntime.ModelMetadata`."
  155. return self._model_meta
  156. def get_input_memory_infos(self) -> Sequence[onnxruntime.MemoryInfo]:
  157. "Return the memory info for the inputs."
  158. return self._input_meminfos
  159. def get_output_memory_infos(self) -> Sequence[onnxruntime.MemoryInfo]:
  160. "Return the memory info for the outputs."
  161. return self._output_meminfos
  162. def get_input_epdevices(self) -> Sequence[onnxruntime.OrtEpDevice]:
  163. "Return the execution providers for the inputs."
  164. return self._input_epdevices
  165. def get_providers(self) -> Sequence[str]:
  166. "Return list of registered execution providers."
  167. return self._providers
  168. def get_provider_options(self):
  169. "Return registered execution providers' configurations."
  170. return self._provider_options
  171. def set_providers(self, providers=None, provider_options=None) -> None:
  172. """
  173. Register the input list of execution providers. The underlying session is re-created.
  174. :param providers: Optional sequence of providers in order of decreasing
  175. precedence. Values can either be provider names or tuples of
  176. (provider name, options dict). If not provided, then all available
  177. providers are used with the default precedence.
  178. :param provider_options: Optional sequence of options dicts corresponding
  179. to the providers listed in 'providers'.
  180. 'providers' can contain either names or names and options. When any options
  181. are given in 'providers', 'provider_options' should not be used.
  182. The list of providers is ordered by precedence. For example
  183. `['CUDAExecutionProvider', 'CPUExecutionProvider']`
  184. means execute a node using CUDAExecutionProvider if capable,
  185. otherwise execute using CPUExecutionProvider.
  186. """
  187. # recreate the underlying C.InferenceSession
  188. self._reset_session(providers, provider_options)
  189. def disable_fallback(self) -> None:
  190. """
  191. Disable session.run() fallback mechanism.
  192. """
  193. self._enable_fallback = False
  194. def enable_fallback(self) -> None:
  195. """
  196. Enable session.Run() fallback mechanism. If session.Run() fails due to an internal Execution Provider failure,
  197. reset the Execution Providers enabled for this session.
  198. If GPU is enabled, fall back to CUDAExecutionProvider.
  199. otherwise fall back to CPUExecutionProvider.
  200. """
  201. self._enable_fallback = True
  202. def _validate_input(self, feed_input_names):
  203. missing_input_names = []
  204. for input in self._inputs_meta:
  205. if input.name not in feed_input_names and not input.type.startswith("optional"):
  206. missing_input_names.append(input.name)
  207. if missing_input_names:
  208. raise ValueError(
  209. f"Required inputs ({missing_input_names}) are missing from input feed ({feed_input_names})."
  210. )
  211. def run(self, output_names, input_feed, run_options=None) -> Sequence[np.ndarray | SparseTensor | list | dict]:
  212. """
  213. Compute the predictions.
  214. :param output_names: name of the outputs
  215. :param input_feed: dictionary ``{ input_name: input_value }``
  216. :param run_options: See :class:`onnxruntime.RunOptions`.
  217. :return: list of results, every result is either a numpy array,
  218. a sparse tensor, a list or a dictionary.
  219. ::
  220. sess.run([output_name], {input_name: x})
  221. """
  222. self._validate_input(list(input_feed.keys()))
  223. if not output_names:
  224. output_names = [output.name for output in self._outputs_meta]
  225. try:
  226. return self._sess.run(output_names, input_feed, run_options)
  227. except C.EPFail as err:
  228. if self._enable_fallback:
  229. print(f"EP Error: {err!s} using {self._providers}")
  230. print(f"Falling back to {self._fallback_providers} and retrying.")
  231. self.set_providers(self._fallback_providers)
  232. # Fallback only once.
  233. self.disable_fallback()
  234. return self._sess.run(output_names, input_feed, run_options)
  235. raise
  236. def run_async(self, output_names, input_feed, callback, user_data, run_options=None):
  237. """
  238. Compute the predictions asynchronously in a separate cxx thread from ort intra-op threadpool.
  239. :param output_names: name of the outputs
  240. :param input_feed: dictionary ``{ input_name: input_value }``
  241. :param callback: python function that accept array of results, and a status string on error.
  242. The callback will be invoked by a cxx thread from ort intra-op threadpool.
  243. :param run_options: See :class:`onnxruntime.RunOptions`.
  244. ::
  245. class MyData:
  246. def __init__(self):
  247. # ...
  248. def save_results(self, results):
  249. # ...
  250. def callback(results: np.ndarray, user_data: MyData, err: str) -> None:
  251. if err:
  252. print (err)
  253. else:
  254. # save results to user_data
  255. sess.run_async([output_name], {input_name: x}, callback)
  256. """
  257. self._validate_input(list(input_feed.keys()))
  258. if not output_names:
  259. output_names = [output.name for output in self._outputs_meta]
  260. return self._sess.run_async(output_names, input_feed, callback, user_data, run_options)
  261. def run_with_ort_values(self, output_names, input_dict_ort_values, run_options=None) -> Sequence[OrtValue]:
  262. """
  263. Compute the predictions.
  264. :param output_names: name of the outputs
  265. :param input_dict_ort_values: dictionary ``{ input_name: input_ort_value }``
  266. See ``OrtValue`` class how to create `OrtValue`
  267. from numpy array or `SparseTensor`
  268. :param run_options: See :class:`onnxruntime.RunOptions`.
  269. :return: an array of `OrtValue`
  270. ::
  271. sess.run([output_name], {input_name: x})
  272. """
  273. def invoke(sess, output_names, input_dict_ort_values, run_options):
  274. input_dict = {}
  275. for n, v in input_dict_ort_values.items():
  276. input_dict[n] = v._get_c_value()
  277. result = sess.run_with_ort_values(input_dict, output_names, run_options)
  278. if not isinstance(result, C.OrtValueVector):
  279. raise TypeError("run_with_ort_values() must return a instance of type 'OrtValueVector'.")
  280. ort_values = [OrtValue(v) for v in result]
  281. return ort_values
  282. self._validate_input(list(input_dict_ort_values.keys()))
  283. if not output_names:
  284. output_names = [output.name for output in self._outputs_meta]
  285. try:
  286. return invoke(self._sess, output_names, input_dict_ort_values, run_options)
  287. except C.EPFail as err:
  288. if self._enable_fallback:
  289. print(f"EP Error: {err!s} using {self._providers}")
  290. print(f"Falling back to {self._fallback_providers} and retrying.")
  291. self.set_providers(self._fallback_providers)
  292. # Fallback only once.
  293. self.disable_fallback()
  294. return invoke(self._sess, output_names, input_dict_ort_values, run_options)
  295. raise
  296. def end_profiling(self):
  297. """
  298. End profiling and return results in a file.
  299. The results are stored in a filename if the option
  300. :meth:`onnxruntime.SessionOptions.enable_profiling`.
  301. """
  302. return self._sess.end_profiling()
  303. def get_profiling_start_time_ns(self):
  304. """
  305. Return the nanoseconds of profiling's start time
  306. Comparable to time.monotonic_ns() after Python 3.3
  307. On some platforms, this timer may not be as precise as nanoseconds
  308. For instance, on Windows and MacOS, the precision will be ~100ns
  309. """
  310. return self._sess.get_profiling_start_time_ns
  311. def io_binding(self) -> IOBinding:
  312. "Return an onnxruntime.IOBinding object`."
  313. return IOBinding(self)
  314. def run_with_iobinding(self, iobinding, run_options=None):
  315. """
  316. Compute the predictions.
  317. :param iobinding: the iobinding object that has graph inputs/outputs bind.
  318. :param run_options: See :class:`onnxruntime.RunOptions`.
  319. """
  320. self._sess.run_with_iobinding(iobinding._iobinding, run_options)
  321. def get_tuning_results(self):
  322. return self._sess.get_tuning_results()
  323. def set_tuning_results(self, results, *, error_on_invalid=False):
  324. return self._sess.set_tuning_results(results, error_on_invalid)
  325. def run_with_ortvaluevector(self, run_options, feed_names, feeds, fetch_names, fetches, fetch_devices):
  326. """
  327. Compute the predictions similar to other run_*() methods but with minimal C++/Python conversion overhead.
  328. :param run_options: See :class:`onnxruntime.RunOptions`.
  329. :param feed_names: list of input names.
  330. :param feeds: list of input OrtValue.
  331. :param fetch_names: list of output names.
  332. :param fetches: list of output OrtValue.
  333. :param fetch_devices: list of output devices.
  334. """
  335. self._sess.run_with_ortvaluevector(run_options, feed_names, feeds, fetch_names, fetches, fetch_devices)
  336. class InferenceSession(Session):
  337. """
  338. This is the main class used to run a model.
  339. """
  340. def __init__(
  341. self,
  342. path_or_bytes: str | bytes | os.PathLike,
  343. sess_options: onnxruntime.SessionOptions | None = None,
  344. providers: Sequence[str | tuple[str, dict[Any, Any]]] | None = None,
  345. provider_options: Sequence[dict[Any, Any]] | None = None,
  346. **kwargs,
  347. ) -> None:
  348. """
  349. :param path_or_bytes: Filename or serialized ONNX or ORT format model in a byte string.
  350. :param sess_options: Session options.
  351. :param providers: Optional sequence of providers in order of decreasing
  352. precedence. Values can either be provider names or tuples of
  353. (provider name, options dict). If not provided, then all available
  354. providers are used with the default precedence.
  355. :param provider_options: Optional sequence of options dicts corresponding
  356. to the providers listed in 'providers'.
  357. The model type will be inferred unless explicitly set in the SessionOptions.
  358. To explicitly set:
  359. ::
  360. so = onnxruntime.SessionOptions()
  361. # so.add_session_config_entry('session.load_model_format', 'ONNX') or
  362. so.add_session_config_entry('session.load_model_format', 'ORT')
  363. A file extension of '.ort' will be inferred as an ORT format model.
  364. All other filenames are assumed to be ONNX format models.
  365. 'providers' can contain either names or names and options. When any options
  366. are given in 'providers', 'provider_options' should not be used.
  367. The list of providers is ordered by precedence. For example
  368. `['CUDAExecutionProvider', 'CPUExecutionProvider']`
  369. means execute a node using `CUDAExecutionProvider`
  370. if capable, otherwise execute using `CPUExecutionProvider`.
  371. """
  372. super().__init__(enable_fallback=int(kwargs.get("enable_fallback", 1)) == 1)
  373. if isinstance(path_or_bytes, (str, os.PathLike)):
  374. self._model_path = os.fspath(path_or_bytes)
  375. self._model_bytes = None
  376. elif isinstance(path_or_bytes, bytes):
  377. self._model_path = None
  378. self._model_bytes = path_or_bytes # TODO: This is bad as we're holding the memory indefinitely
  379. else:
  380. raise TypeError(f"Unable to load from type '{type(path_or_bytes)}'")
  381. self._sess_options = sess_options
  382. self._sess_options_initial = sess_options
  383. if "read_config_from_model" in kwargs:
  384. self._read_config_from_model = int(kwargs["read_config_from_model"]) == 1
  385. else:
  386. self._read_config_from_model = os.environ.get("ORT_LOAD_CONFIG_FROM_MODEL") == "1"
  387. # internal parameters that we don't expect to be used in general so aren't documented
  388. disabled_optimizers = kwargs.get("disabled_optimizers")
  389. try:
  390. self._create_inference_session(providers, provider_options, disabled_optimizers)
  391. except (ValueError, RuntimeError) as e:
  392. if self._enable_fallback:
  393. try:
  394. print("*************** EP Error ***************")
  395. print(f"EP Error {e} when using {providers}")
  396. print(f"Falling back to {self._fallback_providers} and retrying.")
  397. print("****************************************")
  398. self._create_inference_session(self._fallback_providers, None)
  399. # Fallback only once.
  400. self.disable_fallback()
  401. return
  402. except Exception as fallback_error:
  403. raise fallback_error from e
  404. # Fallback is disabled. Raise the original error.
  405. raise e
  406. def _create_inference_session(self, providers, provider_options, disabled_optimizers=None):
  407. available_providers = C.get_available_providers()
  408. # Tensorrt can fall back to CUDA if it's explicitly assigned. All others fall back to CPU.
  409. if "TensorrtExecutionProvider" in available_providers:
  410. if (
  411. providers
  412. and any(
  413. provider == "CUDAExecutionProvider"
  414. or (isinstance(provider, tuple) and provider[0] == "CUDAExecutionProvider")
  415. for provider in providers
  416. )
  417. and any(
  418. provider == "TensorrtExecutionProvider"
  419. or (isinstance(provider, tuple) and provider[0] == "TensorrtExecutionProvider")
  420. for provider in providers
  421. )
  422. ):
  423. self._fallback_providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
  424. else:
  425. self._fallback_providers = ["CPUExecutionProvider"]
  426. if "NvTensorRTRTXExecutionProvider" in available_providers:
  427. if (
  428. providers
  429. and any(
  430. provider == "CUDAExecutionProvider"
  431. or (isinstance(provider, tuple) and provider[0] == "CUDAExecutionProvider")
  432. for provider in providers
  433. )
  434. and any(
  435. provider == "NvTensorRTRTXExecutionProvider"
  436. or (isinstance(provider, tuple) and provider[0] == "NvExecutionProvider")
  437. for provider in providers
  438. )
  439. ):
  440. self._fallback_providers = ["CUDAExecutionProvider", "CPUExecutionProvider"]
  441. else:
  442. self._fallback_providers = ["CPUExecutionProvider"]
  443. # MIGraphX can fall back to ROCM if it's explicitly assigned. All others fall back to CPU.
  444. elif "MIGraphXExecutionProvider" in available_providers:
  445. if providers and any(
  446. provider == "ROCMExecutionProvider"
  447. or (isinstance(provider, tuple) and provider[0] == "ROCMExecutionProvider")
  448. for provider in providers
  449. ):
  450. self._fallback_providers = ["ROCMExecutionProvider", "CPUExecutionProvider"]
  451. else:
  452. self._fallback_providers = ["CPUExecutionProvider"]
  453. else:
  454. self._fallback_providers = ["CPUExecutionProvider"]
  455. # validate providers and provider_options before other initialization
  456. providers, provider_options = check_and_normalize_provider_args(
  457. providers, provider_options, available_providers
  458. )
  459. # Print a warning if user passed providers to InferenceSession() but the SessionOptions instance
  460. # already has provider information (e.g., via add_provider_for_devices()). The providers specified
  461. # here will take precedence.
  462. if self._sess_options is not None and (providers or provider_options) and self._sess_options.has_providers():
  463. warnings.warn(
  464. "Specified 'providers'/'provider_options' when creating InferenceSession but SessionOptions has "
  465. "already been configured with providers. InferenceSession will only use the providers "
  466. "passed to InferenceSession()."
  467. )
  468. session_options = self._sess_options if self._sess_options else C.get_default_session_options()
  469. self._register_ep_custom_ops(session_options, providers, provider_options, available_providers)
  470. if self._model_path:
  471. sess = C.InferenceSession(session_options, self._model_path, True, self._read_config_from_model)
  472. else:
  473. sess = C.InferenceSession(session_options, self._model_bytes, False, self._read_config_from_model)
  474. if disabled_optimizers is None:
  475. disabled_optimizers = set()
  476. elif not isinstance(disabled_optimizers, set):
  477. # convert to set. assumes iterable
  478. disabled_optimizers = set(disabled_optimizers)
  479. # initialize the C++ InferenceSession
  480. sess.initialize_session(providers, provider_options, disabled_optimizers)
  481. self._sess = sess
  482. self._sess_options = self._sess.session_options
  483. self._inputs_meta = self._sess.inputs_meta
  484. self._outputs_meta = self._sess.outputs_meta
  485. self._overridable_initializers = self._sess.overridable_initializers
  486. self._input_meminfos = self._sess.input_meminfos
  487. self._output_meminfos = self._sess.output_meminfos
  488. self._input_epdevices = self._sess.input_epdevices
  489. self._model_meta = self._sess.model_meta
  490. self._providers = self._sess.get_providers()
  491. self._provider_options = self._sess.get_provider_options()
  492. self._profiling_start_time_ns = self._sess.get_profiling_start_time_ns
  493. def _reset_session(self, providers, provider_options) -> None:
  494. "release underlying session object."
  495. # meta data references session internal structures
  496. # so they must be set to None to decrement _sess reference count.
  497. self._sess_options = None
  498. self._inputs_meta = None
  499. self._outputs_meta = None
  500. self._overridable_initializers = None
  501. self._input_meminfos = None
  502. self._output_meminfos = None
  503. self._input_epdevices = None
  504. self._model_meta = None
  505. self._providers = None
  506. self._provider_options = None
  507. self._profiling_start_time_ns = None
  508. # create a new C.InferenceSession
  509. self._sess = None
  510. self._sess_options = self._sess_options_initial
  511. self._create_inference_session(providers, provider_options)
  512. def _register_ep_custom_ops(self, session_options, providers, provider_options, available_providers):
  513. for i in range(len(providers)):
  514. if providers[i] in available_providers and providers[i] == "TensorrtExecutionProvider":
  515. C.register_tensorrt_plugins_as_custom_ops(session_options, provider_options[i])
  516. elif (
  517. isinstance(providers[i], tuple)
  518. and providers[i][0] in available_providers
  519. and providers[i][0] == "TensorrtExecutionProvider"
  520. ):
  521. C.register_tensorrt_plugins_as_custom_ops(session_options, providers[i][1])
  522. if providers[i] in available_providers and providers[i] == "NvTensorRTRTXExecutionProvider":
  523. C.register_nv_tensorrt_rtx_plugins_as_custom_ops(session_options, provider_options[i])
  524. elif (
  525. isinstance(providers[i], tuple)
  526. and providers[i][0] in available_providers
  527. and providers[i][0] == "NvTensorrtRTXExecutionProvider"
  528. ):
  529. C.register_nv_tensorrt_rtx_plugins_as_custom_ops(session_options, providers[i][1])
  530. def make_get_initializer_location_func_wrapper(
  531. get_initializer_location_func: GetInitializerLocationFunc,
  532. ) -> GetInitializerLocationWrapperFunc:
  533. """
  534. Wraps a user's "get initializer location" function. The returned wrapper function adheres to the
  535. signature expected by ORT.
  536. Need this wrapper to:
  537. - Convert the `initializer_value` parameter from `C.OrtValue` to `onnxruntime.OrtValue`, which is more
  538. convenient for the user's function to use.
  539. - Allow the user's function to return the original `external_info` parameter (this wrapper makes a copy)
  540. """
  541. def get_initializer_location_func_wrapper(
  542. initializer_name: str,
  543. initializer_value: C.OrtValue,
  544. external_info: C.OrtExternalInitializerInfo | None,
  545. ) -> C.OrtExternalInitializerInfo | None:
  546. ret_val: C.OrtExternalInitializerInfo | None = get_initializer_location_func(
  547. initializer_name, OrtValue(initializer_value), external_info
  548. )
  549. if ret_val is not None and ret_val == external_info:
  550. # User returned `external_info` (const and owned by ORT). ORT expects the returned value to be
  551. # a new instance (that it deletes), so make a copy.
  552. ret_val = C.OrtExternalInitializerInfo(ret_val.filepath, ret_val.file_offset, ret_val.byte_size)
  553. return ret_val
  554. return get_initializer_location_func_wrapper
  555. class ModelCompiler:
  556. """
  557. This class is used to compile an ONNX model. A compiled ONNX model has EPContext nodes that each
  558. encapsulates a subgraph compiled/optimized for a specific execution provider.
  559. Refer to the EPContext design document for more information about EPContext models:
  560. https://onnxruntime.ai/docs/execution-providers/EP-Context-Design.html
  561. ::
  562. sess_options = onnxruntime.SessionOptions()
  563. sess_options.add_provider("SomeExecutionProvider", {"option1": "value1"})
  564. # Alternatively, allow ONNX Runtime to select the provider automatically given a policy:
  565. # sess_options.set_provider_selection_policy(onnxrt.OrtExecutionProviderDevicePolicy.PREFER_NPU)
  566. model_compiler = onnxruntime.ModelCompiler(sess_options, "input_model.onnx")
  567. model_compiler.compile_to_file("output_model.onnx")
  568. """
  569. def __init__(
  570. self,
  571. sess_options: onnxruntime.SessionOptions,
  572. input_model_path_or_bytes: str | os.PathLike | bytes,
  573. embed_compiled_data_into_model: bool = False,
  574. external_initializers_file_path: str | os.PathLike | None = None,
  575. external_initializers_size_threshold: int = 1024,
  576. flags: int = C.OrtCompileApiFlags.NONE,
  577. graph_optimization_level: C.GraphOptimizationLevel = C.GraphOptimizationLevel.ORT_DISABLE_ALL,
  578. get_initializer_location_func: GetInitializerLocationFunc | None = None,
  579. ):
  580. """
  581. Creates a ModelCompiler instance.
  582. :param sess_options: Session options containing the providers for which the model will be compiled.
  583. Refer to SessionOptions.add_provider() and SessionOptions.set_provider_selection_policy().
  584. :param input_model_path_or_bytes: The path to the input model file or bytes representing a serialized
  585. ONNX model.
  586. :param embed_compiled_data_into_model: Defaults to False. Set to True to embed compiled binary data into
  587. EPContext nodes in the compiled model.
  588. :param external_initializers_file_path: Defaults to None. Set to a path for a file that will store the
  589. initializers for non-compiled nodes.
  590. :param external_initializers_size_threshold: Defaults to 1024. Ignored if `external_initializers_file_path`
  591. is None or empty. Initializers larger than this threshold are stored in the external initializers file.
  592. :param flags: Additional boolean options to enable. Set this parameter to a bitwise OR of
  593. flags in onnxruntime.OrtCompileApiFlags.
  594. :param graph_optimization_level: The graph optimization level.
  595. Defaults to onnxruntime.GraphOptimizationLevel.ORT_DISABLE_ALL.
  596. :param get_initializer_location_func: Optional function called for every initializer to allow user to specify
  597. whether an initializer should be stored within the model or externally. Example:
  598. ```
  599. def get_initializer_location(
  600. initializer_name: str,
  601. initializer_value: onnxrt.OrtValue,
  602. external_info: onnxrt.OrtExternalInitializerInfo | None,
  603. ) -> onnxrt.OrtExternalInitializerInfo | None:
  604. byte_size = initializer_value.tensor_size_in_bytes()
  605. if byte_size < 64:
  606. return None # Store small initializer within compiled model.
  607. # Else, write initializer to new external file.
  608. value_np = initializer_value.numpy()
  609. file_offset = ext_init_file.tell()
  610. ext_init_file.write(value_np.tobytes())
  611. return onnxrt.OrtExternalInitializerInfo(initializer_file_path, file_offset, byte_size)
  612. ```
  613. """
  614. input_model_path: str | os.PathLike | None = None
  615. input_model_bytes: bytes | None = None
  616. if isinstance(input_model_path_or_bytes, (str, os.PathLike)):
  617. if not input_model_path_or_bytes:
  618. raise ValueError("Input model path is empty")
  619. input_model_path = os.fspath(input_model_path_or_bytes)
  620. elif isinstance(input_model_path_or_bytes, bytes):
  621. if len(input_model_path_or_bytes) == 0:
  622. raise ValueError("Input model bytes array is empty")
  623. input_model_bytes = input_model_path_or_bytes
  624. else:
  625. raise TypeError(f"Unable to load from type '{type(input_model_path_or_bytes)}'")
  626. if external_initializers_file_path:
  627. if not isinstance(external_initializers_file_path, (str, os.PathLike)):
  628. arg_type = type(external_initializers_file_path)
  629. raise TypeError(f"Output external initializer filepath is of unexpected type '{arg_type}'")
  630. external_initializers_file_path = os.fspath(external_initializers_file_path)
  631. else:
  632. external_initializers_file_path = ""
  633. if get_initializer_location_func is not None:
  634. if external_initializers_file_path:
  635. raise ValueError(
  636. "Cannot initialize ModelCompiler with both `external_initializers_file_path` "
  637. "and `get_initializer_location_func`"
  638. )
  639. self.get_initializer_location_func_wrapper = make_get_initializer_location_func_wrapper(
  640. get_initializer_location_func
  641. )
  642. else:
  643. self.get_initializer_location_func_wrapper = None
  644. if input_model_path:
  645. self._model_compiler = C.ModelCompiler(
  646. sess_options,
  647. input_model_path,
  648. True, # is path
  649. embed_compiled_data_into_model,
  650. external_initializers_file_path,
  651. external_initializers_size_threshold,
  652. flags,
  653. graph_optimization_level,
  654. self.get_initializer_location_func_wrapper,
  655. )
  656. else:
  657. self._model_compiler = C.ModelCompiler(
  658. sess_options,
  659. input_model_bytes,
  660. False, # is bytes
  661. embed_compiled_data_into_model,
  662. external_initializers_file_path,
  663. external_initializers_size_threshold,
  664. flags,
  665. graph_optimization_level,
  666. self.get_initializer_location_func_wrapper,
  667. )
  668. def compile_to_file(self, output_model_path: str | None = None):
  669. """
  670. Compiles to an output file. If an output file path is not provided,
  671. the output file path is generated based on the input model path by replacing
  672. '.onnx' with '_ctx.onnx'. Ex: The generated output file is 'model_ctx.onnx' for
  673. an input model with path 'model.onnx'.
  674. Raises an 'InvalidArgument' exception if the compilation options are invalid.
  675. :param output_model_path: Defaults to None. The path for the output/compiled model.
  676. """
  677. if output_model_path:
  678. if not isinstance(output_model_path, (str, os.PathLike)):
  679. raise TypeError(f"Output model's filepath is of unexpected type '{type(output_model_path)}'")
  680. output_model_path = os.fspath(output_model_path)
  681. self._model_compiler.compile_to_file(output_model_path)
  682. def compile_to_bytes(self) -> bytes:
  683. """
  684. Compiles to bytes representing the serialized compiled ONNX model.
  685. Raises an 'InvalidArgument' exception if the compilation options are invalid.
  686. :return: A bytes object representing the compiled ONNX model.
  687. """
  688. return self._model_compiler.compile_to_bytes()
  689. def compile_to_stream(self, write_function: Callable[[bytes], None]):
  690. """
  691. Compiles the input model and writes the serialized ONNX bytes to a stream using the provided write function.
  692. Raises an 'InvalidArgument' exception if the compilation options are invalid.
  693. :param write_function: A callable that accepts a bytes buffer to write.
  694. """
  695. self._model_compiler.compile_to_stream(write_function)
  696. class IOBinding:
  697. """
  698. This class provides API to bind input/output to a specified device, e.g. GPU.
  699. """
  700. def __init__(self, session: Session):
  701. self._iobinding = C.SessionIOBinding(session._sess)
  702. self._numpy_obj_references = {}
  703. def bind_cpu_input(self, name, arr_on_cpu):
  704. """
  705. bind an input to array on CPU
  706. :param name: input name
  707. :param arr_on_cpu: input values as a python array on CPU
  708. """
  709. # Hold a reference to the numpy object as the bound OrtValue is backed
  710. # directly by the data buffer of the numpy object and so the numpy object
  711. # must be around until this IOBinding instance is around
  712. self._numpy_obj_references[name] = arr_on_cpu
  713. self._iobinding.bind_input(name, arr_on_cpu)
  714. def bind_input(self, name, device_type, device_id, element_type, shape, buffer_ptr):
  715. """
  716. :param name: input name
  717. :param device_type: e.g. cpu, cuda, cann
  718. :param device_id: device id, e.g. 0
  719. :param element_type: input element type. It can be either numpy type (like numpy.float32) or an integer for onnx type (like onnx.TensorProto.BFLOAT16)
  720. :param shape: input shape
  721. :param buffer_ptr: memory pointer to input data
  722. """
  723. self._iobinding.bind_input(
  724. name,
  725. C.OrtDevice(
  726. get_ort_device_type(device_type),
  727. C.OrtDevice.default_memory(),
  728. device_id,
  729. ),
  730. element_type,
  731. shape,
  732. buffer_ptr,
  733. )
  734. def bind_ortvalue_input(self, name, ortvalue):
  735. """
  736. :param name: input name
  737. :param ortvalue: OrtValue instance to bind
  738. """
  739. self._iobinding.bind_ortvalue_input(name, ortvalue._ortvalue)
  740. def synchronize_inputs(self):
  741. self._iobinding.synchronize_inputs()
  742. def bind_output(
  743. self,
  744. name,
  745. device_type="cpu",
  746. device_id=0,
  747. element_type=None,
  748. shape=None,
  749. buffer_ptr=None,
  750. ):
  751. """
  752. :param name: output name
  753. :param device_type: e.g. cpu, cuda, cann, cpu by default
  754. :param device_id: device id, e.g. 0
  755. :param element_type: output element type. It can be either numpy type (like numpy.float32) or an integer for onnx type (like onnx.TensorProto.BFLOAT16)
  756. :param shape: output shape
  757. :param buffer_ptr: memory pointer to output data
  758. """
  759. # Follow the `if` path when the user has not provided any pre-allocated buffer but still
  760. # would like to bind an output to a specific device (e.g. cuda).
  761. # Pre-allocating an output buffer may not be an option for the user as :
  762. # (1) They may not want to use a custom allocator specific to the device they want to bind the output to,
  763. # in which case ORT will allocate the memory for the user
  764. # (2) The output has a dynamic shape and hence the size of the buffer may not be fixed across runs
  765. if buffer_ptr is None:
  766. self._iobinding.bind_output(
  767. name,
  768. C.OrtDevice(
  769. get_ort_device_type(device_type),
  770. C.OrtDevice.default_memory(),
  771. device_id,
  772. ),
  773. )
  774. else:
  775. if element_type is None or shape is None:
  776. raise ValueError("`element_type` and `shape` are to be provided if pre-allocated memory is provided")
  777. self._iobinding.bind_output(
  778. name,
  779. C.OrtDevice(
  780. get_ort_device_type(device_type),
  781. C.OrtDevice.default_memory(),
  782. device_id,
  783. ),
  784. element_type,
  785. shape,
  786. buffer_ptr,
  787. )
  788. def bind_ortvalue_output(self, name, ortvalue):
  789. """
  790. :param name: output name
  791. :param ortvalue: OrtValue instance to bind
  792. """
  793. self._iobinding.bind_ortvalue_output(name, ortvalue._ortvalue)
  794. def synchronize_outputs(self):
  795. self._iobinding.synchronize_outputs()
  796. def get_outputs(self):
  797. """
  798. Returns the output OrtValues from the Run() that preceded the call.
  799. The data buffer of the obtained OrtValues may not reside on CPU memory
  800. """
  801. outputs = self._iobinding.get_outputs()
  802. if not isinstance(outputs, C.OrtValueVector):
  803. raise TypeError("get_outputs() must return an instance of type 'OrtValueVector'.")
  804. return [OrtValue(ortvalue) for ortvalue in outputs]
  805. def get_outputs_as_ortvaluevector(self):
  806. return self._iobinding.get_outputs()
  807. def copy_outputs_to_cpu(self):
  808. """Copy output contents to CPU."""
  809. return self._iobinding.copy_outputs_to_cpu()
  810. def clear_binding_inputs(self):
  811. self._iobinding.clear_binding_inputs()
  812. def clear_binding_outputs(self):
  813. self._iobinding.clear_binding_outputs()
  814. class OrtValue:
  815. """
  816. A data structure that supports all ONNX data formats (tensors and non-tensors) that allows users
  817. to place the data backing these on a device, for example, on a CUDA supported device.
  818. This class provides APIs to construct and deal with OrtValues.
  819. """
  820. def __init__(self, ortvalue: C.OrtValue, numpy_obj: np.ndarray | None = None):
  821. if isinstance(ortvalue, C.OrtValue):
  822. self._ortvalue = ortvalue
  823. # Hold a ref count to the numpy object if the OrtValue is backed directly
  824. # by its data buffer so that it isn't destroyed when the OrtValue is in use
  825. self._numpy_obj = numpy_obj
  826. else:
  827. # An end user won't hit this error
  828. raise ValueError(
  829. "`Provided ortvalue` needs to be of type `onnxruntime.capi.onnxruntime_pybind11_state.OrtValue`"
  830. )
  831. def _get_c_value(self) -> C.OrtValue:
  832. return self._ortvalue
  833. @classmethod
  834. def ortvalue_from_numpy(cls, numpy_obj: np.ndarray, /, device_type="cpu", device_id=0, vendor_id=-1) -> OrtValue:
  835. """
  836. Factory method to construct an OrtValue (which holds a Tensor) from a given Numpy object
  837. A copy of the data in the Numpy object is held by the OrtValue only if the device is NOT cpu
  838. :param numpy_obj: The Numpy object to construct the OrtValue from
  839. :param device_type: e.g. cpu, cuda, cann, cpu by default
  840. :param device_id: device id, e.g. 0
  841. :param vendor_id: The device's PCI vendor id. If provided, the device_type should be "gpu" or "npu".
  842. """
  843. # Hold a reference to the numpy object (if device_type is 'cpu') as the OrtValue
  844. # is backed directly by the data buffer of the numpy object and so the numpy object
  845. # must be around until this OrtValue instance is around
  846. return cls(
  847. C.OrtValue.ortvalue_from_numpy(
  848. numpy_obj,
  849. OrtDevice.make(device_type, device_id, vendor_id)._get_c_device(),
  850. ),
  851. numpy_obj if device_type.lower() == "cpu" else None,
  852. )
  853. @classmethod
  854. def ortvalue_from_numpy_with_onnx_type(cls, data: np.ndarray, /, onnx_element_type: int) -> OrtValue:
  855. """
  856. This method creates an instance of OrtValue on top of the numpy array.
  857. No data copy is made and the lifespan of the resulting OrtValue should never
  858. exceed the lifespan of bytes object. The API attempts to reinterpret
  859. the data type which is expected to be the same size. This is useful
  860. when we want to use an ONNX data type that is not supported by numpy.
  861. :param data: numpy.ndarray.
  862. :param onnx_element_type: a valid onnx TensorProto::DataType enum value
  863. """
  864. return cls(C.OrtValue.ortvalue_from_numpy_with_onnx_type(data, onnx_element_type), data)
  865. @classmethod
  866. def ortvalue_from_shape_and_type(
  867. cls, shape: Sequence[int], element_type, device_type: str = "cpu", device_id: int = 0, vendor_id: int = -1
  868. ) -> OrtValue:
  869. """
  870. Factory method to construct an OrtValue (which holds a Tensor) from given shape and element_type
  871. :param shape: List of integers indicating the shape of the OrtValue
  872. :param element_type: The data type of the elements. It can be either numpy type (like numpy.float32) or an integer for onnx type (like onnx.TensorProto.BFLOAT16).
  873. :param device_type: e.g. cpu, cuda, cann, cpu by default
  874. :param device_id: device id, e.g. 0
  875. :param vendor_id: If provided the device type should be "gpu" or "npu".
  876. """
  877. device = OrtDevice.make(device_type, device_id, vendor_id)._get_c_device()
  878. # Integer for onnx element type (see https://onnx.ai/onnx/api/mapping.html).
  879. # This is helpful for some data type (like TensorProto.BFLOAT16) that is not available in numpy.
  880. if isinstance(element_type, int):
  881. return cls(
  882. C.OrtValue.ortvalue_from_shape_and_onnx_type(
  883. shape,
  884. element_type,
  885. device,
  886. )
  887. )
  888. return cls(
  889. C.OrtValue.ortvalue_from_shape_and_type(
  890. shape,
  891. element_type,
  892. device,
  893. )
  894. )
  895. @classmethod
  896. def ort_value_from_sparse_tensor(cls, sparse_tensor: SparseTensor) -> OrtValue:
  897. """
  898. The function will construct an OrtValue instance from a valid SparseTensor
  899. The new instance of OrtValue will assume the ownership of sparse_tensor
  900. """
  901. return cls(C.OrtValue.ort_value_from_sparse_tensor(sparse_tensor._get_c_tensor()))
  902. def as_sparse_tensor(self) -> SparseTensor:
  903. """
  904. The function will return SparseTensor contained in this OrtValue
  905. """
  906. return SparseTensor(self._ortvalue.as_sparse_tensor())
  907. def data_ptr(self) -> int:
  908. """
  909. Returns the address of the first element in the OrtValue's data buffer
  910. """
  911. return self._ortvalue.data_ptr()
  912. def device_name(self) -> str:
  913. """
  914. Returns the name of the device where the OrtValue's data buffer resides e.g. cpu, cuda, cann
  915. """
  916. return self._ortvalue.device_name().lower()
  917. def shape(self) -> Sequence[int]:
  918. """
  919. Returns the shape of the data in the OrtValue
  920. """
  921. return self._ortvalue.shape()
  922. def data_type(self) -> str:
  923. """
  924. Returns the data type of the data in the OrtValue. E.g. 'tensor(int64)'
  925. """
  926. return self._ortvalue.data_type()
  927. def element_type(self) -> int:
  928. """
  929. Returns the proto type of the data in the OrtValue
  930. if the OrtValue is a tensor.
  931. """
  932. return self._ortvalue.element_type()
  933. def tensor_size_in_bytes(self) -> int:
  934. """
  935. Returns the size of the data in the OrtValue in bytes
  936. if the OrtValue is a tensor.
  937. """
  938. return self._ortvalue.tensor_size_in_bytes()
  939. def has_value(self) -> bool:
  940. """
  941. Returns True if the OrtValue corresponding to an
  942. optional type contains data, else returns False
  943. """
  944. return self._ortvalue.has_value()
  945. def is_tensor(self) -> bool:
  946. """
  947. Returns True if the OrtValue contains a Tensor, else returns False
  948. """
  949. return self._ortvalue.is_tensor()
  950. def is_sparse_tensor(self) -> bool:
  951. """
  952. Returns True if the OrtValue contains a SparseTensor, else returns False
  953. """
  954. return self._ortvalue.is_sparse_tensor()
  955. def is_tensor_sequence(self) -> bool:
  956. """
  957. Returns True if the OrtValue contains a Tensor Sequence, else returns False
  958. """
  959. return self._ortvalue.is_tensor_sequence()
  960. def numpy(self) -> np.ndarray:
  961. """
  962. Returns a Numpy object from the OrtValue.
  963. Valid only for OrtValues holding Tensors. Throws for OrtValues holding non-Tensors.
  964. Use accessors to gain a reference to non-Tensor objects such as SparseTensor
  965. """
  966. return self._ortvalue.numpy()
  967. def update_inplace(self, np_arr) -> None:
  968. """
  969. Update the OrtValue in place with a new Numpy array. The numpy contents
  970. are copied over to the device memory backing the OrtValue. It can be used
  971. to update the input valuess for an InferenceSession with CUDA graph
  972. enabled or other scenarios where the OrtValue needs to be updated while
  973. the memory address can not be changed.
  974. """
  975. self._ortvalue.update_inplace(np_arr)
  976. def copy_tensors(src: Sequence[OrtValue], dst: Sequence[OrtValue], stream=None) -> None:
  977. """
  978. Copy tensor data from source OrtValue sequence to destination OrtValue sequence.
  979. """
  980. c_sources = [s._get_c_value() for s in src]
  981. c_dsts = [d._get_c_value() for d in dst]
  982. C.copy_tensors(c_sources, c_dsts, stream)
  983. class OrtDevice:
  984. """
  985. A data structure that exposes the underlying C++ OrtDevice
  986. """
  987. def __init__(self, c_ort_device):
  988. """
  989. Internal constructor
  990. """
  991. if isinstance(c_ort_device, C.OrtDevice):
  992. self._ort_device = c_ort_device
  993. else:
  994. # An end user won't hit this error
  995. raise ValueError(
  996. "`Provided object` needs to be of type `onnxruntime.capi.onnxruntime_pybind11_state.OrtDevice`"
  997. )
  998. def _get_c_device(self):
  999. """
  1000. Internal accessor to underlying object
  1001. """
  1002. return self._ort_device
  1003. @staticmethod
  1004. def make(ort_device_name, device_id, vendor_id=-1):
  1005. if vendor_id < 0:
  1006. # backwards compatibility with predefined OrtDevice names
  1007. return OrtDevice(
  1008. C.OrtDevice(
  1009. get_ort_device_type(ort_device_name),
  1010. C.OrtDevice.default_memory(),
  1011. device_id,
  1012. )
  1013. )
  1014. else:
  1015. # generic. use GPU or NPU for ort_device_name and provide a vendor id.
  1016. # vendor id of 0 is valid in some cases (e.g. webgpu is generic and does not have a vendor id)
  1017. return OrtDevice(
  1018. C.OrtDevice(
  1019. get_ort_device_type(ort_device_name),
  1020. C.OrtDevice.default_memory(),
  1021. vendor_id,
  1022. device_id,
  1023. )
  1024. )
  1025. def device_id(self):
  1026. return self._ort_device.device_id()
  1027. def device_type(self):
  1028. return self._ort_device.device_type()
  1029. def device_vendor_id(self):
  1030. return self._ort_device.vendor_id()
  1031. def device_mem_type(self):
  1032. return self._ort_device.mem_type()
  1033. class SparseTensor:
  1034. """
  1035. A data structure that project the C++ SparseTensor object
  1036. The class provides API to work with the object.
  1037. Depending on the format, the class will hold more than one buffer
  1038. depending on the format
  1039. """
  1040. def __init__(self, sparse_tensor: C.SparseTensor):
  1041. """
  1042. Internal constructor
  1043. """
  1044. if isinstance(sparse_tensor, C.SparseTensor):
  1045. self._tensor = sparse_tensor
  1046. else:
  1047. # An end user won't hit this error
  1048. raise ValueError(
  1049. "`Provided object` needs to be of type `onnxruntime.capi.onnxruntime_pybind11_state.SparseTensor`"
  1050. )
  1051. def _get_c_tensor(self) -> C.SparseTensor:
  1052. return self._tensor
  1053. @classmethod
  1054. def sparse_coo_from_numpy(
  1055. cls,
  1056. dense_shape: npt.NDArray[np.int64],
  1057. values: np.ndarray,
  1058. coo_indices: npt.NDArray[np.int64],
  1059. ort_device: OrtDevice,
  1060. ) -> SparseTensor:
  1061. """
  1062. Factory method to construct a SparseTensor in COO format from given arguments
  1063. :param dense_shape: 1-D numpy array(int64) or a python list that contains a dense_shape of the sparse tensor
  1064. must be on cpu memory
  1065. :param values: a homogeneous, contiguous 1-D numpy array that contains non-zero elements of the tensor
  1066. of a type.
  1067. :param coo_indices: contiguous numpy array(int64) that contains COO indices for the tensor. coo_indices may
  1068. have a 1-D shape when it contains a linear index of non-zero values and its length must be equal to
  1069. that of the values. It can also be of 2-D shape, in which has it contains pairs of coordinates for
  1070. each of the nnz values and its length must be exactly twice of the values length.
  1071. :param ort_device: - describes the backing memory owned by the supplied nummpy arrays. Only CPU memory is
  1072. suppored for non-numeric data types.
  1073. For primitive types, the method will map values and coo_indices arrays into native memory and will use
  1074. them as backing storage. It will increment the reference count for numpy arrays and it will decrement it
  1075. on GC. The buffers may reside in any storage either CPU or GPU.
  1076. For strings and objects, it will create a copy of the arrays in CPU memory as ORT does not support those
  1077. on other devices and their memory can not be mapped.
  1078. """
  1079. return cls(C.SparseTensor.sparse_coo_from_numpy(dense_shape, values, coo_indices, ort_device._get_c_device()))
  1080. @classmethod
  1081. def sparse_csr_from_numpy(
  1082. cls,
  1083. dense_shape: npt.NDArray[np.int64],
  1084. values: np.ndarray,
  1085. inner_indices: npt.NDArray[np.int64],
  1086. outer_indices: npt.NDArray[np.int64],
  1087. ort_device: OrtDevice,
  1088. ) -> SparseTensor:
  1089. """
  1090. Factory method to construct a SparseTensor in CSR format from given arguments
  1091. :param dense_shape: 1-D numpy array(int64) or a python list that contains a dense_shape of the
  1092. sparse tensor (rows, cols) must be on cpu memory
  1093. :param values: a contiguous, homogeneous 1-D numpy array that contains non-zero elements of the tensor
  1094. of a type.
  1095. :param inner_indices: contiguous 1-D numpy array(int64) that contains CSR inner indices for the tensor.
  1096. Its length must be equal to that of the values.
  1097. :param outer_indices: contiguous 1-D numpy array(int64) that contains CSR outer indices for the tensor.
  1098. Its length must be equal to the number of rows + 1.
  1099. :param ort_device: - describes the backing memory owned by the supplied nummpy arrays. Only CPU memory is
  1100. suppored for non-numeric data types.
  1101. For primitive types, the method will map values and indices arrays into native memory and will use them as
  1102. backing storage. It will increment the reference count and it will decrement then count when it is GCed.
  1103. The buffers may reside in any storage either CPU or GPU.
  1104. For strings and objects, it will create a copy of the arrays in CPU memory as ORT does not support those
  1105. on other devices and their memory can not be mapped.
  1106. """
  1107. return cls(
  1108. C.SparseTensor.sparse_csr_from_numpy(
  1109. dense_shape,
  1110. values,
  1111. inner_indices,
  1112. outer_indices,
  1113. ort_device._get_c_device(),
  1114. )
  1115. )
  1116. def values(self) -> np.ndarray:
  1117. """
  1118. The method returns a numpy array that is backed by the native memory
  1119. if the data type is numeric. Otherwise, the returned numpy array that contains
  1120. copies of the strings.
  1121. """
  1122. return self._tensor.values()
  1123. def as_coo_view(self):
  1124. """
  1125. The method will return coo representation of the sparse tensor which will enable
  1126. querying COO indices. If the instance did not contain COO format, it would throw.
  1127. You can query coo indices as:
  1128. ::
  1129. coo_indices = sparse_tensor.as_coo_view().indices()
  1130. which will return a numpy array that is backed by the native memory.
  1131. """
  1132. return self._tensor.get_coo_data()
  1133. def as_csrc_view(self):
  1134. """
  1135. The method will return CSR(C) representation of the sparse tensor which will enable
  1136. querying CRS(C) indices. If the instance dit not contain CSR(C) format, it would throw.
  1137. You can query indices as:
  1138. ::
  1139. inner_ndices = sparse_tensor.as_csrc_view().inner()
  1140. outer_ndices = sparse_tensor.as_csrc_view().outer()
  1141. returning numpy arrays backed by the native memory.
  1142. """
  1143. return self._tensor.get_csrc_data()
  1144. def as_blocksparse_view(self):
  1145. """
  1146. The method will return coo representation of the sparse tensor which will enable
  1147. querying BlockSparse indices. If the instance did not contain BlockSparse format, it would throw.
  1148. You can query coo indices as:
  1149. ::
  1150. block_sparse_indices = sparse_tensor.as_blocksparse_view().indices()
  1151. which will return a numpy array that is backed by the native memory
  1152. """
  1153. return self._tensor.get_blocksparse_data()
  1154. def to_cuda(self, ort_device):
  1155. """
  1156. Returns a copy of this instance on the specified cuda device
  1157. :param ort_device: with name 'cuda' and valid gpu device id
  1158. The method will throw if:
  1159. - this instance contains strings
  1160. - this instance is already on GPU. Cross GPU copy is not supported
  1161. - CUDA is not present in this build
  1162. - if the specified device is not valid
  1163. """
  1164. return SparseTensor(self._tensor.to_cuda(ort_device._get_c_device()))
  1165. def format(self):
  1166. """
  1167. Returns a OrtSparseFormat enumeration
  1168. """
  1169. return self._tensor.format
  1170. def dense_shape(self) -> npt.NDArray[np.int64]:
  1171. """
  1172. Returns a numpy array(int64) containing a dense shape of a sparse tensor
  1173. """
  1174. return self._tensor.dense_shape()
  1175. def data_type(self) -> str:
  1176. """
  1177. Returns a string data type of the data in the OrtValue
  1178. """
  1179. return self._tensor.data_type()
  1180. def device_name(self) -> str:
  1181. """
  1182. Returns the name of the device where the SparseTensor data buffers reside e.g. cpu, cuda
  1183. """
  1184. return self._tensor.device_name().lower()
  1185. # Type hint for user-specified function that allows the user to specify initializer locations when compiling a model.
  1186. GetInitializerLocationFunc = Callable[
  1187. [str, OrtValue, C.OrtExternalInitializerInfo | None], C.OrtExternalInitializerInfo | None
  1188. ]
  1189. # Type hint that adheres to the signature expected by ORT.
  1190. GetInitializerLocationWrapperFunc = Callable[
  1191. [str, C.OrtValue, C.OrtExternalInitializerInfo | None], C.OrtExternalInitializerInfo | None
  1192. ]