API#

IO#

class kedro_onnx.io.datasets.OnnxDataSet(filepath, backend='onnx', load_args=None, version=None, credentials=None, fs_args=None)[source]#

Bases: FsspecDataSet[object, ModelProto]

Loads and saves ONNX models.

backend#

ONNX backend to use.

Type:

OnnxFrameworks

Example

>>> from kedro_onnx.io import OnnxDataSet, OnnxSaveModel
>>> from kedro_onnx.io import FloatTensorType
>>> from sklearn.linear_model import LinearRegression
>>>
>>> path = fs.path('test.onnx')
>>> data_set = OnnxDataSet(path, backend='sklearn')
>>>
>>> model = LinearRegression()
>>> model = model.fit([[1], [2], [3]], [2, 4, 6])
>>>
>>> save_model = OnnxSaveModel(model=model,
...     kwargs={'initial_types': (
...                 ('input', FloatTensorType([None, 1])),)})
>>> data_set.save(save_model)
>>> onnx_model = data_set.load()
>>> onnx_model.producer_name
'skl2onnx'
>>>
>>> from kedro_onnx.inference import run
>>> run(onnx_model, [[4]])
array([[8.]], dtype=float32)

For some backends, you may have to specify additional kwargs:

In the example above, we used the sklearn backend. This backend requires the initial_types kwarg to be specified. For more information, see the skl2onnx documentation.

>>> data_set.save(model)
Traceback (most recent call last):
...
kedro.io.core.DataSetError: You need to specify `initial_types` for
 `sklearn` backend.
Use the `kedro_onnx.OnnxSaveModel` `kwargs` to specify additional
 arguments to the conversion function.

You can even use the bare onnx model as the backend:

>>> data_set = OnnxDataSet(path, backend='onnx')
>>> data_set.save(onnx_model)  # already an onnx model
>>> onnx_model = data_set.load()
>>> onnx_model.producer_name
'skl2onnx'

Initialises OnnxDataSet.

Parameters:

  • filepath (str) – Filepath in POSIX format to a ONNX file prefixed with a protocol like s3://. If prefix is not provided, file protocol (local filesystem) will be used. The prefix should be any protocol supported by fsspec. Note: http(s) doesn’t support versioning.

  • backend (OnnxFrameworks) – ONNX backend to use. To see the list of supported backends, look at kedro_onnx.typing.OnnxFrameworks. Defaults to onnx.

  • load_args (Dict[str, Any], optional) – Arguments for the conversion function from onnxmltools. Defaults to None.

  • version (Version, optional) – If specified, should be an instance of kedro.io.core.Version. If its load attribute is None, the latest version will be loaded. If its save attribute is None, save version will be autogenerated. Defaults to None.

  • credentials (Dict[str, Any], optional) – Credentials required to get access to the underlying filesystem. E.g. for GCSFileSystem it should look like {“token”: None}.

  • fs_args (Dict[str, Any], optional) – Extra arguments to pass into underlying filesystem class constructor (e.g. {“project”: “my-project”} for GCSFileSystem), as well as to pass to the filesystem’s open method through nested keys open_args_load and open_args_save.

Note

Here you can find all available arguments for open:

https://filesystem-spec.readthedocs.io/en/latest/api.html#fsspec.spec.AbstractFileSystem.open # noqa: E501

All defaults will be preserved, unless overwritten. mode argument is ignored and overwritten with rb or wb. Defaults to None.

class kedro_onnx.io.datasets.OnnxSaveModel(model, kwargs=<factory>)[source]#

Bases: object

Object to store an ONNX model and kwargs for the converter function.

Execution#

kedro_onnx.inference.run(model, inputs, output_names=None, inference_options=None, run_options=None)[source]#

Runs an ONNX model.

Parameters:

  • model (ModelProto) – ONNX model.

  • inputs (Union[Dict[str, Any], Any]) – Inputs to the model. Keys are the input names defined in the initial_types or in the model itself. Values are the input data. If the input passed is not a dictionary, the function creates a dictionary like this {“input”: inputs}.

  • inference_options (Dict[str, Any], optional) – Options for the inference sess_options: Session options. providers: List of providers. provider_options: List of provider options. kwargs (Dict[str, Any]): Other options.

  • run_options (Optional[Dict[str, Any]]) – Options for the inference run.

Returns:

Output of the model.

Return type:

Any