neural_compressor.data.datasets

Built-in datasets class for multiple framework backends.

Submodules

• neural_compressor.data.datasets.bert_dataset
• neural_compressor.data.datasets.coco_dataset
• neural_compressor.data.datasets.dataset
• neural_compressor.data.datasets.dummy_dataset
• neural_compressor.data.datasets.dummy_dataset_v2
• neural_compressor.data.datasets.imagenet_dataset
• neural_compressor.data.datasets.style_transfer_dataset

Package Contents

Classes

Datasets	A base class for all framework datasets.
Dataset	The base class of dataset.
IterableDataset	An iterable Dataset.

Functions

dataset_registry(dataset_type, framework[, dataset_format])

Register dataset subclasses.

class neural_compressor.data.datasets.Datasets(framework)

Bases: object

A base class for all framework datasets.

Parameters:

framework (str) – framework name, like:”tensorflow”, “tensorflow_itex”, “mxnet”, “onnxrt_qdq”, “onnxrt_qlinearops”, “onnxrt_integerops”, “pytorch”, “pytorch_ipex”, “pytorch_fx”, “onnxrt_qoperator”.

class neural_compressor.data.datasets.Dataset

Bases: object

The base class of dataset.

Subclass datasets should overwrite two methods: __getitem__ for indexing to data sample and `__len__`for the size of the dataset

class neural_compressor.data.datasets.IterableDataset

Bases: object

An iterable Dataset.

Subclass iterable dataset should also implement a method: __iter__ for interating over the samples of the dataset.

neural_compressor.data.datasets.dataset_registry(dataset_type, framework, dataset_format='')

Register dataset subclasses.

Parameters:

• cls (class) – The class of register.

• dataset_type (str) – The dataset registration name

• framework (str) – support 3 framework including ‘tensorflow’, ‘pytorch’, ‘mxnet’

• data_format (str) – The format dataset saved, eg ‘raw_image’, ‘tfrecord’

Returns:

The class of register.

Return type:

cls