kedro_onnx.io.datasets

ONNX datasets.

Module Contents

Classes

FsspecDataSet	An abstract DataSet for creating a new DataSet using fsspec.
OnnxSaveModel	Object to store an ONNX model and kwargs for the converter function.
OnnxDataSet	Loads and saves ONNX models.

Attributes

logger

kedro_onnx.io.datasets.logger

class kedro_onnx.io.datasets.FsspecDataSet(filepath: str, load_args: Dict[str, Any] = None, save_args: Dict[str, Any] = None, version: kedro.io.core.Version = None, credentials: Dict[str, Any] = None, fs_args: Dict[str, Any] = None)

Bases: kedro.io.core.AbstractVersionedDataSet[kedro_onnx.typing.IT, kedro_onnx.typing.OT]

An abstract DataSet for creating a new DataSet using fsspec.

DEFAULT_LOAD_ARGS :Dict[str, Any]

DEFAULT_SAVE_ARGS :Dict[str, Any]

_describe() → Dict[str, Any]

abstract _load_fp(fp: io.IOBase) → kedro_onnx.typing.OT

_load() → kedro_onnx.typing.OT

abstract _save_fp(fp: io.IOBase, data: Any) → None

_save(data: kedro_onnx.typing.IT) → None

_exists() → bool

_release() → None

_invalidate_cache() → None

Invalidate underlying filesystem caches.

class kedro_onnx.io.datasets.OnnxSaveModel

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

model :Any

kwargs :dict

class kedro_onnx.io.datasets.OnnxDataSet(filepath: str, backend: kedro_onnx.typing.OnnxFrameworks = 'onnx', load_args: Dict[str, Any] = None, version: kedro.io.core.Version = None, credentials: Dict[str, Any] = None, fs_args: Dict[str, Any] = None)

Bases: FsspecDataSet[object, kedro_onnx.typing.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'

_describe() → Dict[str, Any]

_load_fp(fp: io.IOBase) → kedro_onnx.typing.ModelProto

_validate_kwarg(kwargs: dict, key: str)

_validate_sklearn(model: Any, kwargs: dict)

_validate_lightbgm(model: Any, kwargs: dict)

_validate_sparkml(model: Any, kwargs: dict)

_validate_xgboost(model: Any, kwargs: dict)

_validate(model: Any, kwargs: dict)

_convert(model: Any, kwargs: Any) → kedro_onnx.typing.ModelProto

_save_fp(fp: io.IOBase, data: Union[OnnxSaveModel, onnx.ModelProto, Any]) → None