diff --git a/examples/tensorflow/nlp/bert_base_mrpc/quantization/ptq/README.md b/examples/tensorflow/nlp/bert_base_mrpc/quantization/ptq/README.md index 663a9cd89ee..d978b0ed335 100644 --- a/examples/tensorflow/nlp/bert_base_mrpc/quantization/ptq/README.md +++ b/examples/tensorflow/nlp/bert_base_mrpc/quantization/ptq/README.md @@ -173,7 +173,7 @@ After prepare step is done, we add tune and benchmark code to generate quantized #### Benchmark ```python from neural_compressor.experimental import Benchmark, common - from neural_compressor.model.model import get_model_type + from neural_compressor.model.tensorflow_model import get_model_type evaluator = Benchmark(FLAGS.config) dataset = Dataset(eval_file, FLAGS.eval_batch_size) evaluator.b_dataloader = common.DataLoader(\ diff --git a/examples/tensorflow/nlp/bert_base_mrpc/quantization/ptq/run_classifier.py b/examples/tensorflow/nlp/bert_base_mrpc/quantization/ptq/run_classifier.py index 9ffd1c6c1a0..5620ab3775a 100644 --- a/examples/tensorflow/nlp/bert_base_mrpc/quantization/ptq/run_classifier.py +++ b/examples/tensorflow/nlp/bert_base_mrpc/quantization/ptq/run_classifier.py @@ -1109,7 +1109,7 @@ def result(self): evaluator.metric = Accuracy() - from neural_compressor.model.model import get_model_type + from neural_compressor.model.tensorflow_model import get_model_type model_type = get_model_type(FLAGS.input_model) if model_type == 'frozen_pb': evaluator.model = FLAGS.input_model diff --git a/neural_compressor/adaptor/mxnet_utils/util.py b/neural_compressor/adaptor/mxnet_utils/util.py index 98577937f56..4aad408179c 100644 --- a/neural_compressor/adaptor/mxnet_utils/util.py +++ b/neural_compressor/adaptor/mxnet_utils/util.py @@ -24,7 +24,7 @@ from enum import Enum from tempfile import TemporaryDirectory from neural_compressor.utils.utility import LazyImport -from neural_compressor.model.model import MXNetModel as NCModel +from neural_compressor.model.mxnet_model import MXNetModel as NCModel mx = LazyImport("mxnet") diff --git a/neural_compressor/adaptor/tensorflow.py b/neural_compressor/adaptor/tensorflow.py index d4956ce4167..22ee3c88fe4 100644 --- a/neural_compressor/adaptor/tensorflow.py +++ b/neural_compressor/adaptor/tensorflow.py @@ -136,7 +136,7 @@ def train(self, model, dataloader, optimizer_tuple, criterion_tuple, hooks, postprocess, **kwargs): # check model is savedmodel or not import tensorflow as tf - from neural_compressor.model.model import get_model_type + from neural_compressor.model.tensorflow_model import get_model_type tf.random.set_seed(1) self.model_type = get_model_type(model._model) optimizer = optimizer_tuple[0](**optimizer_tuple[1]) @@ -1204,7 +1204,7 @@ def inspect_tensor(self, model, dataloader=None, op_list=[], iteration_list=[], ] } """ - from neural_compressor.model.model import TensorflowBaseModel + from neural_compressor.model.tensorflow_model import TensorflowBaseModel from neural_compressor.utils.utility import load_data_from_pkl, dump_data_to_local from neural_compressor.adaptor.tf_utils.graph_util import GraphAnalyzer from .tf_utils.util import int8_node_name_reverse @@ -1586,7 +1586,8 @@ def _get_mse_order(self, fp32_model, tune_cfg, replace_cfgs, ops_lst, dataloader def _partial_dataset_of(self, dataloader, confidence_batches): from neural_compressor.experimental.data.datasets.dummy_dataset import DummyDataset - if isinstance(dataloader.dataset, DummyDataset): + from neural_compressor.data.datasets.dummy_dataset import DummyDataset as DummyDataset_v2_x + if isinstance(dataloader.dataset, DummyDataset) or isinstance(dataloader.dataset, DummyDataset_v2_x): assert(isinstance(confidence_batches, int)) ds = copy.deepcopy(dataloader.dataset) ds.dataset = ds.dataset[:confidence_batches] diff --git a/neural_compressor/data/__init__.py b/neural_compressor/data/__init__.py index 2883a446e7b..6729875aa67 100644 --- a/neural_compressor/data/__init__.py +++ b/neural_compressor/data/__init__.py @@ -14,26 +14,29 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +# ============================================================================== +"""Built-in dataloaders, datasets, transforms, filters for multiple framework backends.""" -from .dataloaders import DataLoader import neural_compressor.data.datasets import neural_compressor.data.transforms -from ..experimental.data.datasets import DATASETS, Dataset, IterableDataset, dataset_registry -from ..experimental.data.transforms import TRANSFORMS, BaseTransform, transform_registry -from ..experimental.data.dataloaders import DATALOADERS -from ..experimental.data.filters import FILTERS, Filter, filter_registry +from .datasets import Datasets, Dataset, IterableDataset, dataset_registry +from .dataloaders import DATALOADERS, DataLoader +from .transforms import TRANSFORMS, BaseTransform, transform_registry, Postprocess + +from .filters import FILTERS, Filter, filter_registry __all__ = [ "DataLoader", "DATALOADERS", - "DATASETS", + "Datasets", "Dataset", "IterableDataset", "dataset_registry", "TRANSFORMS", "BaseTransform", "transform_registry", + "Postprocess", "FILTERS", "Filter", "filter_registry",] diff --git a/neural_compressor/data/dataloaders/__init__.py b/neural_compressor/data/dataloaders/__init__.py index 560070e31e6..b2568bad678 100644 --- a/neural_compressor/data/dataloaders/__init__.py +++ b/neural_compressor/data/dataloaders/__init__.py @@ -14,9 +14,11 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +# ============================================================================== -from .dataloader import DataLoader +from .dataloader import DataLoader, DATALOADERS __all__ = [ "DataLoader", -] + "DATALOADERS" +] \ No newline at end of file diff --git a/neural_compressor/data/dataloaders/base_dataloader.py b/neural_compressor/data/dataloaders/base_dataloader.py new file mode 100644 index 00000000000..9349760239e --- /dev/null +++ b/neural_compressor/data/dataloaders/base_dataloader.py @@ -0,0 +1,119 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""BaseDataloder of all dataloaders.""" + +from abc import abstractmethod + + +class BaseDataLoader: # pragma: no cover + """Base class for all DataLoaders. + + _generate_dataloader is needed to create a dataloader object + from the general params like batch_size and sampler. The dynamic batching is just to + generate a new dataloader by setting batch_size and last_batch. + + """ + + def __init__(self, dataset, batch_size=1, last_batch='rollover', collate_fn=None, + sampler=None, batch_sampler=None, num_workers=0, pin_memory=False, + shuffle=False, distributed=False): + """Initialize BaseDataLoader. + + Args: + dataset (object): dataset from which to load the data + batch_size (int, optional): number of samples per batch. Defaults to 1. + last_batch (str, optional): whether to drop the last batch if it is incomplete. + Support ['rollover', 'discard'], rollover means False, discard means True. + Defaults to 'rollover'. + collate_fn (callable, optional): merge data with outer dimension batch size. Defaults to None. + sampler (Sampler, optional): Sampler object to sample data. Defaults to None. + batch_sampler (BatchSampler, optional): BatchSampler object to generate batch of indices. Defaults to None. + num_workers (int, optional): number of subprocesses to use for data loading. Defaults to 0. + pin_memory (bool, optional): whether to copy data into pinned memory before returning. Defaults to False. + shuffle (bool, optional): whether to shuffle data. Defaults to False. + distributed (bool, optional): whether the dataloader is distributed. Defaults to False. + """ + self.dataset = dataset + self.collate_fn = collate_fn + self.sampler = sampler + self.batch_sampler = batch_sampler + self.num_workers = num_workers + self.pin_memory = pin_memory + self._batch_size = batch_size + self.shuffle = shuffle + self.distributed = distributed + self.last_batch = last_batch + self.drop_last = False if last_batch == 'rollover' else True + + self.dataloader = self._generate_dataloader( + self.dataset, + batch_size=batch_size, + last_batch=last_batch, + collate_fn=collate_fn, + sampler=sampler, + batch_sampler=batch_sampler, + num_workers=num_workers, + pin_memory=pin_memory, + shuffle=shuffle, + distributed=distributed) + + def batch(self, batch_size, last_batch=None): + """Set batch size for dataloader. + + Args: + batch_size (int): number of samples per batch. + last_batch (str, optional): whether to drop the last batch if it is incomplete. + Support ['rollover', 'discard'], rollover means False, discard means True. + Defaults to None. + """ + self._batch_size = batch_size + if last_batch is not None: + self.last_batch = last_batch + self.dataloader = self._generate_dataloader( + self.dataset, + batch_size, + self.last_batch, + self.collate_fn, + self.sampler, + self.batch_sampler, + self.num_workers, + self.pin_memory, + self.shuffle, + self.distributed) + + @property + def batch_size(self): + """Get dataloader's batch_size. + + Returns: + int: batch_size + """ + return self._batch_size + + def __iter__(self): + """Yield data in iterative order. + + Returns: + iterator: iterator for dataloder + """ + return iter(self.dataloader) + + @abstractmethod + def _generate_dataloader(self, dataset, batch_size, last_batch, collate_fn, sampler, + batch_sampler, num_workers, pin_memory, shuffle, distributed): + raise NotImplementedError diff --git a/neural_compressor/data/dataloaders/dataloader.py b/neural_compressor/data/dataloaders/dataloader.py index 3c7078af916..5683b6c8e12 100644 --- a/neural_compressor/data/dataloaders/dataloader.py +++ b/neural_compressor/data/dataloaders/dataloader.py @@ -15,13 +15,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Built-in dataloaders for multiple framework backends.""" + from neural_compressor.experimental.data.dataloaders import DATALOADERS # THIS API IS TO BE DEPRECATED! class DataLoader(object): """Entrance of all configured DataLoaders. Will dispatch the DataLoaders to framework specific one. Users will be not aware of the dispatching, and the Interface is unified. - """ def __new__(cls, framework, dataset, batch_size=1, collate_fn=None, diff --git a/neural_compressor/data/dataloaders/default_dataloader.py b/neural_compressor/data/dataloaders/default_dataloader.py new file mode 100644 index 00000000000..d9a2d74fb26 --- /dev/null +++ b/neural_compressor/data/dataloaders/default_dataloader.py @@ -0,0 +1,143 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Default dataloader for multiple framework backends.""" + +import collections +import numpy as np +from math import ceil, floor +from abc import abstractmethod +from .sampler import IterableSampler, SequentialSampler, BatchSampler +from .fetcher import FETCHERS +from .base_dataloader import BaseDataLoader + +def default_collate(batch): # pragma: no cover + """Merge data with outer dimension batch size.""" + elem = batch[0] + if isinstance(elem, collections.abc.Mapping): + return {key: default_collate([d[key] for d in batch]) for key in elem} + elif isinstance(elem, collections.abc.Sequence): + batch = zip(*batch) + return [default_collate(samples) for samples in batch] + elif isinstance(elem, np.ndarray): + try: + return np.stack(batch) + except: + return batch + else: + return batch + +class DefaultDataLoader(BaseDataLoader): # pragma: no cover + """DefaultDataLoader for multiple framework backends.""" + + def __init__(self, dataset, batch_size=1, last_batch='rollover', collate_fn=None, + sampler=None, batch_sampler=None, num_workers=0, pin_memory=False, + shuffle=False, distributed=False): + """Initialize DefaultDataLoader. + + Args: + dataset (object): dataset from which to load the data + batch_size (int, optional): number of samples per batch. Defaults to 1. + last_batch (str, optional): whether to drop the last batch if it is incomplete. + Support ['rollover', 'discard'], rollover means False, discard means True. + Defaults to 'rollover'. + collate_fn (callable, optional): merge data with outer dimension batch size. Defaults to None. + sampler (Sampler, optional): Sampler object to sample data. Defaults to None. + batch_sampler (BatchSampler, optional): BatchSampler object to generate batch of indices. Defaults to None. + num_workers (int, optional): number of subprocesses to use for data loading. Defaults to 0. + pin_memory (bool, optional): whether to copy data into pinned memory before returning. Defaults to False. + shuffle (bool, optional): whether to shuffle data. Defaults to False. + distributed (bool, optional): whether the dataloader is distributed. Defaults to False. + """ + self.dataset = dataset + self.last_batch = last_batch + self.sampler = sampler + self.batch_sampler = batch_sampler + self.num_workers = num_workers + self.pin_memory = pin_memory + self.collate_fn = collate_fn + self._batch_size = batch_size + self.shuffle = shuffle + self.distributed = distributed + self.drop_last = False if last_batch == 'rollover' else True + if self.collate_fn == None: + self.collate_fn = default_collate + + def batch(self, batch_size, last_batch='rollover'): + """Set batch_size and last_batch.""" + self._batch_size = batch_size + self.last_batch = last_batch + + @property + def dataloader(self): + """Return dataloader.""" + return self + + def __iter__(self): + """Yield data in iterative order.""" + return self._generate_dataloader( + self.dataset, + batch_size=self.batch_size, + last_batch=self.last_batch, + collate_fn=self.collate_fn, + sampler=self.sampler, + batch_sampler=self.batch_sampler, + num_workers=self.num_workers, + pin_memory=self.pin_memory, + shuffle=self.shuffle, + distributed=self.distributed) + + def __len__(self): + """Get dataset length.""" + try: + dataset_len = self.dataset.__len__() + except (AttributeError, TypeError): + dataset_len = 0 + for _ in self.dataset: + dataset_len += 1 + except Exception: + raise ValueError(f"{self.dataset} is invalid, {self.dataset}" \ + " does not support calculating the length of its dataloader") + if self.drop_last == False: + dataloader_len = ceil(dataset_len / self.batch_size) + else: + dataloader_len = floor(dataset_len / self.batch_size) + return dataloader_len + + def _generate_dataloader(self, dataset, batch_size, last_batch, collate_fn, sampler, + batch_sampler, num_workers, pin_memory, shuffle, distributed): + + sampler = self._generate_sampler(dataset, distributed) + self.batch_sampler = BatchSampler(sampler, batch_size, self.drop_last) + self.fetcher = FETCHERS[self.dataset_type](dataset, collate_fn, self.drop_last, distributed) + + for batched_indices in self.batch_sampler: + try: + data = self.fetcher(batched_indices) + yield data + except StopIteration: + return + + def _generate_sampler(self, dataset, distributed): + if hasattr(dataset, "__getitem__"): + self.dataset_type = 'index' + return SequentialSampler(dataset, distributed) + elif hasattr(dataset, "__iter__"): + self.dataset_type = 'iter' + return IterableSampler(dataset) + else: + raise ValueError("dataset type only support (index, iter)") diff --git a/neural_compressor/data/dataloaders/fetcher.py b/neural_compressor/data/dataloaders/fetcher.py new file mode 100644 index 00000000000..01ab6d895fa --- /dev/null +++ b/neural_compressor/data/dataloaders/fetcher.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Definitions of the methods to fetch data from an iterable-style or list-style dataset.""" + +from abc import abstractmethod + +class Fetcher(object): # pragma: no cover + """Base class for different fetchers.""" + + def __init__(self, dataset, collate_fn, drop_last): + """Initialize Fetcher. + + Args: + dataset (object): dataset object from which to get data + collate_fn (callable): merge data with outer dimension batch size + drop_last (bool): whether to drop the last batch if it is incomplete + """ + self.dataset = dataset + self.collate_fn = collate_fn + self.drop_last = drop_last + + @abstractmethod + def __call__(self, batched_indices): + """Fetch data. + + Args: + batched_indices (list): fetch data according to batched_indices + + """ + raise NotImplementedError + +class IterableFetcher(Fetcher): # pragma: no cover + """Iterate to get next batch-size samples as a batch.""" + + def __init__(self, dataset, collate_fn, drop_last, distributed): + """Initialize IterableFetcher. + + Args: + dataset (object): dataset object from which to get data + collate_fn (callable): merge data with outer dimension batch size + drop_last (bool): whether to drop the last batch if it is incomplete + distributed (bool): whether the dataloader is distributed + + """ + super(IterableFetcher, self).__init__(dataset, collate_fn, drop_last) + self.dataset_iter = iter(dataset) + self.index_whole = 0 + self.process_rank = 0 # The default rank is 0, which represents the main process + self.process_size = 1 # By default, process_size=1, only the main process is running + if distributed: + import horovod.tensorflow as hvd + hvd.init() + self.process_rank = hvd.rank() + self.process_size = hvd.size() + if self.process_size < 2: + raise EnvironmentError("The program is now trying to traverse" \ + " the distributed TensorFlow DefaultDataLoader in only one process." \ + " If you do not want to use distributed DataLoader, please set" \ + " 'distributed: False'. Or If you want to use distributed DataLoader," \ + " please set 'distributed: True' and launch multiple processes.") + + def __call__(self, batched_indices): + """Fetch data. + + Args: + batched_indices (list): fetch data according to batched_indices + + """ + batch_data = [] + batch_size = len(batched_indices) + while True: + try: + iter_data = next(self.dataset_iter) + if (self.index_whole-self.process_rank)%self.process_size == 0: + batch_data.append(iter_data) + self.index_whole += 1 + if len(batch_data) == batch_size: + break + except StopIteration: + break + if len(batch_data) == 0 or (self.drop_last and len(batch_data) < len(batched_indices)): + raise StopIteration + return self.collate_fn(batch_data) + +class IndexFetcher(Fetcher): # pragma: no cover + """Take single index or a batch of indices to fetch samples as a batch.""" + + def __init__(self, dataset, collate_fn, drop_last, distributed): + """Initialize IndexFetcher. + + Args: + dataset (object): dataset object from which to get data + collate_fn (callable): merge data with outer dimension batch size + drop_last (bool): whether to drop the last batch if it is incomplete + distributed (bool): whether the dataloader is distributed + """ + super(IndexFetcher, self).__init__(dataset, collate_fn, drop_last) + + def __call__(self, batched_indices): + """Fetch data. + + Args: + batched_indices (list): fetch data according to batched_indices + + """ + data = [self.dataset[idx] for idx in batched_indices] + return self.collate_fn(data) + +FETCHERS = {"index": IndexFetcher, "iter": IterableFetcher, } diff --git a/neural_compressor/data/dataloaders/mxnet_dataloader.py b/neural_compressor/data/dataloaders/mxnet_dataloader.py new file mode 100644 index 00000000000..352f63fc731 --- /dev/null +++ b/neural_compressor/data/dataloaders/mxnet_dataloader.py @@ -0,0 +1,42 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""MXNet Dataloader implementation.""" + +from neural_compressor.utils.utility import LazyImport +from .base_dataloader import BaseDataLoader +import logging +mx = LazyImport('mxnet') + +class MXNetDataLoader(BaseDataLoader): # pragma: no cover + """Subclass of BaseDataLoader.""" + def _generate_dataloader(self, dataset, batch_size, last_batch, collate_fn, + sampler, batch_sampler, num_workers, pin_memory, + shuffle, distributed): + """Overwrite _generate_dataloader function.""" + if shuffle: + logging.warning('Shuffle is not supported yet in MXNetDataLoader, ' \ + 'ignoring shuffle keyword.') + return mx.gluon.data.DataLoader( + dataset, + batch_size=batch_size, + batchify_fn=collate_fn, + last_batch=last_batch, + num_workers=num_workers, + pin_memory=pin_memory, + sampler=sampler, + batch_sampler=batch_sampler) diff --git a/neural_compressor/data/dataloaders/onnxrt_dataloader.py b/neural_compressor/data/dataloaders/onnxrt_dataloader.py new file mode 100644 index 00000000000..028bcdb6981 --- /dev/null +++ b/neural_compressor/data/dataloaders/onnxrt_dataloader.py @@ -0,0 +1,82 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Built-in dataloaders for onnxruntime framework backends.""" + +from neural_compressor.utils.utility import LazyImport +from .base_dataloader import BaseDataLoader +from .default_dataloader import DefaultDataLoader +from ..datasets.bert_dataset import ONNXRTBertDataset +import logging +torch = LazyImport('torch') + +class ONNXRTBertDataLoader(DefaultDataLoader): # pragma: no cover + """Built-in dataloader for onnx bert model and its varients.""" + + def _generate_dataloader(self, dataset, batch_size, last_batch, collate_fn, + sampler, batch_sampler, num_workers, pin_memory, + shuffle, distributed): + import numpy as np + from torch.utils.data import DataLoader, SequentialSampler + sampler = SequentialSampler(dataset) + dataloader = DataLoader(dataset, sampler=sampler, \ + batch_size=batch_size) + dynamic_length = dataset.dynamic_length + model_type = dataset.model_type + max_seq_length = dataset.max_seq_length + + for batch in dataloader: + try: + batch_seq_length = max_seq_length if not dynamic_length \ + else torch.max(batch[-2], 0)[0].item() + batch = tuple(t.detach().cpu().numpy() \ + if not isinstance(t, np.ndarray) else t \ + for t in batch) + if model_type == 'bert': + data = [ + batch[0][:,:batch_seq_length], + batch[1][:,:batch_seq_length], + batch[2][:,:batch_seq_length] + ] + else: + data = [ + batch[0][:,:batch_seq_length], + batch[1][:,:batch_seq_length] + ] + label = batch[-1] + yield data, label + except StopIteration: + return + +class ONNXRTDataLoader(BaseDataLoader): # pragma: no cover + """Built-in dataloader for onnxruntime framework backends.""" + + def _generate_dataloader(self, dataset, batch_size, last_batch, collate_fn, + sampler, batch_sampler, num_workers, pin_memory, + shuffle, distributed): + if shuffle: + logging.warning('Shuffle is not supported yet in ONNXRTDataLoader, ' \ + 'ignoring shuffle keyword.') + + if isinstance(dataset, ONNXRTBertDataset): + return ONNXRTBertDataLoader(dataset, batch_size, last_batch, collate_fn, + sampler, batch_sampler, num_workers, pin_memory, + shuffle, distributed) + else: + return DefaultDataLoader(dataset, batch_size, last_batch, collate_fn, + sampler, batch_sampler, num_workers, pin_memory, + shuffle, distributed) diff --git a/neural_compressor/data/dataloaders/pytorch_dataloader.py b/neural_compressor/data/dataloaders/pytorch_dataloader.py new file mode 100644 index 00000000000..301519a8acf --- /dev/null +++ b/neural_compressor/data/dataloaders/pytorch_dataloader.py @@ -0,0 +1,69 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Initialize the DATASETS class.""" + +import numpy as np +from neural_compressor.utils.utility import LazyImport +from .base_dataloader import BaseDataLoader +torch = LazyImport('torch') +hvd = LazyImport('horovod.torch') + +class PyTorchDataLoader(BaseDataLoader): # pragma: no cover + """PyTorchDataLoader inherits from BaseDataLoader.""" + + def _generate_dataloader(self, dataset, batch_size, last_batch, collate_fn, + sampler, batch_sampler, num_workers, pin_memory, + shuffle, distributed): + """Generate PyTorch dataloader. + + Args: + dataset: dataset + batch_size (int): batch size + last_batch (string): rollover last batch or not. + collate_fn: collate_fn + sampler: sampler + batch_sampler: batch_sampler + num_workers (int): num_workers + pin_memory (bool): pin_memory + shuffle (bool): shuffle + distributed (bool): distributed + + Returns: + _type_: _description_ + """ + drop_last = False if last_batch == 'rollover' else True + assert len(dataset) != 0, \ + "Warning: Dataset is empty, Please check dataset path!" + if distributed and sampler is None: + # TODO: lazy init here + hvd.init() + # sampler option is mutually exclusive with shuffle pytorch + self.sampler = sampler = torch.utils.data.distributed.DistributedSampler( + dataset, num_replicas=hvd.size(), rank=hvd.rank()) + + return torch.utils.data.DataLoader( + dataset, + shuffle=shuffle, + batch_size=batch_size, + collate_fn=collate_fn, + drop_last=drop_last, + num_workers=num_workers, + pin_memory=pin_memory, + sampler=sampler, + batch_sampler=batch_sampler) + diff --git a/neural_compressor/data/dataloaders/sampler.py b/neural_compressor/data/dataloaders/sampler.py new file mode 100644 index 00000000000..a383e6d9891 --- /dev/null +++ b/neural_compressor/data/dataloaders/sampler.py @@ -0,0 +1,136 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Definitions of the methods to sample data.""" + +from abc import abstractmethod + +class Sampler(object): # pragma: no cover + """Base class for all Samplers. + + __iter__ is needed no matter whether you use IterableSampler + or Squential sampler, if you want implement your own sampler, make clear what the type is + your Dataset, if IterableDataset(method __iter__ implemented), try to use IterableSampler, + else if you have an IndexDataset(method __getitem__ implemented), your dataset should have + method __len__ implemented. + """ + + def __init__(self, data_source): + """Initialize Sampler.""" + pass + + @abstractmethod + def __iter__(self): + """Convert dataloder to an iterator.""" + raise NotImplementedError + + +class IterableSampler(Sampler): # pragma: no cover + """Interally samples elements. + + Used for datasets retrieved element by interator. Yield None to act as a placeholder for each iteration. + """ + + def __init__(self, dataset): + """Initialize IterableSampler. + + Args: + dataset (object): dataset object from which to get data + """ + super(IterableSampler, self).__init__(None) + self.whole_dataset = dataset + + def __iter__(self): + """Yield data in iterative order.""" + while True: + yield None + + def __len__(self): + """Return the length of dataset.""" + raise NotImplementedError("'__len__' for IterableDataset object has not defined") + +class SequentialSampler(Sampler): # pragma: no cover + """Sequentially samples elements, used for datasets retrieved element by index.""" + + def __init__(self, dataset, distributed): + """Initialize SequentialSampler. + + Args: + dataset (object): dataset object from which to get data + distributed (bool): whether the dataloader is distributed + """ + self.whole_dataset = dataset + self.distributed = distributed + + def __iter__(self): + """Yield data in iterative order.""" + self.process_rank = 0 # The default rank is 0, which represents the main process + self.process_size = 1 # By default, process_size=1, only the main process is running + if self.distributed: + import horovod.tensorflow as hvd + hvd.init() + self.process_rank = hvd.rank() + self.process_size = hvd.size() + if self.process_size < 2: + raise EnvironmentError("The program is now trying to traverse" \ + " the distributed TensorFlow DefaultDataLoader in only one process." \ + " If you do not want to use distributed DataLoader, please set" \ + " 'distributed: False'. Or If you want to use distributed DataLoader," \ + " please set 'distributed: True' and launch multiple processes.") + return iter(range(self.process_rank, len(self.whole_dataset), self.process_size)) + + def __len__(self): + """Return the length of dataset.""" + return len(self.whole_dataset) + +class BatchSampler(Sampler): # pragma: no cover + """Yield a batch of indices and number of batches.""" + + def __init__(self, sampler, batch_size, drop_last=True): + """Initialize BatchSampler. + + Args: + sampler (Sampler): sampler used for generating batches + batch_size (int): size of batch + drop_last (bool, optional): whether to drop the last batch if it is incomplete. Defaults to True. + """ + if isinstance(drop_last, bool): + self.drop_last = drop_last + else: + raise ValueError("last_batch only support bool as input") + + self.sampler = sampler + self.batch_size = batch_size + self.drop_last = drop_last + + def __iter__(self): + """Yield data in iterative order.""" + batch = [] + for idx in self.sampler: + batch.append(idx) + if len(batch) == self.batch_size: + yield batch + batch = [] + if len(batch) > 0 and not self.drop_last: + yield batch + + def __len__(self): + """Return the number of batches.""" + if self.drop_last: + return len(self.sampler) // self.batch_size + else: + return (len(self.sampler) + self.batch_size - 1) // self.batch_size diff --git a/neural_compressor/data/dataloaders/tensorflow_dataloader.py b/neural_compressor/data/dataloaders/tensorflow_dataloader.py new file mode 100644 index 00000000000..ddc010841ac --- /dev/null +++ b/neural_compressor/data/dataloaders/tensorflow_dataloader.py @@ -0,0 +1,322 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""TensorFlow Dataloader implementation.""" + +from neural_compressor.experimental.data.datasets import dataset +from neural_compressor.utils.utility import LazyImport +from abc import abstractmethod +import collections +import numpy as np +import sys +from math import ceil, floor +from .sampler import IterableSampler, SequentialSampler, BatchSampler +from .fetcher import FETCHERS +from .default_dataloader import default_collate +from .default_dataloader import DefaultDataLoader +from ..datasets.bert_dataset import TensorflowBertDataset, TensorflowModelZooBertDataset +from .base_dataloader import BaseDataLoader +import logging + +tf = LazyImport('tensorflow') +neural_compressor = LazyImport('neural_compressor') + +class TFDataDataLoader(BaseDataLoader): # pragma: no cover + """Tensorflow dataloader class. + + In tensorflow1.x dataloader is coupled with the graph, but it also support feed_dict + method to do session run, this dataloader is designed to satisfy the usage of feed dict + in tf1.x. Although it's a general dataloader and can be used in MXNet and PyTorch. + + Args: + dataset: obj. wrapper of needed data. + batch_size: int. batch size + """ + + def __init__(self, dataset, batch_size=1, last_batch='rollover'): + """Initialize `TFDataDataLoader` class.""" + self.dataset = dataset + self.last_batch = last_batch + self._batch_size = batch_size + dataset = dataset.batch(batch_size) + + def batch(self, batch_size, last_batch='rollover'): + """Dataset return data per batch.""" + drop_last = False if last_batch == 'rollover' else True + self._batch_size = batch_size + self.dataset = self.dataset.batch(batch_size, drop_last) + + def __iter__(self): + """Iterate dataloader.""" + return self._generate_dataloader( + self.dataset, + batch_size=self.batch_size, + last_batch=self.last_batch,) + + def _generate_dataloader(self, dataset, batch_size=1, last_batch='rollover', \ + collate_fn=None, sampler=None, batch_sampler=None, \ + num_workers=None, pin_memory=None, shuffle=False, \ + distributed=False): + """Yield data.""" + drop_last = False if last_batch == 'rollover' else True + if shuffle: + logging.warning('Shuffle is not supported yet in TFDataLoader, ' \ + 'ignoring shuffle keyword.') + + def check_dynamic_shape(element_spec): + if isinstance(element_spec, collections.abc.Sequence): + return any([check_dynamic_shape(ele) for ele in element_spec]) + elif isinstance(element_spec, tf.TensorSpec): + return True if element_spec.shape.num_elements() is None else False + else: + raise ValueError('unrecognized element spec...') + + def squeeze_output(output): + if isinstance(output, collections.abc.Sequence): + return [squeeze_output(ele) for ele in output] + elif isinstance(output, np.ndarray): + return np.squeeze(output, axis=0) + else: + raise ValueError('not supported output format....') + + if tf.executing_eagerly(): + index = 0 + outputs = [] + for iter_tensors in dataset: + samples = [] + iter_inputs, iter_labels = iter_tensors[0],iter_tensors[1] + if isinstance(iter_inputs, tf.Tensor): + samples.append(iter_inputs.numpy()) + else: + samples.append((iter_input.numpy() for iter_input in iter_inputs)) + if isinstance(iter_labels,tf.Tensor): + samples.append(iter_labels.numpy()) + else: + samples.append([np.array(l) for l in iter_labels]) + index += 1 + outputs.append(samples) + if index == batch_size: + outputs = default_collate(outputs) + yield outputs + outputs = [] + index = 0 + if len(outputs) > 0: + outputs = default_collate(outputs) + yield outputs + else: + try_single_batch = check_dynamic_shape(dataset.element_spec) + dataset = dataset.batch(1 if try_single_batch else batch_size, drop_last) + ds_iterator = tf.compat.v1.data.make_one_shot_iterator(dataset) + iter_tensors = ds_iterator.get_next() + data_config = tf.compat.v1.ConfigProto() + data_config.use_per_session_threads = 1 + data_config.intra_op_parallelism_threads = 1 + data_config.inter_op_parallelism_threads = 16 + data_sess = tf.compat.v1.Session(config=data_config) + # pylint: disable=no-name-in-module + from tensorflow.python.framework.errors_impl import OutOfRangeError + while True: + if not try_single_batch: + try: + outputs = data_sess.run(iter_tensors) + yield outputs + except OutOfRangeError: + data_sess.close() + return + else: + try: + outputs = [] + for i in range(0, batch_size): + outputs.append(squeeze_output(data_sess.run(iter_tensors))) + outputs = default_collate(outputs) + yield outputs + except OutOfRangeError: + if len(outputs) == 0: + data_sess.close() + return + else: + outputs = default_collate(outputs) + yield outputs + data_sess.close() + return + +class TensorflowBertDataLoader(DefaultDataLoader): # pragma: no cover + """Subclass of DefaultDataLoader. + + this dataloader is designed to satisfy the usage of Bert models. + """ + + def _generate_dataloader(self, dataset, batch_size, last_batch, collate_fn, + sampler, batch_sampler, num_workers, pin_memory, shuffle, + distributed): + + if shuffle: + logging.warning('Shuffle is not supported yet in TensorflowBertDataLoader, ' \ + 'ignoring shuffle keyword.') + def bert_collate_fn(batch): + elem = batch[0] + return elem + drop_last = False if last_batch == 'rollover' else True + sampler = self._generate_sampler(dataset, distributed) + self.batch_sampler = BatchSampler(sampler, batch_size, drop_last) + self.fetcher = FETCHERS[self.dataset_type]\ + (dataset, bert_collate_fn, drop_last, distributed) + + for batched_indices in self.batch_sampler: + try: + data = self.fetcher(batched_indices) + yield (data[0], batch_size), data[1] + except StopIteration: + return + +class TensorflowModelZooBertDataLoader(DefaultDataLoader): # pragma: no cover + """Subclass of DefaultDataLoader. + + this dataloader is designed to satisfy the usage of Model Zoo Bert models. + """ + + def _generate_dataloader(self, dataset, batch_size, last_batch, collate_fn, + sampler, batch_sampler, num_workers, pin_memory, shuffle, + distributed): + + if shuffle: + logging.warning('Shuffle is not supported yet in TensorflowBertDataLoader, ' \ + 'ignoring shuffle keyword.') + def bert_collate_fn(batch): + input_ids = [] + input_mask = [] + segment_ids = [] + for elem in batch: + input_ids.append(elem[0][0][0]) + input_mask.append(elem[0][1][0]) + segment_ids.append(elem[0][2][0]) + inputs = [input_ids, input_mask, segment_ids] + return inputs, batch[0][1] + drop_last = False if last_batch == 'rollover' else True + sampler = self._generate_sampler(dataset, distributed) + self.batch_sampler = BatchSampler(sampler, batch_size, drop_last) + self.fetcher = FETCHERS[self.dataset_type]\ + (dataset, bert_collate_fn, drop_last, distributed) + + inputs = [] + for batched_indices in self.batch_sampler: + try: + data = self.fetcher(batched_indices) + yield data + except StopIteration: + return + +class TensorflowDataLoader(BaseDataLoader): # pragma: no cover + """DataLoader for framework Tensorflow. + + if it's a tf.data.Dataset we will directly use the dataloader in the other case + will use DefaultDataLoader instead. + """ + + def _generate_dataloader(self, dataset, batch_size, last_batch, collate_fn, \ + sampler, batch_sampler, num_workers, pin_memory, shuffle, distributed): + + if shuffle: + logging.warning('Shuffle is not supported yet in TensorflowDataLoader, ' \ + 'ignoring shuffle keyword.') + if isinstance(dataset, tf.data.Dataset): + if int(tf.__version__[0]) > 1: + has_batch = hasattr(dataset, '_batch_size') + else: + has_batch = hasattr(dataset._dataset, '_batch_size') + if has_batch: + raise TypeError(f"Parameter 'batch_size={batch_size}'" \ + " conflicts with 'tf.data.Dataset'," \ + f" because {dataset} is already a BatchDataset." \ + f" Please pass in 'tf.data.Dataset' without batch attributes.") + process_rank = 0 # The default rank is 0, which represents the main process + process_size = 1 # By default, process_size=1, only the main process is running + if self.distributed: + import horovod.tensorflow as hvd + hvd.init() + process_rank = hvd.rank() + process_size = hvd.size() + if process_size < 2: + raise EnvironmentError("The program is now trying to generate" \ + " the distributed TensorflowDataLoader in only one process." \ + " If you do not want to use distributed DataLoader, please set" \ + " 'distributed: False'. Or If you want to use distributed DataLoader," \ + " please set 'distributed: True' and launch multiple processes.") + dataset = dataset.shard(process_size, process_rank) + tf_dataloader = TFDataDataLoader(dataset, batch_size, last_batch=last_batch) + return tf_dataloader + elif isinstance(dataset, TensorflowBertDataset): + if distributed: + raise NotImplementedError("Distributed TensorflowBertDataLoader" \ + " is not yet supported, please set 'distributed: False'") + tf_bert_dataloader = TensorflowBertDataLoader(dataset, batch_size, \ + last_batch, collate_fn, sampler, batch_sampler, \ + num_workers, pin_memory, shuffle, distributed) + return tf_bert_dataloader + elif isinstance(dataset, TensorflowModelZooBertDataset): + if distributed: + raise NotImplementedError("Distributed TensorflowBertDataLoader" \ + " is not yet supported, please set 'distributed: False'") + tf_bert_dataloader = TensorflowModelZooBertDataLoader(dataset, batch_size, \ + last_batch, collate_fn, sampler, batch_sampler, \ + num_workers, pin_memory, shuffle, distributed) + return tf_bert_dataloader + else: + return DefaultDataLoader(dataset, batch_size, last_batch, collate_fn, + sampler, batch_sampler, num_workers, + pin_memory, shuffle, distributed) + + def __bool__(self): + """Judgement if the dataloader exists.""" + # workaround in assert dataloader which will overload __len__() without __bool__() + # provided. Calling __len__() in asserting is not supposed and may cause issues. + return True + + def __len__(self): + """Total number of dataset.""" + try: + dataset_len = self.dataset.__len__() + except (AttributeError, TypeError): + try: + dataset_len = 0 + for _ in self.dataset: + dataset_len += 1 + except RuntimeError: return sum([1 for _ in self]) + except Exception: + raise ValueError(f"{self.dataset} is invalid, {self.dataset}" \ + " does not support calculating the length of its dataloader") + process_rank = 0 # The default rank is 0, which represents the main process + process_size = 1 # By default, process_size=1, only the main process is running + if self.distributed: + import horovod.tensorflow as hvd + hvd.init() + process_rank = hvd.rank() + process_size = hvd.size() + if process_size < 2: + raise EnvironmentError("The program is now trying to get length of" \ + " the distributed TensorflowDataLoader in only one process." \ + " If you do not want to use distributed DataLoader, please set" \ + " 'distributed: False'. Or If you want to use distributed DataLoader," \ + " please set 'distributed: True' and launch multiple processes.") + if process_rank < (dataset_len % process_size): + self.dis_dataset_len = dataset_len // process_size + 1 + else: + self.dis_dataset_len = dataset_len // process_size + if self.drop_last == False: + dataloader_len = ceil(self.dis_dataset_len / self.batch_size) + else: + dataloader_len = floor(self.dis_dataset_len / self.batch_size) + return sys.maxsize if dataloader_len > sys.maxsize else dataloader_len diff --git a/neural_compressor/data/datasets/__init__.py b/neural_compressor/data/datasets/__init__.py index 77edda38b30..c2460d737ed 100644 --- a/neural_compressor/data/datasets/__init__.py +++ b/neural_compressor/data/datasets/__init__.py @@ -15,6 +15,9 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Built-in datasets class for multiple framework backends.""" + +from .dataset import Datasets, Dataset, IterableDataset, dataset_registry from os.path import dirname, basename, isfile, join import glob @@ -24,3 +27,5 @@ if isfile(f) and not f.startswith('__') and not f.endswith('__init__.py'): __import__(basename(f)[:-3], globals(), locals(), level=1) + +__all__ = ["Datasets", "Dataset", "IterableDataset", "dataset_registry"] diff --git a/neural_compressor/data/datasets/bert_dataset.py b/neural_compressor/data/datasets/bert_dataset.py new file mode 100644 index 00000000000..dc6cfc896b4 --- /dev/null +++ b/neural_compressor/data/datasets/bert_dataset.py @@ -0,0 +1,467 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Built-in BERT datasets class for multiple framework backends.""" + +import os +import logging +import json +import dataclasses +from dataclasses import dataclass +from typing import List, Optional, Union +from neural_compressor.utils.utility import LazyImport +from .dataset import dataset_registry, Dataset +torch = LazyImport('torch') +transformers = LazyImport('transformers') + +logger = logging.getLogger("neural_compressor") + +@dataset_registry(dataset_type="bert", framework="pytorch", dataset_format='') +class PytorchBertDataset(Dataset): # pragma: no cover + """PyTorch dataset used for model Bert. + + This Dataset is to construct from the Bert TensorDataset and not a full implementation + from yaml config. The original repo link is: https://github.com/huggingface/transformers. + When you want use this Dataset, you should add it before you initialize your DataLoader. + (TODO) add end to end support for easy config by yaml by adding the method of + load examples and process method. + + Args: dataset (list): list of data. + task (str): the task of the model, support "classifier", "squad". + model_type (str, default='bert'): model type, support 'distilbert', 'bert', + 'xlnet', 'xlm'. + transform (transform object, default=None): transform to process input data. + filter (Filter objects, default=None): filter out examples according + to specific conditions. + + Examples: + dataset = [[ + [101,2043,2001], + [1,1,1], + [[0,0,0,0,0,0,0], + [0,0,0,0,0,0,0], + [0,0,0,0,0,0,0]], + [1,1,1], + [1,1,1], + [[0,0,0,0,0,0,0], + [0,0,0,0,0,0,0], + [0,0,0,0,0,0,0]] + ]] + dataset = PytorchBertDataset(dataset=dataset, task='classifier', model_type='bert', + transform=preprocess, filter=filter) + """ + + def __init__(self, dataset, task, model_type='bert', transform=None, filter=None): + """Initialize the attributes of class.""" + self.dataset = dataset + assert task in ("classifier", "squad"), "Bert task support only classifier squad" + self.task = task + self.transform = transform + self.model_type = model_type + + def __len__(self): + """Length of the dataset.""" + return len(self.dataset) + + def __getitem__(self, index): + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + sample = self.dataset[index] + if self.transform is not None: + sample = self.transform(sample) + if self.task == 'classifier': + inputs = { + 'input_ids': sample[0], + 'attention_mask': sample[1], + 'labels': sample[3]} + + if self.model_type != 'distilbert': + # XLM, DistilBERT and RoBERTa don't use segment_ids + if self.model_type in ['bert', 'xlnet']: + inputs['token_type_ids'] = sample[2] + sample = (inputs, inputs['labels']) + + elif self.task == 'squad': + inputs = { + 'input_ids': sample[0], + 'attention_mask': sample[1], } + if self.model_type != 'distilbert': + # XLM, DistilBERT and RoBERTa don't use segment_ids + inputs['token_type_ids'] = sample[2] if self.model_type in [ + 'bert', 'xlnet'] else None + if self.model_type in ['xlnet', 'xlm']: + inputs.update({'cls_index': sample[4], 'p_mask': sample[5]}) + example_indices = sample[3] + sample = (inputs, example_indices) + return sample + + +@dataset_registry(dataset_type="GLUE", framework="onnxrt_qlinearops, \ + onnxrt_integerops", dataset_format='') +class ONNXRTBertDataset(Dataset): # pragma: no cover + """ONNXRT dataset used for model Bert. + + Args: data_dir (str): The input data dir. + model_name_or_path (str): Path to pre-trained student model or shortcut name, + selected in the list: + max_seq_length (int, default=128): The maximum length after tokenization. + Sequences longer than this will be truncated, + sequences shorter will be padded. + do_lower_case (bool, default=True): Whether to lowercase the input when tokenizing. + task (str, default=mrpc): The name of the task to fine-tune. + Choices include mrpc, qqp, qnli, rte, + sts-b, cola, mnli, wnli. + model_type (str, default='bert'): model type, support 'distilbert', 'bert', + 'mobilebert', 'roberta'. + dynamic_length (bool, default=False): Whether to use fixed sequence length. + evaluate (bool, default=True): Whether do evaluation or training. + transform (transform object, default=None): transform to process input data. + filter (Filter objects, default=None): filter out examples according + to specific conditions. + + Examples: + dataset = ONNXRTBertDataset(data_dir=data_dir, model_name_or_path='bert-base-uncase', + transform=preprocess, filter=filter) + """ + def __init__(self, data_dir, model_name_or_path, max_seq_length=128,\ + do_lower_case=True, task='mrpc', model_type='bert', dynamic_length=False,\ + evaluate=True, transform=None, filter=None): + """Initialize the attributes of class.""" + task = task.lower() + model_type = model_type.lower() + assert task in ['mrpc', 'qqp', 'qnli', 'rte', 'sts-b', 'cola', \ + 'mnli', 'wnli'], 'Unsupported task type' + assert model_type in ['distilbert', 'bert', 'mobilebert', 'roberta'], 'Unsupported \ + model type' + + self.dynamic_length = dynamic_length + self.model_type = model_type + self.max_seq_length = max_seq_length + tokenizer = transformers.AutoTokenizer.from_pretrained(model_name_or_path, + do_lower_case=do_lower_case) + self.dataset = load_and_cache_examples(data_dir, model_name_or_path, \ + max_seq_length, task, model_type, tokenizer, evaluate) + + def __len__(self): + """Length of the dataset.""" + return len(self.dataset) + + def __getitem__(self, index): + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + return self.dataset[index] + + +def load_and_cache_examples(data_dir, model_name_or_path, max_seq_length, task, \ + model_type, tokenizer, evaluate): # pragma: no cover + """Load and cache the examples. + + Helper Function for ONNXRTBertDataset. + """ + from torch.utils.data import TensorDataset + + processor = transformers.glue_processors[task]() + output_mode = transformers.glue_output_modes[task] + # Load data features from cache or dataset file + if not os.path.exists("./dataset_cached"): + os.makedirs("./dataset_cached") + cached_features_file = os.path.join("./dataset_cached", 'cached_{}_{}_{}_{}'.format( + 'dev' if evaluate else 'train', + list(filter(None, model_name_or_path.split('/'))).pop(), + str(max_seq_length), + str(task))) + if os.path.exists(cached_features_file): + logger.info("Load features from cached file {}.".format(cached_features_file)) + features = torch.load(cached_features_file) + else: + logger.info("Create features from dataset file at {}.".format(data_dir)) + label_list = processor.get_labels() + if task in ['mnli', 'mnli-mm'] and model_type in ['roberta']: + # HACK(label indices are swapped in RoBERTa pretrained model) + label_list[1], label_list[2] = label_list[2], label_list[1] + examples = processor.get_dev_examples(data_dir) if evaluate else \ + processor.get_train_examples(data_dir) + features = convert_examples_to_features(examples, + tokenizer, + task=task, + label_list=label_list, + max_length=max_seq_length, + output_mode=output_mode, + ) + logger.info("Save features into cached file {}.".format(cached_features_file)) + torch.save(features, cached_features_file) + # Convert to Tensors and build dataset + all_input_ids = torch.tensor([f.input_ids for f in features], dtype=torch.long) + all_attention_mask = torch.tensor([f.attention_mask for f in features], dtype=torch.long) + all_token_type_ids = torch.tensor([f.token_type_ids for f in features], dtype=torch.long) + all_seq_lengths = torch.tensor([f.seq_length for f in features], dtype=torch.long) + if output_mode == "classification": + all_labels = torch.tensor([f.label for f in features], dtype=torch.long) + elif output_mode == "regression": + all_labels = torch.tensor([f.label for f in features], dtype=torch.float) + dataset = TensorDataset(all_input_ids, all_attention_mask, all_token_type_ids, \ + all_seq_lengths, all_labels) + return dataset + + +def convert_examples_to_features( + examples, + tokenizer, + max_length=128, + task=None, + label_list=None, + output_mode="classification", + pad_token=0, + pad_token_segment_id=0, + mask_padding_with_zero=True, + ): # pragma: no cover + """Convert examples to features. + + Helper function for load_and_cache_examples. + """ + processor = transformers.glue_processors[task]() + if label_list is None: + label_list = processor.get_labels() + logger.info("Use label list {} for task {}.".format(label_list, task)) + label_map = {label: i for i, label in enumerate(label_list)} + features = [] + for (ex_index, example) in enumerate(examples): + inputs = tokenizer.encode_plus( + example.text_a, + example.text_b, + add_special_tokens=True, + max_length=max_length, + return_token_type_ids=True, + truncation=True, + ) + input_ids, token_type_ids = inputs["input_ids"], inputs["token_type_ids"] + # The mask has 1 for real tokens and 0 for padding tokens. Only real + # tokens are attended to. + attention_mask = [1 if mask_padding_with_zero else 0] * len(input_ids) + + # Zero-pad up to the sequence length. + seq_length = len(input_ids) + padding_length = max_length - len(input_ids) + + input_ids = input_ids + ([pad_token] * padding_length) + attention_mask = attention_mask + \ + ([0 if mask_padding_with_zero else 1] * padding_length) + token_type_ids = token_type_ids + ([pad_token_segment_id] * padding_length) + + assert len(input_ids) == max_length, \ + "Error with input_ids length {} vs {}".format( + len(input_ids), max_length) + assert len(attention_mask) == max_length, \ + "Error with attention_mask length {} vs {}".format( + len(attention_mask), max_length + ) + assert len(token_type_ids) == max_length, \ + "Error with token_type_ids length {} vs {}".format( + len(token_type_ids), max_length + ) + if output_mode == "classification": + label = label_map[example.label] + elif output_mode == "regression": + label = float(example.label) + else: + raise KeyError(output_mode) + + feats = InputFeatures( + input_ids=input_ids, + attention_mask=attention_mask, + token_type_ids=token_type_ids, + label=label, + seq_length=seq_length, + ) + features.append(feats) + return features + + +@dataclass(frozen=True) +class InputFeatures: # pragma: no cover + """Single set of features of data. + + Property names are the same names as the corresponding inputs to a model. + + Args: + input_ids: Indices of input sequence tokens in the vocabulary. + attention_mask: Mask to avoid performing attention on padding token indices. + Mask values selected in ``[0, 1]``: Usually ``1`` for tokens that are NOT MASKED, + ``0`` for MASKED (padded) tokens. + token_type_ids: (Optional) Segment token indices to indicate first and second + portions of the inputs. Only some models use them. + label: (Optional) Label corresponding to the input. Int for classification problems, + float for regression problems. + seq_length: (Optional) The length of input sequence before padding. + """ + + input_ids: List[int] + attention_mask: Optional[List[int]] = None + token_type_ids: Optional[List[int]] = None + label: Optional[Union[int, float]] = None + seq_length: Optional[List[int]] = None + + def to_json_string(self): + """Serialize this instance to a JSON string.""" + return json.dumps(dataclasses.asdict(self)) + "\n" + + +@dataset_registry(dataset_type="bert", framework="tensorflow, tensorflow_itex", dataset_format='') +class TensorflowBertDataset(Dataset): # pragma: no cover + """Tensorflow dataset used for model Bert. + + This dataset supports tfrecord data, please refer to Guide to create tfrecord file first. + + Args: root (str): path of dataset. + label_file (str): path of label file. + task (str, default='squad'): task type of model. + model_type (str, default='bert'): model type, support 'bert'. + transform (transform object, default=None): transform to process input data. + filter (Filter objects, default=None): filter out examples according + to specific conditions + """ + + def __init__(self, root, label_file, task='squad', + model_type='bert', transform=None, filter=None): + """Initialize the attributes of class.""" + import json + with open(label_file) as lf: + label_json = json.load(lf) + assert label_json['version'] == '1.1', 'only support squad 1.1' + self.label = label_json['data'] + self.root = root + self.transform = transform + self.filter = filter + + def __getitem__(self, index): + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index). + """ + return self.root, self.label + + def __len__(self): + """Length of the dataset.""" + return 1 + + +class ParseDecodeBert(): # pragma: no cover + """Helper function for TensorflowModelZooBertDataset. + + Parse the features from sample. + """ + + def __call__(self, sample): + """Parse the sample data. + + Args: + sample: Data to be parsed. + """ + import tensorflow as tf + # Dense features in Example proto. + feature_map = { + 'input_ids': + tf.compat.v1.VarLenFeature(dtype=tf.int64), + 'input_mask': + tf.compat.v1.VarLenFeature(dtype=tf.int64), + 'segment_ids': + tf.compat.v1.VarLenFeature(dtype=tf.int64), + } + + features = tf.io.parse_single_example(sample, feature_map) + + input_ids = features['input_ids'].values + input_mask = features['input_mask'].values + segment_ids = features['segment_ids'].values + + return (input_ids, input_mask, segment_ids) + +@dataset_registry(dataset_type="mzbert", framework="tensorflow, tensorflow_itex", dataset_format='') +class TensorflowModelZooBertDataset(Dataset): # pragma: no cover + """Tensorflow dataset for three-input Bert in tf record format. + + Root is a full path to tfrecord file, which contains the file name. + Please use Resize transform when batch_size > 1 + Args: root (str): path of dataset. + label_file (str): path of label file. + task (str, default='squad'): task type of model. + model_type (str, default='bert'): model type, support 'bert'. + transform (transform object, default=None): transform to process input data. + filter (Filter objects, default=None): filter out examples according. + """ + + def __init__(self, root, label_file, task='squad', + model_type='bert', transform=None, filter=None, num_cores=28): + """Initialize the attributes of class.""" + import json + with open(label_file) as lf: + label_json = json.load(lf) + assert label_json['version'] == '1.1', 'only support squad 1.1' + self.label = label_json['data'] + import tensorflow as tf + record_iterator = tf.compat.v1.python_io.tf_record_iterator(root) + example = tf.train.SequenceExample() + for element in record_iterator: + example.ParseFromString(element) + break + feature = example.context.feature + if len(feature['input_ids'].int64_list.value) == 0 \ + and len(feature['input_mask'].int64_list.value) == 0: + raise ValueError("Tfrecord format is incorrect, please refer\ + 'https://github.com/tensorflow/models/blob/master/research/\ + object_detection/dataset_tools/' to create correct tfrecord") + # pylint: disable=no-name-in-module + from tensorflow.python.data.experimental import parallel_interleave + tfrecord_paths = [root] + ds = tf.data.TFRecordDataset.list_files(tfrecord_paths) + ds = ds.apply( + parallel_interleave(tf.data.TFRecordDataset, + cycle_length=num_cores, + block_length=5, + sloppy=True, + buffer_output_elements=10000, + prefetch_input_elements=10000)) + if transform is not None: + transform.transform_list.insert(0, ParseDecodeBert()) + else: + transform = ParseDecodeBert() + ds = ds.map(transform, num_parallel_calls=None) + if filter is not None: + ds = ds.filter(filter) + ds = ds.prefetch(buffer_size=1000) + from ..dataloaders.tensorflow_dataloader import TFDataDataLoader + ds = TFDataDataLoader(ds) + self.root = [] + for inputs in ds: + self.root.append(inputs) + self.transform = transform + self.filter = filter + + def __getitem__(self, index): + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + return self.root[index], self.label + + def __len__(self): + """Length of the dataset.""" + return len(self.root) diff --git a/neural_compressor/data/datasets/coco_dataset.py b/neural_compressor/data/datasets/coco_dataset.py new file mode 100644 index 00000000000..650c0648bc8 --- /dev/null +++ b/neural_compressor/data/datasets/coco_dataset.py @@ -0,0 +1,301 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Copyright 2017 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== + +"""Built-in COCO datasets class for multiple framework backends.""" + +import numpy as np +from PIL import Image +from neural_compressor.utils.utility import LazyImport +from .dataset import dataset_registry, IterableDataset, Dataset + +tf = LazyImport('tensorflow') +mx = LazyImport('mxnet') +torch = LazyImport('torch') + +class ParseDecodeCoco(): # pragma: no cover + """Helper function for TensorflowModelZooBertDataset. + + Parse the features from sample. + """ + + def __call__(self, sample): + """Parse the sample data. + + Args: + sample: Data to be parsed. + """ + # Dense features in Example proto. + feature_map = { + 'image/encoded': + tf.compat.v1.FixedLenFeature([], dtype=tf.string, default_value=''), + 'image/object/class/text': + tf.compat.v1.VarLenFeature(dtype=tf.string), + 'image/object/class/label': + tf.compat.v1.VarLenFeature(dtype=tf.int64), + 'image/source_id':tf.compat.v1.FixedLenFeature([], dtype=tf.string, default_value=''), + } + sparse_float32 = tf.compat.v1.VarLenFeature(dtype=tf.float32) + # Sparse features in Example proto. + feature_map.update({ + k: sparse_float32 + for k in [ + 'image/object/bbox/xmin', 'image/object/bbox/ymin', + 'image/object/bbox/xmax', 'image/object/bbox/ymax' + ] + }) + + features = tf.io.parse_single_example(sample, feature_map) + + xmin = tf.expand_dims(features['image/object/bbox/xmin'].values, 0) + ymin = tf.expand_dims(features['image/object/bbox/ymin'].values, 0) + xmax = tf.expand_dims(features['image/object/bbox/xmax'].values, 0) + ymax = tf.expand_dims(features['image/object/bbox/ymax'].values, 0) + + bbox = tf.concat([ymin, xmin, ymax, xmax], 0) + # Force the variable number of bounding boxes into the shape + # [1, num_boxes, coords]. + bbox = tf.expand_dims(bbox, 0) + bbox = tf.transpose(bbox, [0, 2, 1]) + + encoded_image = features['image/encoded'] + image_tensor = tf.image.decode_image(encoded_image, channels=3) + image_tensor.set_shape([None, None, 3]) + + str_label = features['image/object/class/text'].values + int_label = features['image/object/class/label'].values + image_id = features['image/source_id'] + + return image_tensor, (bbox[0], str_label, int_label, image_id) + +@dataset_registry(dataset_type="COCORecord", framework="tensorflow, tensorflow_itex", dataset_format='') +class COCORecordDataset(IterableDataset): # pragma: no cover + """Tensorflow COCO dataset in tf record format. + + Root is a full path to tfrecord file, which contains the file name. + Please use Resize transform when batch_size > 1 + + Args: root (str): Root directory of dataset. + num_cores (int, default=28):The number of input Datasets to interleave from in parallel. + transform (transform object, default=None): transform to process input data. + filter (Filter objects, default=None): filter out examples according + to specific conditions. + """ + + def __new__(cls, root, num_cores=28, transform=None, filter=filter): + """Build a new object.""" + record_iterator = tf.compat.v1.python_io.tf_record_iterator(root) + example = tf.train.SequenceExample() + for element in record_iterator: + example.ParseFromString(element) + break + feature = example.context.feature + if len(feature['image/object/class/text'].bytes_list.value) == 0 \ + and len(feature['image/object/class/label'].int64_list.value) == 0: + raise ValueError("Tfrecord format is incorrect, please refer\ + 'https://github.com/tensorflow/models/blob/master/research/\ + object_detection/dataset_tools/create_coco_tf_record.py' to\ + create correct tfrecord") + # pylint: disable=no-name-in-module + from tensorflow.python.data.experimental import parallel_interleave + tfrecord_paths = [root] + ds = tf.data.TFRecordDataset.list_files(tfrecord_paths) + ds = ds.apply( + parallel_interleave(tf.data.TFRecordDataset, + cycle_length=num_cores, + block_length=5, + sloppy=True, + buffer_output_elements=10000, + prefetch_input_elements=10000)) + if transform is not None: + transform.transform_list.insert(0, ParseDecodeCoco()) + else: + transform = ParseDecodeCoco() + ds = ds.map(transform, num_parallel_calls=None) + if filter is not None: + ds = ds.filter(filter) + ds = ds.prefetch(buffer_size=1000) + return ds + +@dataset_registry(dataset_type="COCORaw", framework="onnxrt_qlinearops, \ + onnxrt_integerops, pytorch, mxnet, tensorflow, \ + tensorflow_itex", dataset_format='') +class COCORaw(Dataset): # pragma: no cover + """Coco raw dataset. + + Please arrange data in this way: + /root/img_dir/1.jpg + /root/img_dir/2.jpg + ... + /root/img_dir/n.jpg + /root/anno_dir + Please use Resize transform when batch_size > 1 + + Args: root (str): Root directory of dataset. + img_dir (str, default='val2017'): image file directory. + anno_dir (str, default='annotations/instances_val2017.json'): annotation file directory. + transform (transform object, default=None): transform to process input data. + filter (Filter objects, default=None): filter out examples according + to specific conditions. + """ + + def __init__(self, root, img_dir='val2017', \ + anno_dir='annotations/instances_val2017.json', transform=None, filter=filter): + """Initialize the attributes of class.""" + import json + import os + import numpy as np + from pycocotools.coco import COCO + self.image_list = [] + self.transform = transform + img_path = os.path.join(root, img_dir) + anno_path = os.path.join(root, anno_dir) + coco = COCO(anno_path) + img_ids = coco.getImgIds() + cat_ids = coco.getCatIds() + for idx, img_id in enumerate(img_ids): + img_info = {} + bboxes = [] + labels = [] + ids = [] + img_detail = coco.loadImgs(img_id)[0] + ids.append(img_detail['file_name'].encode('utf-8')) + pic_height = img_detail['height'] + pic_width = img_detail['width'] + + ann_ids = coco.getAnnIds(imgIds=img_id,catIds=cat_ids) + anns = coco.loadAnns(ann_ids) + for ann in anns: + bbox = ann['bbox'] + if len(bbox) == 0: + continue + bbox = [bbox[0]/float(pic_width), bbox[1]/float(pic_height),\ + bbox[2]/float(pic_width), bbox[3]/float(pic_height)] + bboxes.append([bbox[1], bbox[0], bbox[1]+bbox[3], bbox[0]+bbox[2]]) + labels.append(coco.cats[ann['category_id']]['name'].encode('utf8')) + img_file = os.path.join(img_path, img_detail['file_name']) + if not os.path.exists(img_file) or len(bboxes) == 0: + continue + + if filter and not filter(None, bboxes): + continue + + with Image.open(img_file) as image: + image = np.array(image.convert('RGB')) + self.image_list.append( + (image, [np.array(bboxes), np.array(labels), np.array([]),\ + np.array(img_detail['file_name'].encode('utf-8'))])) + + def __len__(self): + """Length of the dataset.""" + return len(self.image_list) + + def __getitem__(self, index): + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + sample = self.image_list[index] + if self.transform is not None: + sample= self.transform(sample) + return sample + +@dataset_registry(dataset_type="COCONpy", framework="onnxrt_qlinearops, \ + onnxrt_integerops, pytorch, mxnet, tensorflow, \ + tensorflow_itex", dataset_format='') +class COCONpy(Dataset): # pragma: no cover + """COCO npy dataset. + + Please arrange data in this way: + /root/npy_dir/1.jpg.npy + /root/npy_dir/2.jpg.npy + ... + /root/npy_dir/n.jpg.npy + /root/anno_dir + + Args: root (str): Root directory of dataset. + npy_dir (str, default='val2017'): npy file directory. + anno_dir (str, default='annotations/instances_val2017.json'): annotation file directory. + transform (transform object, default=None): transform to process input data. + filter (Filter objects, default=None): filter out examples according + to specific conditions. + """ + + def __init__(self, root, npy_dir='val2017', \ + anno_dir='annotations/instances_val2017.json', transform=None, filter=None): + """Initialize the attributes of class.""" + import json + import os + import numpy as np + from pycocotools.coco import COCO + self.image_list = [] + npy_path = os.path.join(root, npy_dir) + anno_path = os.path.join(root, anno_dir) + coco = COCO(anno_path) + img_ids = coco.getImgIds() + cat_ids = coco.getCatIds() + for idx, img_id in enumerate(img_ids): + img_info = {} + labels = [] + ids = [] + img_detail = coco.loadImgs(img_id)[0] + ids.append(img_detail['file_name'].encode('utf-8')) + pic_height = img_detail['height'] + pic_width = img_detail['width'] + + ann_ids = coco.getAnnIds(imgIds=img_id,catIds=cat_ids) + anns = coco.loadAnns(ann_ids) + for ann in anns: + bbox = ann['bbox'] + category_id = ann['category_id'] + if len(bbox) == 0: + continue + labels.append((np.array(category_id), np.array(bbox))) + npy_file = os.path.join(npy_path, img_detail['file_name']) + npy_file = npy_file + ".npy" + if not os.path.exists(npy_file): + continue + + image = np.load(npy_file) + self.image_list.append( + (image, labels)) + + def __len__(self): + """Length of the dataset.""" + return len(self.image_list) + + def __getitem__(self, index): + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + sample = self.image_list[index] + return sample diff --git a/neural_compressor/data/datasets/dataset.py b/neural_compressor/data/datasets/dataset.py new file mode 100644 index 00000000000..8cba94c54ca --- /dev/null +++ b/neural_compressor/data/datasets/dataset.py @@ -0,0 +1,1114 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""This is the base class for each framework.""" + +from abc import abstractmethod +import os +from neural_compressor.utils.utility import LazyImport, singleton +from PIL import Image +torch = LazyImport('torch') +torchvision = LazyImport('torchvision') +tf = LazyImport('tensorflow') +mx = LazyImport('mxnet') +np = LazyImport('numpy') +hashlib = LazyImport('hashlib') +gzip = LazyImport('gzip') +tarfile = LazyImport('tarfile') +zipfile = LazyImport('zipfile') +pickle = LazyImport('pickle') +glob = LazyImport('glob') + + +@singleton +class TensorflowDatasets(object): # pragma: no cover + """The base class of Tensorflow datasets class.""" + + def __init__(self): + """Initialize the attributes of class.""" + self.datasets = {} + self.datasets.update(TENSORFLOW_DATASETS) + + +@singleton +class PyTorchDatasets(object): # pragma: no cover + """The base class of PyTorch datasets class.""" + + def __init__(self): + """Initialize the attributes of class.""" + self.datasets = { + 'ImageFolder': PytorchMxnetWrapDataset( + torchvision.datasets.ImageFolder), + } + self.datasets.update(PYTORCH_DATASETS) + + +@singleton +class MXNetDatasets(object): # pragma: no cover + """The base class of MXNet datasets class.""" + + def __init__(self): + """Initialize the attributes of class.""" + self.datasets = {} + self.datasets.update(MXNET_DATASETS) + + +@singleton +class ONNXRTQLDatasets(object): # pragma: no cover + """The base class of ONNXRT QLinear datasets class.""" + + def __init__(self): + """Initialize the attributes of class.""" + self.datasets = {} + self.datasets.update(ONNXRTQL_DATASETS) + + +@singleton +class ONNXRTITDatasets(object): # pragma: no cover + """The base class of ONNXRT IT datasets class.""" + + def __init__(self): + """Initialize the attributes of class.""" + self.datasets = {} + self.datasets.update(ONNXRTIT_DATASETS) + + +class PytorchMxnetWrapDataset(): # pragma: no cover + """The base class for PyTorch and MXNet frameworks. + + Args: + datafunc: The datasets class of PyTorch or MXNet. + """ + + def __init__(self, datafunc): + """Initialize the attributes of class.""" + self.datafunc = datafunc + + def __call__(self, transform=None, filter=None, *args, **kwargs): + """Wrap the dataset for PyTorch and MXNet framework.""" + return PytorchMxnetWrapFunction(self.datafunc, transform=transform, \ + filter=filter, *args, **kwargs) + + +class PytorchMxnetWrapFunction(): # pragma: no cover + """The Helper class for PytorchMxnetWrapDataset. + + Args: + dataset (datasets class): The datasets class of PyTorch or MXNet. + transform (transform object): transform to process input data. + filter (Filter objects): filter out examples according to specific + conditions. + """ + + def __init__(self, dataset, transform, filter, *args, **kwargs): + """Initialize the attributes of class.""" + self.dataset = dataset(*args, **kwargs) + self.transform = transform + self.filter = filter + + def __len__(self): + """Length of the dataset.""" + return len(self.dataset) + + def __getitem__(self, index): + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + sample = self.dataset[index] + if self.transform is not None: + sample = self.transform(sample) + return sample + + +framework_datasets = {"tensorflow": TensorflowDatasets, + "tensorflow_itex": TensorflowDatasets, + "mxnet": MXNetDatasets, + "pytorch": PyTorchDatasets, + "pytorch_ipex": PyTorchDatasets, + "pytorch_fx": PyTorchDatasets, + "onnxrt_qdq": ONNXRTQLDatasets, + "onnxrt_qlinearops": ONNXRTQLDatasets, + "onnxrt_qoperator": ONNXRTQLDatasets, + "onnxrt_integerops": ONNXRTITDatasets, + } + +"""The datasets supported by neural_compressor, it's model specific and can be configured by yaml file. + + User could add new datasets by implementing new Dataset subclass under this directory. + The naming convention of new dataset subclass should be something like ImageClassifier, user + could choose this dataset by setting "imageclassifier" string in tuning.strategy field of yaml. + + Datasets variable is used to store all implemented Dataset subclasses to support + model specific dataset. +""" + + +class Datasets(object): # pragma: no cover + """A base class for all framework datasets. + + Args: + framework (str): framework name, like:"tensorflow", "tensorflow_itex", + "mxnet", "onnxrt_qdq", "onnxrt_qlinearops", "onnxrt_integerops", + "pytorch", "pytorch_ipex", "pytorch_fx", "onnxrt_qoperator". + """ + + def __init__(self, framework): + """Initialize the attributes of class.""" + assert framework in ["tensorflow", "tensorflow_itex", \ + "mxnet", "onnxrt_qdq", "onnxrt_qlinearops", "onnxrt_integerops", \ + "pytorch", "pytorch_ipex", "pytorch_fx", "onnxrt_qoperator"], \ + "framework support tensorflow pytorch mxnet onnxrt" + self.datasets = framework_datasets[framework]().datasets + + def __getitem__(self, dataset_type): + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + assert dataset_type in self.datasets.keys(), "dataset type only support {}".\ + format(self.datasets.keys()) + return self.datasets[dataset_type] + + +# user/model specific datasets will be registered here +TENSORFLOW_DATASETS = {} +TENSORFLOWITEX_DATASETS = {} +MXNET_DATASETS = {} +PYTORCH_DATASETS = {} +PYTORCHIPEX_DATASETS = {} +PYTORCHFX_DATASETS = {} +ONNXRTQL_DATASETS = {} +ONNXRTIT_DATASETS = {} + +registry_datasets = {"tensorflow": TENSORFLOW_DATASETS, + "tensorflow_itex": TENSORFLOWITEX_DATASETS, + "mxnet": MXNET_DATASETS, + "pytorch": PYTORCH_DATASETS, + "pytorch_ipex": PYTORCHIPEX_DATASETS, + "pytorch_fx": PYTORCHFX_DATASETS, + "onnxrt_integerops": ONNXRTIT_DATASETS, + "onnxrt_qdq": ONNXRTQL_DATASETS, + "onnxrt_qoperator": ONNXRTQL_DATASETS, + "onnxrt_qlinearops": ONNXRTQL_DATASETS, + } + + +def dataset_registry(dataset_type, framework, dataset_format=''): # pragma: no cover + """Register dataset subclasses. + + Args: + 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: + cls: The class of register. + """ + def decorator_dataset(cls): + for single_framework in [fwk.strip() for fwk in framework.split(',')]: + assert single_framework in [ + "tensorflow", + "tensorflow_itex", + "mxnet", + "pytorch", + "pytorch_ipex", + "pytorch_fx", + "onnxrt_qlinearops", + "onnxrt_integerops", + "onnxrt_qdq", + "onnxrt_qoperator", + ], "The framework support tensorflow mxnet pytorch onnxrt" + dataset_name = dataset_type + dataset_format + if dataset_name in registry_datasets[single_framework].keys(): + raise ValueError('Cannot have two datasets with the same name') + registry_datasets[single_framework][dataset_name] = cls + return cls + return decorator_dataset + + +class Dataset(object): # pragma: no cover + """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 + """ + + @abstractmethod + def __getitem__(self, index): + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + raise NotImplementedError + + # it's suggested to implement your __len__ method though we do not set it in abstract class + # @abstractmethod + # def __len__(self): + # raise NotImplementedError + + +class IterableDataset(object): # pragma: no cover + """An iterable Dataset. + + Subclass iterable dataset should also implement a method: + `__iter__` for interating over the samples of the dataset. + """ + + @abstractmethod + def __iter__(self): + """Magic method. + + Returns the iterator object itself. + """ + raise NotImplementedError + + +def download_url(url, root, filename=None, md5=None): # pragma: no cover + """Download from url. + + Args: + url (str): the address to download from. + root (str): the path for saving. + filename (str): the file name for saving. + md5 (str): the md5 string. + """ + import urllib + root = os.path.expanduser(root) + if not filename: + filename = os.path.basename(url) + fpath = os.path.join(root, filename) + + os.makedirs(root, exist_ok=True) + + if check_integrity(fpath, md5): + print('Using downloaded and verified file: ' + fpath) + else: + try: + print('Downloading ' + url + ' to ' + fpath) + urllib.request.urlretrieve( + url, fpath, + reporthook=gen_bar_updater() + ) + except (urllib.error.URLError, IOError) as e: + if url[:5] == 'https': + url = url.replace('https:', 'http:') + print('Failed download. Trying https -> http instead.' + ' Downloading ' + url + ' to ' + fpath) + urllib.request.urlretrieve( + url, fpath, + reporthook=gen_bar_updater() + ) + else: + raise e + if not check_integrity(fpath, md5): + raise RuntimeError("File not found or corrupted.") + + +def gen_bar_updater(): # pragma: no cover + """Generate progress bar.""" + from tqdm import tqdm + pbar = tqdm(total=None) + + def bar_update(count, block_size, total_size): + """Update progress bar.""" + if pbar.total is None and total_size: + pbar.total = total_size + progress_bytes = count * block_size + pbar.update(progress_bytes - pbar.n) + return bar_update + + +def check_integrity(fpath, md5): # pragma: no cover + """Check MD5 checksum.""" + if not os.path.isfile(fpath): + return False + if md5 is None: + return True + return md5 == calculate_md5(fpath) + + +def calculate_md5(fpath, chunk_size=1024*1024): # pragma: no cover + """Generate MD5 checksum for a file.""" + md5 = hashlib.md5() + with open(fpath, 'rb') as f: + for chunk in iter(lambda: f.read(chunk_size), b''): + md5.update(chunk) + return md5.hexdigest() + +@dataset_registry(dataset_type="CIFAR10", framework="onnxrt_qlinearops, \ + onnxrt_integerops", dataset_format='') +class CIFAR10(Dataset): # pragma: no cover + """The CIFAR10 and CIFAR100 database. + + For CIFAR10: If download is True, it will download dataset to root/ and extract it + automatically, otherwise user can download file from + https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz manually to + root/ and extract it. + For CIFAR100: If download is True, it will download dataset to root/ and extract it + automatically, otherwise user can download file from + https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz manually to + root/ and extract it. + + Args: + root (str): Root directory of dataset. + train (bool, default=False): If True, creates dataset from train subset, + otherwise from validation subset. + transform (transform object, default=None): transform to process input data. + filter (Filter objects, default=None): filter out examples according to specific + conditions. + download (bool, default=True): If true, downloads the dataset from the internet + and puts it in root directory. If dataset is already + downloaded, it is not downloaded again. + """ + + url = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz" + filename = "cifar-10-python.tar.gz" + tgz_md5 = 'c58f30108f718f92721af3b95e74349a' + train_list = [ + ['data_batch_1', 'c99cafc152244af753f735de768cd75f'], + ['data_batch_2', 'd4bba439e000b95fd0a9bffe97cbabec'], + ['data_batch_3', '54ebc095f3ab1f0389bbae665268c751'], + ['data_batch_4', '634d18415352ddfa80567beed471001a'], + ['data_batch_5', '482c414d41f54cd18b22e5b47cb7c3cb'], + ] + + test_list = [ + ['test_batch', '40351d587109b95175f43aff81a1287e'], + ] + + meta = { + 'filename': 'batches.meta', + 'key': 'label_names', + 'md5': '5ff9c542aee3614f3951f8cda6e48888', + } + + def __init__(self, + root, + train=False, + transform=None, + filter=None, + download=True): # pragma: no cover + """Initialize the attributes of class.""" + self.root = root + if download: + self.download() + + if not self._check_integrity(): + raise RuntimeError( + 'Dataset not found or corrupted. You can use download=True to download it') + if train: + downloaded_list = self.train_list + else: + downloaded_list = self.test_list + + self.data = [] + self.targets = [] + for file_name, checksum in downloaded_list: + file_path = os.path.join(self.root, file_name) + with open(file_path, 'rb') as f: + entry = pickle.load(f, encoding='latin1') + self.data.append(entry['data']) + if 'labels' in entry: + self.targets.extend(entry['labels']) + else: + self.targets.extend(entry['fine_labels']) + self.data = np.vstack(self.data).reshape(-1, 3, 32, 32) + self.data = self.data.transpose((0, 2, 3, 1)) # convert to HWC + + self.load_meta() + self.transform = transform + + def load_meta(self): # pragma: no cover + """Load meta.""" + path = os.path.join(self.root, self.meta['filename']) + if not check_integrity(path, self.meta['md5']): + raise RuntimeError('Dataset metadata file not found or corrupted.' + + ' You can use download=True to download it') + with open(path, 'rb') as infile: + data = pickle.load(infile, encoding='latin1') + self.classes = data[self.meta['key']] + self.class_to_idx = {_class: i for i, _class in enumerate(self.classes)} + + def __getitem__(self, index): # pragma: no cover + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + image, label = self.data[index], self.targets[index] + if self.transform is not None: + image, label = self.transform((image, label)) + return image, label + + def __len__(self): # pragma: no cover + """Length of the dataset.""" + return len(self.data) + + def download(self): # pragma: no cover + """Download a file.""" + if self._check_integrity(): + print('Files already downloaded and verified') + return + download_root = os.path.expanduser(self.root) + filename = os.path.basename(self.url) + download_url(self.url, download_root, filename, self.tgz_md5) + archive = os.path.join(download_root, filename) + print("Extracting {} to {}".format(archive, download_root)) + with tarfile.open(archive, 'r:gz') as tar: + tar.extractall(path=download_root) + + def _check_integrity(self): # pragma: no cover + """Check MD5 checksum.""" + root = self.root + for fentry in (self.train_list + self.test_list): + filename, md5 = fentry[0], fentry[1] + fpath = os.path.join(root, filename) + if not check_integrity(fpath, md5): + return False + return True + + +@dataset_registry(dataset_type="CIFAR10", framework="pytorch", dataset_format='') +class PytorchCIFAR10(CIFAR10): + """The PyTorch datasets for CIFAR10.""" + + def __getitem__(self, index): # pragma: no cover + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + image, label = self.data[index], self.targets[index] + image = Image.fromarray(image) + if self.transform is not None: + image, label = self.transform((image, label)) + return (image, label) + + +@dataset_registry(dataset_type="CIFAR10", framework="mxnet", dataset_format='') +class MXNetCIFAR10(CIFAR10): + """The MXNet datasets for CIFAR10.""" + + def __getitem__(self, index): # pragma: no cover + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + image, label = self.data[index], self.targets[index] + image = mx.nd.array(image) + if self.transform is not None: + image, label = self.transform((image, label)) + return (image, label) + + +@dataset_registry(dataset_type="CIFAR10", framework="tensorflow, tensorflow_itex", dataset_format='') +class TensorflowCIFAR10(CIFAR10): + """The Tensorflow datasets for CIFAR10.""" + + def __getitem__(self, index): # pragma: no cover + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + image, label = self.data[index], self.targets[index] + if self.transform is not None: + image, label = self.transform((image, label)) + if type(image).__name__ == 'Tensor': + with tf.compat.v1.Session() as sess: + image = sess.run(image) + elif type(image).__name__ == 'EagerTensor': + image = image.numpy() + return (image, label) + + +@dataset_registry(dataset_type="CIFAR100", framework="onnxrt_qlinearops, \ + onnxrt_integerops", dataset_format='') +class CIFAR100(CIFAR10): + """CIFAR100 database. + + For CIFAR100: If download is True, it will download dataset to root/ and extract it + automatically, otherwise user can download file from + https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz manually to + root/ and extract it. + + Args: + root (str): Root directory of dataset. + train (bool, default=False): If True, creates dataset from train subset, + otherwise from validation subset. + transform (transform object, default=None): transform to process input data. + filter (Filter objects, default=None): filter out examples according to specific + conditions. + download (bool, default=True): If true, downloads the dataset from the internet + and puts it in root directory. If dataset is already + downloaded, it is not downloaded again. + """ + + url = "https://www.cs.toronto.edu/~kriz/cifar-100-python.tar.gz" + filename = "cifar-100-python.tar.gz" + tgz_md5 = 'eb9058c3a382ffc7106e4002c42a8d85' + train_list = [ + ['train', '16019d7e3df5f24257cddd939b257f8d'], + ] + test_list = [ + ['test', 'f0ef6b0ae62326f3e7ffdfab6717acfc'], + ] + meta = { + 'filename': 'meta', + 'key': 'fine_label_names', + 'md5': '7973b15100ade9c7d40fb424638fde48', + } + + +@dataset_registry(dataset_type="CIFAR100", framework="pytorch", dataset_format='') +class PytorchCIFAR100(CIFAR100): + """The PyTorch datasets for CIFAR100.""" + + def __getitem__(self, index): # pragma: no cover + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + image, label = self.data[index], self.targets[index] + image = Image.fromarray(image) + if self.transform is not None: + image, label = self.transform((image, label)) + image = np.array(image) + return (image, label) + + +@dataset_registry(dataset_type="CIFAR100", framework="mxnet", dataset_format='') +class MXNetCIFAR100(CIFAR100): + """The MXNet datasets for CIFAR100.""" + + def __getitem__(self, index): # pragma: no cover + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + image, label = self.data[index], self.targets[index] + image = mx.nd.array(image) + if self.transform is not None: + image, label = self.transform((image, label)) + return (image, label) + + +@dataset_registry(dataset_type="CIFAR100", framework="tensorflow, tensorflow_itex", dataset_format='') +class TensorflowCIFAR100(CIFAR100): + """The Tensorflow datasets for CIFAR100.""" + + def __getitem__(self, index): # pragma: no cover + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + image, label = self.data[index], self.targets[index] + if self.transform is not None: + image, label = self.transform((image, label)) + if type(image).__name__ == 'Tensor': + with tf.compat.v1.Session() as sess: + image = sess.run(image) + elif type(image).__name__ == 'EagerTensor': + image = image.numpy() + return (image, label) + +@dataset_registry(dataset_type="MNIST", framework="onnxrt_qlinearops, \ + onnxrt_integerops", dataset_format='') +class MNIST(Dataset): # pragma: no cover + """Modified National Institute of Standards and Technology database and FashionMNIST database. + + For MNIST: If download is True, it will download dataset to root/MNIST/, otherwise user + should put mnist.npz under root/MNIST/ manually. + For FashionMNIST: If download is True, it will download dataset to root/FashionMNIST/, + otherwise user should put train-labels-idx1-ubyte.gz, + train-images-idx3-ubyte.gz, t10k-labels-idx1-ubyte.gz + and t10k-images-idx3-ubyte.gz under root/FashionMNIST/ manually. + + Args: + root (str): Root directory of dataset. + train (bool, default=False): If True, creates dataset from train subset, + otherwise from validation subset. + transform (transform object, default=None): transform to process input data. + filter (Filter objects, default=None): filter out examples according to specific + conditions. + download (bool, default=True): If true, downloads the dataset from the internet + and puts it in root directory. If dataset is already + downloaded, it is not downloaded again. + """ + + classes = ['0 - zero', '1 - one', '2 - two', '3 - three', '4 - four', + '5 - five', '6 - six', '7 - seven', '8 - eight', '9 - nine'] + resource = [ + ('https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz', + '8a61469f7ea1b51cbae51d4f78837e45') + ] + + def __init__(self, root, train=False, transform=None, filter=None, download=True): + """Initialize the attributes of class.""" + self.root = root + self.train = train + self.transform = transform + if download: + self.download() + + self.read_data() + + def read_data(self): + """Read data from a file.""" + for file_name, checksum in self.resource: + file_path = os.path.join(self.root, os.path.basename(file_name)) + if not os.path.exists(file_path): + raise RuntimeError( + 'Dataset not found. You can use download=True to download it') + with np.load(file_path, allow_pickle=True) as f: + if self.train: + self.data, self.targets = f['x_train'], f['y_train'] + else: + self.data, self.targets = f['x_test'], f['y_test'] + + def __len__(self): + """Length of the dataset.""" + return len(self.data) + + def __getitem__(self, index): + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + image, label = self.data[index], int(self.targets[index]) + image = np.expand_dims(image, -1) + if self.transform is not None: + image, label = self.transform((image, label)) + return image, label + + @property + def class_to_idx(self): + """Return a dict of class.""" + return {_class: i for i, _class in enumerate(self.classes)} + + def download(self): + """Download a file.""" + for url, md5 in self.resource: + filename = os.path.basename(url) + if os.path.exists(os.path.join(self.root, filename)): + continue + else: + download_url(url, root=self.root, + filename=filename, md5=md5) + + +@dataset_registry(dataset_type="MNIST", framework="pytorch", dataset_format='') +class PytorchMNIST(MNIST): # pragma: no cover + """The PyTorch datasets for MNIST.""" + + def __getitem__(self, index): + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + image, label = self.data[index], int(self.targets[index]) + image = Image.fromarray(image, mode='L') + if self.transform is not None: + image, label = self.transform((image, label)) + image = np.array(image) + return (image, label) + + +@dataset_registry(dataset_type="MNIST", framework="mxnet", dataset_format='') +class MXNetMNIST(MNIST): # pragma: no cover + """The MXNet datasets for MNIST.""" + + def __getitem__(self, index): + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + image, label = self.data[index], int(self.targets[index]) + image = mx.nd.array(image) + image = image.reshape((image.shape[0], image.shape[1], 1)) + if self.transform is not None: + image, label = self.transform((image, label)) + return (image, label) + + +@dataset_registry(dataset_type="MNIST", framework="tensorflow, tensorflow_itex", dataset_format='') +class TensorflowMNIST(MNIST): # pragma: no cover + """The Tensorflow datasets for MNIST.""" + + def __getitem__(self, index): + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + image, label = self.data[index], int(self.targets[index]) + image = np.expand_dims(image, -1) + if self.transform is not None: + image, label = self.transform((image, label)) + if type(image).__name__ == 'Tensor': + with tf.compat.v1.Session() as sess: + image = sess.run(image) + elif type(image).__name__ == 'EagerTensor': + image = image.numpy() + return (image, label) + + +@dataset_registry(dataset_type="FashionMNIST", framework="onnxrt_qlinearops, \ + onnxrt_integerops", dataset_format='') +class FashionMNIST(MNIST): # pragma: no cover + """FashionMNIST database. + + For FashionMNIST: If download is True, it will download dataset to root/FashionMNIST/, + otherwise user should put train-labels-idx1-ubyte.gz, + train-images-idx3-ubyte.gz, t10k-labels-idx1-ubyte.gz + and t10k-images-idx3-ubyte.gz under root/FashionMNIST/ manually. + + Args: + root (str): Root directory of dataset. + train (bool, default=False): If True, creates dataset from train subset, + otherwise from validation subset. + transform (transform object, default=None): transform to process input data. + filter (Filter objects, default=None): filter out examples according to specific + conditions. + download (bool, default=True): If true, downloads the dataset from the internet + and puts it in root directory. If dataset is already + downloaded, it is not downloaded again. + """ + + resource = [ + ('https://storage.googleapis.com/tensorflow/tf-keras-datasets/' + file_name, None) + for file_name in [ + 'train-labels-idx1-ubyte.gz', 'train-images-idx3-ubyte.gz', + 't10k-labels-idx1-ubyte.gz', 't10k-images-idx3-ubyte.gz' + ] + ] + + classes = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', + 'Shirt', 'Sneaker', 'Bag', 'Ankle boot'] + + def read_data(self): + """Read data from a file.""" + import struct + if self.train: + label_path = os.path.join(self.root, 'train-labels-idx1-ubyte.gz') + image_path = os.path.join(self.root, 'train-images-idx3-ubyte.gz') + else: + label_path = os.path.join(self.root, 't10k-labels-idx1-ubyte.gz') + image_path = os.path.join(self.root, 't10k-images-idx3-ubyte.gz') + with gzip.open(label_path, 'rb') as f: + struct.unpack(">II", f.read(8)) + self.targets = np.frombuffer(f.read(), dtype=np.uint8).astype(np.int32) + with gzip.open(image_path, 'rb') as f: + struct.unpack(">IIII", f.read(16)) + data = np.frombuffer(f.read(), dtype=np.uint8) + self.data = data.reshape(len(self.targets), 28, 28) + + +@dataset_registry(dataset_type="FashionMNIST", framework="pytorch", dataset_format='') +class PytorchFashionMNIST(FashionMNIST): # pragma: no cover + """The PyTorch datasets for FashionMNIST.""" + + def __getitem__(self, index): + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + image, label = self.data[index], int(self.targets[index]) + image = Image.fromarray(image, mode='L') + if self.transform is not None: + image, label = self.transform((image, label)) + image = np.array(image) + return (image, label) + + +@dataset_registry(dataset_type="FashionMNIST", framework="mxnet", dataset_format='') +class MXNetFashionMNIST(FashionMNIST): # pragma: no cover + """The MXNet Dataset for FashionMNIST.""" + + def __getitem__(self, index): + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + image, label = self.data[index], int(self.targets[index]) + image = mx.nd.array(image) + image = image.reshape((image.shape[0], image.shape[1], 1)) + if self.transform is not None: + image, label = self.transform((image, label)) + return (image, label) + + +@dataset_registry(dataset_type="FashionMNIST", framework="tensorflow, tensorflow_itex", dataset_format='') +class TensorflowFashionMNIST(FashionMNIST): # pragma: no cover + """The Tensorflow Dataset for FashionMNIST.""" + + def __getitem__(self, index): + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + image, label = self.data[index], int(self.targets[index]) + image = np.expand_dims(image, -1) + if self.transform is not None: + image, label = self.transform((image, label)) + if type(image).__name__ == 'Tensor': + with tf.compat.v1.Session() as sess: + image = sess.run(image) + elif type(image).__name__ == 'EagerTensor': + image = image.numpy() + return (image, label) + + +@dataset_registry(dataset_type="ImageFolder", framework="onnxrt_qlinearops, \ + onnxrt_integerops", dataset_format='') +class ImageFolder(Dataset): # pragma: no cover + """The base class for ImageFolder. + + Expects the data folder to contain subfolders representing the classes to which + its images belong. + + Please arrange data in this way: + root/class_1/xxx.png + root/class_1/xxy.png + root/class_1/xxz.png + ... + root/class_n/123.png + root/class_n/nsdf3.png + root/class_n/asd932_.png + Please put images of different categories into different folders. + + Args: root (str): Root directory of dataset. + transform (transform object, default=None): transform to process input data. + filter (Filter objects, default=None): filter out examples according to specific + conditions. + """ + + def __init__(self, root, transform=None, filter=None): + """Initialize the attributes of class.""" + self.root = root + assert os.path.exists(self.root), "Datapath doesn't exist!" + + self.transform = transform + self.image_list = [] + files = glob.glob(os.path.join(self.root, '*')) + files.sort() + for idx, file in enumerate(files): + imgs = glob.glob(os.path.join(file, '*')) + imgs.sort() + for img in imgs: + self.image_list.append((img, idx)) + + def __len__(self): + """Length of the dataset.""" + return len(self.image_list) + + def __getitem__(self, index): + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + sample = self.image_list[index] + label = sample[1] + with Image.open(sample[0]) as image: + image = np.array(image) + if self.transform is not None: + image, label = self.transform((image, label)) + return (image, label) + + +@dataset_registry(dataset_type="ImageFolder", framework="mxnet", dataset_format='') +class MXNetImageFolder(ImageFolder): # pragma: no cover + """The MXNet Dataset for image folder. + + Expects the data folder to contain subfolders representing the classes to which + its images belong. + + Please arrange data in this way: + root/class_1/xxx.png + root/class_1/xxy.png + root/class_1/xxz.png + ... + root/class_n/123.png + root/class_n/nsdf3.png + root/class_n/asd932_.png + Please put images of different categories into different folders. + + Args: root (str): Root directory of dataset. + transform (transform object, default=None): transform to process input data. + filter (Filter objects, default=None): filter out examples according to specific + conditions. + """ + + def __getitem__(self, index): + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + sample = self.image_list[index] + label = sample[1] + image = mx.image.imread(sample[0]) + if self.transform is not None: + image, label = self.transform((image, label)) + return (image, label) + + +@dataset_registry(dataset_type="ImageFolder", framework="tensorflow, tensorflow_itex", dataset_format='') +class TensorflowImageFolder(ImageFolder): # pragma: no cover + """The Tensorflow Dataset for image folder. + + Expects the data folder to contain subfolders representing the classes to which + its images belong. + + Please arrange data in this way: + root/class_1/xxx.png + root/class_1/xxy.png + root/class_1/xxz.png + ... + root/class_n/123.png + root/class_n/nsdf3.png + root/class_n/asd932_.png + Please put images of different categories into different folders. + + Args: root (str): Root directory of dataset. + transform (transform object, default=None): transform to process input data. + filter (Filter objects, default=None): filter out examples according to specific + conditions. + """ + + def __getitem__(self, index): + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + sample = self.image_list[index] + label = sample[1] + with Image.open(sample[0]) as image: + if image.mode != 'RGB': + image = image.convert('RGB') + image = np.array(image) + if self.transform is not None: + image, label = self.transform((image, label)) + if type(image).__name__ == 'Tensor': + with tf.compat.v1.Session() as sess: + image = sess.run(image) + elif type(image).__name__ == 'EagerTensor': + image = image.numpy() + return (image, label) + + +@dataset_registry(dataset_type="TFRecordDataset", framework="tensorflow, tensorflow_itex", dataset_format='') +class TensorflowTFRecordDataset(IterableDataset): # pragma: no cover + """The Tensorflow TFRecord Dataset. + + Root is a full path to tfrecord file, which contains the file name. + + Args: root (str): filename of dataset. + transform (transform object, default=None): transform to process input data. + filter (Filter objects, default=None): filter out examples according + to specific conditions. + """ + + def __new__(cls, root, transform=None, filter=None): + """Build a new object of TensorflowTFRecordDataset class.""" + # pylint: disable=no-name-in-module + from tensorflow.python.data.experimental import parallel_interleave + from tensorflow.python.platform import gfile + file_names = gfile.Glob(root) + ds = tf.data.Dataset.from_tensor_slices(file_names) + ds = ds.apply(parallel_interleave( + tf.data.TFRecordDataset, cycle_length=len(file_names))) + if transform is not None: + ds = ds.map(transform, num_parallel_calls=None) + ds = ds.prefetch(buffer_size=tf.data.experimental.AUTOTUNE) # this number can be tuned + return ds + + +@dataset_registry(dataset_type="ImageRecord", framework="tensorflow, tensorflow_itex", dataset_format='') +class TensorflowImageRecord(IterableDataset): # pragma: no cover + """Tensorflow imageNet database in tf record format. + + Please arrange data in this way: + root/validation-000-of-100 + root/validation-001-of-100 + ... + root/validation-099-of-100 + The file name needs to follow this pattern: '* - * -of- *' + + Args: root (str): Root directory of dataset. + transform (transform object, default=None): transform to process input data. + filter (Filter objects, default=None): filter out examples according + to specific conditions. + """ + + """Configuration for Imagenet dataset.""" + def __new__(cls, root, transform=None, filter=None): + """Build a new object of TensorflowImageRecord class.""" + from tensorflow.python.platform import gfile # pylint: disable=no-name-in-module + glob_pattern = os.path.join(root, '*-*-of-*') + file_names = gfile.Glob(glob_pattern) + if not file_names: + raise ValueError('Found no files in --root matching: {}'.format(glob_pattern)) + + # pylint: disable=no-name-in-module + from tensorflow.python.data.experimental import parallel_interleave + from neural_compressor.data.transforms.imagenet_transform import ParseDecodeImagenet + ds = tf.data.TFRecordDataset.list_files(file_names, shuffle=False) + ds = ds.apply(parallel_interleave( + tf.data.TFRecordDataset, cycle_length=len(file_names))) + + if transform is not None: + transform.transform_list.insert(0, ParseDecodeImagenet()) + else: + transform = ParseDecodeImagenet() + ds = ds.map(transform, num_parallel_calls=None) + ds = ds.prefetch(buffer_size=tf.data.experimental.AUTOTUNE) # this number can be tuned + return ds + + +@dataset_registry(dataset_type="VOCRecord", framework="tensorflow, tensorflow_itex", dataset_format='') +class TensorflowVOCRecord(IterableDataset): # pragma: no cover + """The Tensorflow PASCAL VOC 2012 database in tf record format. + + Please arrange data in this way: + root/val-00000-of-00004.tfrecord + root/val-00001-of-00004.tfrecord + ... + root/val-00003-of-00004.tfrecord + The file name needs to follow this pattern: 'val-*-of-*' + + Args: root (str): Root directory of dataset. + transform (transform object, default=None): transform to process input data. + filter (Filter objects, default=None): filter out examples according + to specific conditions. + """ + + def __new__(cls, root, transform=None, filter=None): + """Build a new object of TensorflowVOCRecord class.""" + from tensorflow.python.platform import gfile # pylint: disable=no-name-in-module + glob_pattern = os.path.join(root, '%s-*' % 'val') + file_names = gfile.Glob(glob_pattern) + if not file_names: + raise ValueError('Found no files in --root matching: {}'.format(glob_pattern)) + + # pylint: disable=no-name-in-module + from tensorflow.python.data.experimental import parallel_interleave + ds = tf.data.TFRecordDataset.list_files(file_names, shuffle=False) + ds = ds.apply(parallel_interleave( + tf.data.TFRecordDataset, cycle_length=len(file_names))) + + if transform is not None: + ds = ds.map(transform, num_parallel_calls=None) + ds = ds.prefetch(buffer_size=tf.data.experimental.AUTOTUNE) # this number can be tuned + return ds diff --git a/neural_compressor/data/datasets/dummy_dataset.py b/neural_compressor/data/datasets/dummy_dataset.py new file mode 100644 index 00000000000..8b54296015d --- /dev/null +++ b/neural_compressor/data/datasets/dummy_dataset.py @@ -0,0 +1,143 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Dummy dataset for dummy data generation on multiple framework backends.""" + +from .dataset import dataset_registry, Dataset +import numpy as np +from neural_compressor.utils.utility import LazyImport +import logging + +mx = LazyImport('mxnet') +torch = LazyImport('torch') + +logger = logging.getLogger("neural_compressor") + +@dataset_registry(dataset_type="dummy", framework="tensorflow, tensorflow_itex, \ + onnxrt_qlinearops, onnxrt_integerops, \ + pytorch, pytorch_ipex, pytorch_fx, \ + mxnet", + dataset_format='') +class DummyDataset(Dataset): # pragma: no cover + """Dataset used for dummy data generation. + + This Dataset is to construct a dataset from a specific shape. + The value range is calculated from: low * stand_normal(0, 1) + high. + (TODO) construct dummy data from real dataset or iteration of data. + """ + + def __init__(self, shape, low=-128., high=127., dtype='float32', label=True, \ + transform=None, filter=None): + """Initialize `DummyDataset` class. + + Args: + shape (list or tuple): Support create multi shape tensors, use list of tuples + for each tuple in the list, will create a such size tensor. + low (list or float, default=-128.): Low out the tensor value range from [0, 1] + to [0, low] or [low, 0] if low < 0, if float, will implement all tensors with same low value. + high (list or float, default=127.): High the tensor value by add all tensor element + value high. If list, length of list should be same with shape list. + dtype (list or str, default='float32'): Support multi tensor dtype setting. + If list, length of list should be same with shape list. If str, all tensors will + use same dtype. dtype supports 'float32', 'float16', 'uint8', 'int8', 'int32', 'int64', 'bool'. + label (bool, default=True): Whether to return 0 as label. + transform (transform object, default=None): Dummy dataset does not need transform. + If transform is not None, it will ignore it. + filter (Filter objects, default=None): Filter out examples according to specific conditions. + """ + dtype_map = {'float32':np.float32, 'float16':np.float16, 'uint8':np.uint8, \ + 'int8': np.int8, 'int32':np.int32, 'int64':np.int64, 'bool':bool,\ + 'string': str} + + np.random.seed(9527) + self.transform = transform + self.label = label + if len(shape)==0: + logger.info("No data in the dummy dataset.") + elif isinstance(shape, list): + # list tensor should same first demension n + n = shape[0][0] + assert all(isinstance(elem, tuple) and elem[0] == n for elem in shape), \ + 'each tensor shape should be tuple and same fisrt demension' + + if isinstance(low, list): + assert len(low) == len(shape) and all(isinstance(elem, float) for elem in low), \ + 'low list should have same length with shape with element data type float' + else: + low = (low * np.ones(len(shape))).astype(float) + + if isinstance(high, list): + assert len(high) == len(shape) and all(isinstance(elem, float) for elem in high), \ + 'high list should have same length with shape with element data type float' + else: + high = (high * np.ones(len(shape))).astype(float) + + if isinstance(dtype, list): + assert len(dtype) == len(shape) and \ + all(elem in dtype_map.keys() for elem in dtype), \ + 'high list should have same length with shape with element data type float' + else: + dtype = [dtype for i in range(0, len(shape))] + + elif isinstance(shape, tuple): + shape = [shape] + if isinstance(low, float): + low = [low] + else: + assert isinstance(low, list) and len(low) == 1 and isinstance(low[0], float), \ + 'low should be float or list of float with length 1' + + if isinstance(high, float): + high = [high] + else: + assert isinstance(high, list) and len(high) == 1 and isinstance(high[0], float), \ + 'high should be float or list of float with length 1' + + if isinstance(dtype, str): + assert dtype in dtype_map.keys(), 'dtype only support {}'.format(dtype_map.keys()) + dtype = [dtype] + else: + assert isinstance(dtype, list) and \ + len(dtype) == 1 and dtype[0] in dtype_map.keys(), \ + 'dtype should be str or list of str in supported dtypes' + + self.dataset = [] + for idx in range(0, len(shape)): + tensor = np.random.uniform(low=low[idx], high=high[idx], size=shape[idx]) + tensor = tensor.astype(dtype_map[dtype[idx]]) + self.dataset.append(tensor) + + if len(self.dataset) == 1: + self.dataset = self.dataset[0] + else: + self.dataset = [elem for elem in zip(*self.dataset)] + + + def __len__(self): + """Return the length of dataset.""" + return len(self.dataset) + + def __getitem__(self, index): + """Return the item of dataset according to the given index.""" + sample = self.dataset[index] + if self.transform is not None: + logger.warning("Dummy dataset does not need transform.") + + if self.label: + return sample, 0 + else: + return sample diff --git a/neural_compressor/data/datasets/dummy_dataset_v2.py b/neural_compressor/data/datasets/dummy_dataset_v2.py new file mode 100644 index 00000000000..46cf3bba73d --- /dev/null +++ b/neural_compressor/data/datasets/dummy_dataset_v2.py @@ -0,0 +1,275 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Dummy dataset for dummy_v2/sparse_dummy_v2 data generation on multiple framework backends.""" + +import sys +from .dataset import dataset_registry, IterableDataset +import numpy as np +from neural_compressor.utils.utility import LazyImport +from functools import reduce + +mx = LazyImport('mxnet') +torch = LazyImport('torch') + +@dataset_registry(dataset_type="dummy_v2", framework="tensorflow, tensorflow_itex, \ + onnxrt_qlinearops, onnxrt_integerops, \ + pytorch, pytorch_ipex, pytorch_fx, mxnet", + dataset_format='') +class DummyDataset(IterableDataset): # pragma: no cover + """Dataset used for dummy_v2 data generation. + + This Dataset is to construct a dataset from a input shape and label shape. + The value range is calculated from: low * stand_normal(0, 1) + high. + """ + + def __init__(self, input_shape, label_shape=None, low=-128., high=127., \ + dtype='float32', transform=None, filter=None): + """Initialize `DummyDataset` class. + + Args: + sample_size (int): Total size of the dummy samples. + input_shape (list or tuple): Create single or multi input tensors, + tuple reperesent the sample shape of the dataset, e.g. an image size should be + represented as (224, 224, 3), list contains multiple tuple and represent multi input tensors. + label_shape (list or tuple): Create single or multi label tensors, + tuple reperesent the label shape of the dataset, e.g. an label size should be + represented as (1, ), list contains multiple tuple and represent multi label tensors. + low (list or float, default=-128.): Low out the tensor value range from [0, 1] + to [0, low] or [low, 0] if low < 0. If float, will implement all tensors with same low value. + high (list or float, default=127.): High the tensor value by add all tensor element value high. + If list, length of list should be same with shape list. + dtype (list or str, default='float32'): Support multi tensor dtype setting. + If list, length of list should be same with shape list. + If str, all tensors will use same dtype. + dtype supports 'float32', 'float16', 'uint8', 'int8','int32', 'int64', 'bool'. + transform (transform object, default=None): dummy_v2 dataset does not need transform. + If transform is not None, it will ignore it. + filter (Filter objects, default=None): Filter out examples according to specific conditions. + """ + self.dtype_map = {'float32':np.float32, 'float16':np.float16, 'uint8':np.uint8, \ + 'int8':np.int8, 'int32':np.int32, 'int64':np.int64, 'bool':np.bool} + + np.random.seed(9527) + self.transform = transform + self.input_shape = input_shape + self.label_shape = label_shape + self.low = low + self.high = high + self.dtype = dtype + + if label_shape is None: + self.label_dim = 0 + elif isinstance(label_shape, tuple): + self.label_dim = 1 + else: + self.label_dim = len(label_shape) + + self.input_dim = 1 if isinstance(input_shape, tuple) else len(input_shape) + self.total_dim = self.input_dim + self.label_dim + + if isinstance(high, list): + assert len(high) == self.total_dim and \ + all(isinstance(elem, float) for elem in high),\ + 'high value list length should same with label dim + input_dim' + else: + self.high = (high * np.ones(self.total_dim)).astype(np.float) + + if isinstance(low, list): + assert len(low) == self.total_dim and \ + all(isinstance(elem, float) for elem in low), \ + 'low value list length should same with label dim + input_dim' + else: + self.low = (low * np.ones(self.total_dim)).astype(np.float) + + if isinstance(dtype, list): + assert len(dtype) == self.total_dim and \ + all(elem in self.dtype_map.keys() for elem in dtype), \ + 'dtype list length should same with label dim + input_dim' + else: + self.dtype = [self.dtype for i in range(0, self.total_dim)] + + if isinstance(input_shape, tuple): + self.input_shape = [input_shape] + + if isinstance(label_shape, tuple): + self.label_shape = [label_shape] + + def __iter__(self): + """Yield data in iterative order.""" + while True: + input_data = [] + for idx in range(0, self.input_dim): + tensor = np.random.uniform(\ + low=self.low[idx], high=self.high[idx], size=self.input_shape[idx]) + tensor = tensor.astype(self.dtype_map[self.dtype[idx]]) + input_data.append(tensor) + + label = [] + for idx in range(0, self.label_dim): + shift_idx = self.input_dim + idx + tensor = np.random.uniform(low=self.low[shift_idx], + high=self.high[shift_idx], + size=self.label_shape[idx]) + tensor = tensor.astype(self.dtype_map[self.dtype[shift_idx]]) + label.append(tensor) + + if len(input_data) == 1: + input_data = input_data[0] + + if len(label) == 1: + label = label[0] + + if len(label) > 0: + yield input_data, label + else: + yield input_data + + def __len__(self): + """Return the length of dataset.""" + return sys.maxsize + +@dataset_registry(dataset_type="sparse_dummy_v2", framework="tensorflow, tensorflow_itex, \ + onnxrt_qlinearops, onnxrt_integerops, \ + pytorch, pytorch_ipex, pytorch_fx, mxnet", + dataset_format='') +class SparseDummyDataset(IterableDataset): # pragma: no cover + """Dataset used for sparse_dummy_v2 data generation. + + This Dataset is to construct a dataset from a input shape and label shape. + The value range is calculated from: low * stand_normal(0, 1) + high. + """ + + def __init__(self, dense_shape, label_shape=None, sparse_ratio=0.5, low=-128., high=127., \ + dtype='float32', transform=None, filter=None): + """Initialize `SparseDummyDataset` class. + + Args: + sample_size (int): Total size of the dummy samples. + dense_shape (list or tuple): Create single or multi sparse tensors, tuple reperesent + the sample shape of the dataset, e.g. an image size should be represented as (224, 224, 3), + list contains multiple tuple and represent multi input tensors. + label_shape (list or tuple): Create single or multi label tensors, tuple reperesent + the label shape of the dataset, e.g. an label size should be represented as (1, ), + list contains multiple tuple and represent multi label tensors. + sparse_ratio (float, default=0.5): The ratio of sparsity, supports [0, 1]. + low (list or float, default=-128.): Low out the tensor value range from [0, 1] + to [0, low] or [low, 0] if low < 0. If float, will implement all tensors with same low value. + high (list or float, default=127.): High the tensor value by add all tensor element value high. + If list, length of list should be same with shape list. + dtype (list or str, default='float32'): Support multi tensor dtype setting. If list, + length of list should be same with shape list. If str, all tensors will use same dtype. + dtype supports 'float32', 'float16', 'uint8', 'int8', 'int32', 'int64', 'bool'. + transform (transform object, default=None): dummy_v2 dataset does not need transform. + If transform is not None, it will ignore it. + filter (Filter objects, default=None): Filter out examples according to specific conditions. + """ + self.dtype_map = {'float32':np.float32, 'float16':np.float16, 'uint8':np.uint8, \ + 'int8':np.int8, 'int32':np.int32, 'int64':np.int64, 'bool':np.bool} + + np.random.seed(9527) + self.transform = transform + self.dense_shape = dense_shape + self.label_shape = label_shape + self.sparse_ratio = sparse_ratio + self.low = low + self.high = high + self.dtype = dtype + + if isinstance(dense_shape, tuple): + self.dense_shape = [dense_shape] + + if label_shape is None: + self.label_dim = 0 + else: + if isinstance(label_shape, tuple): + self.label_shape = [label_shape] + if len(self.label_shape) == 1 and len(self.label_shape) != len(self.dense_shape): + self.label_shape = len(self.dense_shape) * self.label_shape + assert len(self.label_shape) == len(self.dense_shape), \ + 'length of dense_shape should be euqal to length of label_shape' + self.label_dim = len(self.label_shape) + + self.input_dim = 1 if isinstance(dense_shape, tuple) else len(dense_shape) + self.total_dim = self.input_dim + self.label_dim + + if isinstance(sparse_ratio, list): + assert len(sparse_ratio) == self.input_dim and \ + all(isinstance(elem, float) for elem in sparse_ratio),\ + 'sparse_ratio list length should same with input_dim' + else: + self.sparse_ratio = (sparse_ratio * np.ones(self.input_dim)).astype(np.float) + assert all([0 <= i <= 1 for i in self.sparse_ratio]), 'sparse_ratio should be in [0,1]' + + if isinstance(high, list): + assert len(high) == self.total_dim and \ + all(isinstance(elem, float) for elem in high),\ + 'high value list length should same with label dim + input_dim' + else: + self.high = (high * np.ones(self.total_dim)).astype(np.float) + + if isinstance(low, list): + assert len(low) == self.total_dim and \ + all(isinstance(elem, float) for elem in low), \ + 'low value list length should same with label dim + input_dim' + else: + self.low = (low * np.ones(self.total_dim)).astype(np.float) + + if isinstance(dtype, list): + assert len(dtype) == self.total_dim and \ + all(elem in self.dtype_map.keys() for elem in dtype), \ + 'dtype list length should same with label dim + input_dim' + else: + self.dtype = [self.dtype for i in range(0, self.total_dim)] + + def __iter__(self): + """Yield data in iterative order.""" + while True: + input_data = [] + for idx, shape in enumerate(self.dense_shape): + dim = len(shape) + total = reduce(lambda x, y: x*y, shape) + sparse_num = round(total * (1 - self.sparse_ratio[idx])) + val = np.random.uniform(\ + low=self.low[idx], high=self.high[idx], size=sparse_num) + val = val.astype(self.dtype_map[self.dtype[idx]]) + nums = np.arange(sparse_num) + indices = [] + dim_shape = [reduce(lambda x, y: x*y, shape[i:])/shape[i] \ + for i in range(len(shape))] + for num in nums: + indice = [] + for item in dim_shape: + indice.append(num//item) + num = num - indice[-1] * item if num - indice[-1] * item > 0 else num + indices.append(indice) + + if self.label_dim > 0: + shift_idx = self.input_dim + idx + tensor = np.random.uniform(low=self.low[shift_idx], + high=self.high[shift_idx], + size=self.label_shape[idx]) + tensor = tensor.astype(self.dtype_map[self.dtype[shift_idx]]) + input_data.append([(np.array(indices), val), tensor]) + else: + input_data.append((np.array(indices), val)) + + yield input_data + + def __len__(self): + """Return the length of dataset.""" + return sys.maxsize diff --git a/neural_compressor/data/datasets/imagenet_dataset.py b/neural_compressor/data/datasets/imagenet_dataset.py index 3cf944ad2ef..9d0d7daf2d1 100644 --- a/neural_compressor/data/datasets/imagenet_dataset.py +++ b/neural_compressor/data/datasets/imagenet_dataset.py @@ -29,20 +29,131 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== +"""Dataset for ImageNet data generation on multiple framework backends.""" + import os +import re +import numpy as np from PIL import Image from neural_compressor.utils.utility import LazyImport from neural_compressor.utils import logger -from neural_compressor.experimental.data.datasets import dataset_registry, IterableDataset, Dataset +from .dataset import dataset_registry, IterableDataset, Dataset tf = LazyImport('tensorflow') +mx = LazyImport('mxnet') +torch = LazyImport('torch') + +@dataset_registry(dataset_type="ImagenetRaw", framework="onnxrt_qlinearops, \ + onnxrt_integerops", dataset_format='') +class ImagenetRaw(Dataset): # pragma: no cover + """Configuration for ImageNet raw dataset. + + Please arrange data in this way: + data_path/img1.jpg + data_path/img2.jpg + ... + data_path/imgx.jpg + dataset will read name and label of each image from image_list file, + if user set image_list to None, it will read from data_path/val_map.txt automatically. + """ + + def __init__(self, data_path, image_list, transform=None, filter=None): + """Initialize `ImagenetRaw` class. + + Args: + data_path (str): Root directory of dataset. + image_list (str): Data file, record image_names and their labels. + transform (transform object, default=None): Transform to process input data. + filter (Filter objects, default=None): Filter out examples according to specific conditions. + """ + self.image_list = [] + self.label_list = [] + self.data_path = data_path + self.transform = transform + not_found = 0 + if image_list is None: + # by default look for val.txt + image_list = os.path.join(data_path, "val.txt") + + with open(image_list, 'r') as f: + for s in f: + image_name, label = re.split(r"\s+", s.strip()) + src = os.path.join(data_path, image_name) + if not os.path.exists(src): + # if the image does not exists ignore it + not_found += 1 + continue + self.image_list.append(src) + self.label_list.append(int(label)) + + if not self.image_list: + raise ValueError("no images in image list found") + if not_found > 0: + print("reduced image list, %d images not found", not_found) + + def __getitem__(self, index): + """Return the item of dataset according to the given index.""" + image_path, label = self.image_list[index], self.label_list[index] + with Image.open(image_path) as image: + image = np.array(image.convert('RGB')) + if self.transform is not None: + image, label = self.transform((image, label)) + return (image, label) + + def __len__(self): + """Return the length of dataset.""" + return len(self.image_list) + +@dataset_registry(dataset_type="ImagenetRaw", framework="pytorch", dataset_format='') +class PytorchImagenetRaw(ImagenetRaw): # pragma: no cover + """Dataset for ImageNet data generation on pytorch backend.""" + + def __getitem__(self, index): + """Return the item of dataset according to the given index.""" + image_path, label = self.image_list[index], self.label_list[index] + with Image.open(image_path) as image: + image = image.convert('RGB') + if self.transform is not None: + image, label = self.transform((image, label)) + image = np.array(image) + return (image, label) + +@dataset_registry(dataset_type="ImagenetRaw", framework="mxnet", dataset_format='') +class MXNetImagenetRaw(ImagenetRaw): # pragma: no cover + """Dataset for ImageNet data generation on mxnet backend.""" + + def __getitem__(self, index): + """Return the item of dataset according to the given index.""" + image_path, label = self.image_list[index], self.label_list[index] + image = mx.image.imread(image_path) + if self.transform is not None: + image, label = self.transform((image, label)) + return (image, label) + +@dataset_registry(dataset_type="ImagenetRaw", framework="tensorflow, \ + tensorflow_itex", dataset_format='') +class TensorflowImagenetRaw(ImagenetRaw): # pragma: no cover + """Dataset for ImageNet data generation on tensorflow/inteltensorflow/tensorflow_itex backend.""" + + def __getitem__(self, index): + """Return the item of dataset according to the given index.""" + image_path, label = self.image_list[index], self.label_list[index] + with Image.open(image_path) as image: + image = np.array(image.convert('RGB')) + if self.transform is not None: + image, label = self.transform((image, label)) + if type(image).__name__ == 'Tensor': + with tf.compat.v1.Session() as sess: + image = sess.run(image) + elif type(image).__name__ == 'EagerTensor': + image = image.numpy() + return (image, label) -# BELOW API TO BE DEPRECATED! @dataset_registry(dataset_type="Imagenet", framework="tensorflow", dataset_format='') -class TensorflowImagenetDataset(IterableDataset): +class TensorflowImagenetDataset(IterableDataset): # pragma: no cover """Configuration for Imagenet dataset.""" def __new__(cls, root, subset='validation', num_cores=28, transform=None, filter=None): - + """New a imagenet dataset for tensorflow.""" assert subset in ('validation', 'train'), \ 'only support subset (validation, train)' logger.warning("This api is going to be deprecated, " @@ -55,7 +166,7 @@ def __new__(cls, root, subset='validation', num_cores=28, transform=None, filter raise ValueError('Found no files in --root matching: {}'.format(glob_pattern)) from tensorflow.python.data.experimental import parallel_interleave - from neural_compressor.experimental.data.transforms.imagenet_transform import ParseDecodeImagenet + from neural_compressor.data.transforms.imagenet_transform import ParseDecodeImagenet ds = tf.data.TFRecordDataset.list_files(file_names, shuffle=False) ds = ds.apply( parallel_interleave( @@ -72,10 +183,11 @@ def __new__(cls, root, subset='validation', num_cores=28, transform=None, filter @dataset_registry(dataset_type="Imagenet", framework="onnxrt_qlinearops, \ onnxrt_integerops", dataset_format='') -class ONNXRTImagenetDataset(Dataset): +class ONNXRTImagenetDataset(Dataset): # pragma: no cover """Configuration for Imagenet dataset.""" def __init__(self, root, subset='val', num_cores=28, transform=None, filter=None): + """Initialize `ONNXRTImagenetDataset` class.""" self.val_dir = os.path.join(root, subset) assert os.path.exists(self.val_dir), "find no val dir in {}".format(root) + \ "please make sure there are train/val subfolders" @@ -93,9 +205,11 @@ def __init__(self, root, subset='val', num_cores=28, transform=None, filter=None self.image_list.append((img, idx)) def __len__(self): + """Return the number of images.""" return len(self.image_list) def __getitem__(self, index): + """Return the item of dataset according to the given index.""" from PIL import Image sample = self.image_list[index] image = Image.open(sample[0]) diff --git a/neural_compressor/data/datasets/style_transfer_dataset.py b/neural_compressor/data/datasets/style_transfer_dataset.py new file mode 100644 index 00000000000..8f6f6ff332f --- /dev/null +++ b/neural_compressor/data/datasets/style_transfer_dataset.py @@ -0,0 +1,84 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Dataset used for style transfer task on multiple framework backends.""" + +import os +import numpy as np +import glob +from .dataset import dataset_registry, Dataset + + +@dataset_registry(dataset_type="style_transfer", framework="tensorflow, \ + tensorflow_itex", dataset_format='') +class StyleTransferDataset(Dataset): # pragma: no cover + """Dataset used for style transfer task on tensorflow/inteltensorflow/tensorflow_itex backend. + + This Dataset is to construct a dataset from two specific image holders representing + content image folder and style image folder. + """ + + def __init__(self, content_folder, style_folder, crop_ratio=0.1, + resize_shape=(256, 256), image_format='jpg', transform=None, filter=None): + """Initialize `StyleTransferDataset` class. + + Args: + content_folder (str): Root directory of content images. + style_folder (str): Root directory of style images. + crop_ratio (float, default=0.1): Cropped ratio to each side. + resize_shape (tuple, default=(256, 256)): Target size of image. + image_format (str, default='jpg'): Target image format. + transform (transform object, default=None): Transform to process input data. + filter (Filter objects, default=None): Filter out examples according to specific conditions. + """ + self.transform = transform + self.content_folder = content_folder + self.style_folder = style_folder + self.resize_shape = resize_shape + self.crop_ratio = crop_ratio + self.content_images = glob.glob(os.path.join(content_folder, '*' + image_format)) + self.style_images = glob.glob(os.path.join(style_folder, '*' + image_format)) + self.image_list = [] + for content in self.content_images: + for style in self.style_images: + self.image_list.append((content, style)) + + def __len__(self): + """Return the length of dataset.""" + return len(self.image_list) + + def __getitem__(self, index): + """Return the item of dataset according to the given index.""" + from PIL import Image + content_image, style_image = self.image_list[index] + content_image = Image.open(content_image) + style_image = Image.open(style_image) + width, height = style_image.size + crop_ratio = self.crop_ratio + crop_box = ( + crop_ratio * height, + crop_ratio * width, + (1 - crop_ratio) * height, + (1 - crop_ratio) * width) + content_image = np.asarray(content_image.resize(self.resize_shape)) + style_image = np.asarray(style_image.resize(self.resize_shape)) + if content_image.max() > 1.0: + content_image = content_image / 255. + if style_image.max() > 1.0: + style_image = style_image / 255. + + return (content_image, style_image), 0 diff --git a/neural_compressor/data/filters/__init__.py b/neural_compressor/data/filters/__init__.py new file mode 100644 index 00000000000..6ec13cf416f --- /dev/null +++ b/neural_compressor/data/filters/__init__.py @@ -0,0 +1,31 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Built-in filter.""" + +from .filter import FILTERS, Filter, filter_registry +from os.path import dirname, basename, isfile, join +import glob + +modules = glob.glob(join(dirname(__file__), "*.py")) + +for f in modules: + if isfile(f) and not f.startswith('__') and not f.endswith('__init__.py'): + __import__(basename(f)[:-3], globals(), locals(), level=1) + + +__all__ = ["FILTERS", "Filter", "filter_registry"] diff --git a/neural_compressor/data/filters/coco_filter.py b/neural_compressor/data/filters/coco_filter.py new file mode 100644 index 00000000000..3f9431185ab --- /dev/null +++ b/neural_compressor/data/filters/coco_filter.py @@ -0,0 +1,60 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Built-in COCO filter.""" + +from neural_compressor.utils.utility import LazyImport +from .filter import Filter, filter_registry +tf = LazyImport('tensorflow') + + +@filter_registry(filter_type="LabelBalanceCOCORecord", framework="tensorflow, tensorflow_itex") +class LabelBalanceCOCORecordFilter(Filter): # pragma: no cover + """The label balance filter for COCO Record.""" + + def __init__(self, size=1): + """Initialize the attribute of class.""" + self.size = size + + def __call__(self, image, label): + """Execute the filter. + + Args: + image: Not used. + label: label of a sample. + """ + return tf.math.equal(len(label[0]), self.size) + + +@filter_registry(filter_type="LabelBalanceCOCORaw", framework="tensorflow, \ + tensorflow_itex, pytorch, mxnet, onnxrt_qlinearops, onnxrt_integerops") +class LabelBalanceCOCORawFilter(Filter): # pragma: no cover + """The label balance filter for COCO raw data.""" + + def __init__(self, size=1): + """Initialize the attribute of class.""" + self.size = size + + def __call__(self, image, label): + """Execute the filter. + + Args: + image: Not used. + label: label of a sample. + """ + return len(label) == self.size + diff --git a/neural_compressor/data/filters/filter.py b/neural_compressor/data/filters/filter.py new file mode 100644 index 00000000000..7abda00e054 --- /dev/null +++ b/neural_compressor/data/filters/filter.py @@ -0,0 +1,177 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""The base filter class for all frameworks.""" + +from abc import abstractmethod +from neural_compressor.utils.utility import singleton + + +@singleton +class TensorflowFilters(object): # pragma: no cover + """The base filter class for Tensorflow framework.""" + + def __init__(self): + """Initialize the atrribute of the class.""" + self.filters = {} + self.filters.update(TENSORFLOW_FILTERS) + + +@singleton +class ONNXRTQLFilters(object): # pragma: no cover + """The base filter class for ONNXRT framework QLinear mode.""" + + def __init__(self): + """Initialize the atrribute of the class.""" + self.filters = {} + self.filters.update(ONNXRT_QL_FILTERS) + + +@singleton +class ONNXRTITFilters(object): # pragma: no cover + """The base filter class for ONNXRT framework IT mode.""" + + def __init__(self): + """Initialize the atrribute of the class.""" + self.filters = {} + self.filters.update(ONNXRT_IT_FILTERS) + + +@singleton +class PyTorchFilters(object): # pragma: no cover + """The base filter class for PyTorch framework.""" + + def __init__(self): + """Initialize the atrribute of the class.""" + self.filters = {} + self.filters.update(PYTORCH_FILTERS) + + +@singleton +class MXNetFilters(object): # pragma: no cover + """The base filter class for MXNet framework.""" + + def __init__(self): + """Initialize the atrribute of the class.""" + self.filters = {} + self.filters.update(MXNET_FILTERS) + + +TENSORFLOW_FILTERS = {} +TENSORFLOW_ITEX_FILTERS = {} +ONNXRT_IT_FILTERS = {} +ONNXRT_QL_FILTERS = {} +PYTORCH_FILTERS = {} +MXNET_FILTERS = {} + +framework_filters = {"tensorflow": TensorflowFilters, + "tensorflow_itex": TensorflowFilters, + "pytorch": PyTorchFilters, + "pytorch_ipex": PyTorchFilters, + "pytorch_fx": PyTorchFilters, + "mxnet": MXNetFilters, + "onnxrt_qlinearops": ONNXRTQLFilters, + "onnxrt_qdq": ONNXRTQLFilters, + "onnxrt_qoperator": ONNXRTQLFilters, + "onnxrt_integerops": ONNXRTITFilters, + } + +registry_filters = {"tensorflow": TENSORFLOW_FILTERS, + "tensorflow_itex": TENSORFLOW_ITEX_FILTERS, + "pytorch": PYTORCH_FILTERS, + "pytorch_ipex": PYTORCH_FILTERS, + "pytorch_fx": PYTORCH_FILTERS, + "mxnet": MXNET_FILTERS, + "onnxrt_integerops": ONNXRT_IT_FILTERS, + "onnxrt_qdq": ONNXRT_QL_FILTERS, + "onnxrt_qoperator": ONNXRT_QL_FILTERS, + "onnxrt_qlinearops": ONNXRT_QL_FILTERS} + + +class FILTERS(object): # pragma: no cover + """The filter register for all frameworks. + + Args: + framework (str): frameworks in ["tensorflow", "tensorflow_itex", "mxnet", + "onnxrt_qdq", "pytorch", "pytorch_ipex", + "pytorch_fx", "onnxrt_integerops", + "onnxrt_qlinearops", "onnxrt_qoperator"]. + """ + + def __init__(self, framework): + """Initialize the attribute of class.""" + assert framework in ["tensorflow", "tensorflow_itex", + "mxnet", "onnxrt_qdq", "pytorch", "pytorch_ipex", "pytorch_fx", + "onnxrt_integerops", "onnxrt_qlinearops", "onnxrt_qoperator"], \ + "framework support tensorflow pytorch mxnet onnxrt" + self.filters = framework_filters[framework]().filters + self.framework = framework + + def __getitem__(self, filter_type): + """Magic method. + + x[i] is roughly equivalent to type(x).__getitem__(x, index) + """ + assert filter_type in self.filters.keys(), "filter support {}".\ + format(self.filters.keys()) + return self.filters[filter_type] + + +def filter_registry(filter_type, framework): # pragma: no cover + """Register all transform subclasses. + + Args: + filter_type (str): fILTER registration name. + framework (str): support 4 framework including 'tensorflow', 'pytorch', 'mxnet', 'onnxrt'. + cls (class): The class of register. + + Returns: + cls: The class of register. + """ + def decorator_transform(cls): + """Decorate a class.""" + for single_framework in [fwk.strip() for fwk in framework.split(',')]: + assert single_framework in [ + "tensorflow", + "tensorflow_itex", + "pytorch", + "pytorch_ipex", + "pytorch_fx", + "mxnet", + "onnxrt_integerops", + "onnxrt_qdq", + "onnxrt_qlinearops", + "onnxrt_qoperator" + ], "The framework support tensorflow mxnet pytorch onnxrt" + if filter_type in registry_filters[single_framework].keys(): + raise ValueError('Cannot have two transforms with the same name') + registry_filters[single_framework][filter_type] = cls + return cls + return decorator_transform + + +class Filter(object): # pragma: no cover + """The base class for transform. + + __call__ method is needed when write user specific transform. + + """ + + @abstractmethod + def __call__(self, *args, **kwargs): + """Execute the filter.""" + raise NotImplementedError diff --git a/neural_compressor/data/transforms/__init__.py b/neural_compressor/data/transforms/__init__.py index 77edda38b30..eb849e9f002 100644 --- a/neural_compressor/data/transforms/__init__.py +++ b/neural_compressor/data/transforms/__init__.py @@ -14,7 +14,11 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. +# ============================================================================== +"""Neural Compressor Built-in transforms for multiple framework backends.""" +from .transform import TRANSFORMS, BaseTransform, transform_registry +from .postprocess import Postprocess from os.path import dirname, basename, isfile, join import glob @@ -24,3 +28,5 @@ if isfile(f) and not f.startswith('__') and not f.endswith('__init__.py'): __import__(basename(f)[:-3], globals(), locals(), level=1) + +__all__ = ["TRANSFORMS", "BaseTransform", "transform_registry", "Postprocess"] diff --git a/neural_compressor/data/transforms/coco_transform.py b/neural_compressor/data/transforms/coco_transform.py index 7bef847eb78..6fdac97e94d 100644 --- a/neural_compressor/data/transforms/coco_transform.py +++ b/neural_compressor/data/transforms/coco_transform.py @@ -36,9 +36,11 @@ # BELOW IS TO BE DEPRECATED! @transform_registry(transform_type="ParseDecodeCoco", \ process="preprocess", framework="tensorflow") -class ParseDecodeCocoTransform(BaseTransform): - +class ParseDecodeCocoTransform(BaseTransform): # pragma: no cover + """Coco decoding will be performed automatically from Neural Compressor v1.4. + """ def __call__(self, sample): + """Convert `ParseDecodeCocoTransform` feature.""" logger.warning("This transform is going to be deprecated, " \ "coco decoding will be performed automatically from Neural Compressor v1.4.") return sample diff --git a/neural_compressor/data/transforms/imagenet_transform.py b/neural_compressor/data/transforms/imagenet_transform.py index 842e068a188..251e1aa667b 100644 --- a/neural_compressor/data/transforms/imagenet_transform.py +++ b/neural_compressor/data/transforms/imagenet_transform.py @@ -29,28 +29,130 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== +"""Neural Compressor built-in imagenet transforms.""" import numpy as np -from neural_compressor.utils import logger from neural_compressor.utils.utility import LazyImport -from neural_compressor.experimental.data.transforms import transform_registry, BaseTransform - +from neural_compressor.utils import logger +from .transform import transform_registry, BaseTransform tf = LazyImport('tensorflow') cv2 = LazyImport('cv2') -# BELOW IS TO BE DEPRECATED! +@transform_registry(transform_type="QuantizedInput", \ + process="preprocess", framework="tensorflow, tensorflow_itex") +class QuantizedInput(BaseTransform): # pragma: no cover + """Convert the dtype of input to quantize it. + + Args: + dtype(str): desired image dtype, support 'uint8', 'int8' + scale(float, default=None):scaling ratio of each point in image + + Returns: + tuple of processed image and label + """ + + def __init__(self, dtype, scale=None): + """Initialize `QuantizedInput` class.""" + self.dtype_map = {'uint8': tf.uint8, 'int8': tf.int8} + assert dtype in self.dtype_map.keys(), \ + 'only support cast dtype {}'.format(self.dtype_map.keys()) + self.dtype = dtype + self.scale = scale + + def __call__(self, sample): + """Convert the dtype of input.""" + # scale is not know when tuning, in this case this transform + # do nothing, it's only used when scale is set + if self.scale == None: + return sample + image, label = sample + image = image * self.scale + if self.dtype == 'uint8': + image = image + 128 + image = tf.dtypes.cast(image, dtype=self.dtype_map[self.dtype]) + return image, label + +@transform_registry(transform_type="LabelShift", \ + process="postprocess", framework="pytorch, tensorflow, tensorflow_itex,\ + onnxrt_qlinearops, onnxrt_integerops") +class LabelShift(BaseTransform): # pragma: no cover + """Convert label to label - label_shift. + + Args: + label_shift(int, default=0): number of label shift + + Returns: + tuple of processed image and label + """ + + def __init__(self, label_shift=0): + """Initialize `LabelShift` class.""" + self.label_shift = label_shift + + def __call__(self, sample): + """Convert label to label_shift.""" + images, labels = sample + if isinstance(labels, np.ndarray): + labels = labels - self.label_shift + elif isinstance(labels, list): + if isinstance(labels[0], tuple): + labels = [tuple(np.array(label) - self.label_shift) for label in labels] + elif isinstance(labels[0], np.ndarray): + labels = [label - self.label_shift for label in labels] + else: + labels = np.array(labels) - self.label_shift + labels = labels.tolist() + else: + labels = np.array(labels) - self.label_shift + return images, labels + +class ParseDecodeImagenet(): # pragma: no cover + """Parse features in Example proto. + + Returns: + tuple of parsed image and label + """ + + def __call__(self, sample): + """Parse features in example.""" + # Dense features in Example proto. + feature_map = { + 'image/encoded': tf.io.FixedLenFeature([], dtype=tf.string, default_value=''), + 'image/class/label': tf.io.FixedLenFeature([1], dtype=tf.int64, default_value=-1)} + + sparse_float32 = tf.io.VarLenFeature(dtype=tf.float32) + # Sparse features in Example proto. + feature_map.update( + {k: sparse_float32 for k in ['image/object/bbox/xmin', + 'image/object/bbox/ymin', + 'image/object/bbox/xmax', + 'image/object/bbox/ymax']}) + + features = tf.io.parse_single_example(serialized=sample, features=feature_map) + label = tf.cast(features['image/class/label'], dtype=tf.int32) + image = features['image/encoded'] + image = tf.image.decode_jpeg( + image, channels=3, fancy_upscaling=False, dct_method='INTEGER_FAST') + return (image, label) + @transform_registry(transform_type="ParseDecodeImagenet", \ process="preprocess", framework="tensorflow") -class ParseDecodeImagenetTransform(BaseTransform): +class ParseDecodeImagenetTransform(BaseTransform): # pragma: no cover + """Imagenet decoding will be performed automatically from Neural Compressor v1.4. + + Returns: + sample + """ def __call__(self, sample): + """Convert `ParseDecodeImagenetTransform` feature.""" logger.warning("This transform is going to be deprecated, " \ "imagenet decoding will be performed automatically from Neural Compressor v1.4.") return sample @transform_registry(transform_type="ResizeCropImagenet", \ process="preprocess", framework="tensorflow") -class TensorflowResizeCropImagenetTransform(BaseTransform): +class TensorflowResizeCropImagenetTransform(BaseTransform): # pragma: no cover """Combination of a series of transforms which is applicable to images in Imagenet. Args: @@ -70,7 +172,7 @@ def __init__(self, height, width, random_crop=False, resize_side=256, \ resize_method='bilinear', random_flip_left_right=False, \ mean_value=[0.0,0.0,0.0], scale=1.0, \ data_format='channels_last', subpixels='RGB'): - + """Initialize `TensorflowResizeCropImagenetTransform` class.""" self.height = height self.width = width self.mean_value = mean_value @@ -84,6 +186,7 @@ def __init__(self, height, width, random_crop=False, resize_side=256, \ # sample is (images, labels) def __call__(self, sample): + """Convert `TensorflowResizeCropImagenetTransform` feature.""" image, label = sample shape = tf.shape(input=image) @@ -99,15 +202,6 @@ def __call__(self, sample): new_height = tf.cast(tf.math.rint(height*scale), dtype=tf.int32) new_width = tf.cast(tf.math.rint(width*scale), dtype=tf.int32) - # image = tf.cond(pred=tf.greater(shape[0], shape[1]), \ - # false_fn=lambda: tf.image.resize(image, \ - # tf.convert_to_tensor(value=[self.resize_side*shape[0]/shape[1], \ - # self.resize_side], dtype=tf.int32)), - # true_fn=lambda: tf.image.resize(image, \ - # tf.convert_to_tensor(value=[self.resize_side, \ - # self.resize_side * shape[1] / shape[0]], dtype=tf.int32)), - # ) - if self.subpixels=='BGR' and self.data_format=='channels_first': # 'RGB'->'BGR' image = tf.cond(tf.equal(tf.rank(image), 3), @@ -119,8 +213,7 @@ def __call__(self, sample): image = tf.expand_dims(image, 0) image = tf.image.resize(image, [new_height, new_width], method=self.resize_method) - image = tf.squeeze(image) - + image = tf.squeeze(image) shape = tf.shape(input=image) if self.random_crop: y0 = tf.random.uniform(shape=[], minval=0, maxval=(shape[0] - self.height +1), @@ -141,7 +234,7 @@ def __call__(self, sample): @transform_registry(transform_type="BilinearImagenet", \ process="preprocess", framework="tensorflow") -class BilinearImagenetTransform(BaseTransform): +class BilinearImagenetTransform(BaseTransform): # pragma: no cover """Combination of a series of transforms which is applicable to images in Imagenet. Args: @@ -157,7 +250,7 @@ class BilinearImagenetTransform(BaseTransform): def __init__(self, height, width, central_fraction=0.875, mean_value=[0.0,0.0,0.0], scale=1.0): - + """Initialize `BilinearImagenetTransform` class.""" self.height = height self.width = width self.mean_value = mean_value @@ -166,6 +259,7 @@ def __init__(self, height, width, central_fraction=0.875, # sample is (images, labels) def __call__(self, sample): + """Convert `BilinearImagenetTransform` feature.""" image, label = sample if image.dtype is not tf.float32: image = tf.image.convert_image_dtype(image, dtype=tf.float32) @@ -184,12 +278,11 @@ def __call__(self, sample): image = tf.multiply(image, 2.0) means = tf.broadcast_to(self.mean_value, tf.shape(input=image)) image = (image - means) * self.scale - return (image, label) @transform_registry(transform_type="BilinearImagenet", process="preprocess", \ framework="onnxrt_qlinearops, onnxrt_integerops") -class OnnxBilinearImagenetTransform(BaseTransform): +class OnnxBilinearImagenetTransform(BaseTransform): # pragma: no cover """Combination of a series of transforms which is applicable to images in Imagenet. Args: @@ -205,6 +298,7 @@ class OnnxBilinearImagenetTransform(BaseTransform): def __init__(self, height, width, central_fraction=0.875, mean_value=[0.0,0.0,0.0], scale=1.0): + """Initialize `OnnxBilinearImagenetTransform` class.""" self.height = height self.width = width self.mean_value = mean_value @@ -212,6 +306,7 @@ def __init__(self, height, width, central_fraction=0.875, self.central_fraction = central_fraction def __call__(self, sample): + """Convert `OnnxBilinearImagenetTransform` feature.""" image, label = sample if isinstance(image, np.ndarray): image = image.astype('float32') / 255. @@ -235,12 +330,11 @@ def __call__(self, sample): means = np.broadcast_to(self.mean_value, image.shape) image = (image - means) * self.scale image = image.astype(np.float32) - return (image, label) @transform_registry(transform_type="ResizeCropImagenet", process="preprocess", \ framework="onnxrt_qlinearops, onnxrt_integerops") -class ONNXResizeCropImagenetTransform(BaseTransform): +class ONNXResizeCropImagenetTransform(BaseTransform): # pragma: no cover """Combination of a series of transforms which is applicable to images in Imagenet. Args: @@ -257,7 +351,7 @@ class ONNXResizeCropImagenetTransform(BaseTransform): def __init__(self, height, width, random_crop=False, resize_side=256, \ mean_value=[0.0,0.0,0.0], std_value=[0.229, 0.224, 0.225], \ resize_method='bilinear', data_format='channels_last', subpixels='RGB'): - + """Initialize `ONNXResizeCropImagenetTransform` class.""" self.height = height self.width = width self.mean_value = mean_value @@ -270,6 +364,7 @@ def __init__(self, height, width, random_crop=False, resize_side=256, \ # sample is (images, labels) def __call__(self, sample): + """Convert `ONNXResizeCropImagenetTransform` feature.""" # TODO Support optional resize_method, data_format, subpixels for ONNX image, label = sample height, width = image.shape[0], image.shape[1] @@ -295,14 +390,22 @@ def __call__(self, sample): @transform_registry(transform_type="ResizeWithAspectRatio", process="preprocess", \ framework="onnxrt_qlinearops, onnxrt_integerops") -class ResizeWithAspectRatio(BaseTransform): +class ResizeWithAspectRatio(BaseTransform): # pragma: no cover + """Resize the image with aspect ratio. + + Returns: + image and label + """ + def __init__(self, height, width, scale=87.5, inter_pol=cv2.INTER_AREA): + """Initialize `ResizeWithAspectRatio` class.""" self.height = height self.width = width self.scale = scale self.inter_pol = inter_pol def __call__(self, sample): + """Convert `ResizeWithAspectRatio` feature.""" (img, label) = sample assert len(img.shape) == 3 height, width, _ = img.shape diff --git a/neural_compressor/data/transforms/postprocess.py b/neural_compressor/data/transforms/postprocess.py new file mode 100644 index 00000000000..605417a73ab --- /dev/null +++ b/neural_compressor/data/transforms/postprocess.py @@ -0,0 +1,27 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Common Postprocess.""" + +class Postprocess(object): +# class Transform(object): + """Just collect the infos to construct a Postprocess.""" + + def __init__(self, postprocess_cls, name='user_postprocess', **kwargs): + """Initialize `Postprocess` class.""" + self.postprocess_cls = postprocess_cls + self.name = name + self.kwargs = kwargs diff --git a/neural_compressor/data/transforms/tokenization.py b/neural_compressor/data/transforms/tokenization.py new file mode 100644 index 00000000000..06dc3a86fe4 --- /dev/null +++ b/neural_compressor/data/transforms/tokenization.py @@ -0,0 +1,342 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# coding=utf-8 +# Copyright 2018 The Google AI Language Team Authors. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Tokenization helper classes.""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function +from neural_compressor.utils.utility import LazyImport +import collections +import re +import unicodedata +import six +tf = LazyImport('tensorflow') + +def convert_to_unicode(text): # pragma: no cover + """Convert `text` to Unicode (if it's not already), assuming utf-8 input.""" + if six.PY3: + if isinstance(text, str): + return text + elif isinstance(text, bytes): + return text.decode("utf-8", "ignore") + else: + raise ValueError("Unsupported string type: %s" % (type(text))) + elif six.PY2: + if isinstance(text, str): + return text.decode("utf-8", "ignore") + elif isinstance(text, unicode): # pylint: disable=undefined-variable + return text + else: + raise ValueError("Unsupported string type: %s" % (type(text))) + else: + raise ValueError("Not running on Python2 or Python 3?") + +def load_vocab(vocab_file): + """Load a vocabulary file into a dictionary.""" + vocab = collections.OrderedDict() + index = 0 + with tf.io.gfile.GFile(vocab_file, "r") as reader: + while True: + token = convert_to_unicode(reader.readline()) + if not token: + break + token = token.strip() + vocab[token] = index + index += 1 + return vocab + +def convert_by_vocab(vocab, items): + """Convert a sequence of [tokens|ids] using the vocab.""" + output = [] + for item in items: + output.append(vocab[item]) + return output + +def whitespace_tokenize(text): + """Run basic whitespace cleaning and splitting on a piece of text.""" + text = text.strip() + if not text: + return [] + tokens = text.split() + return tokens + + +class FullTokenizer(object): + """Run end-to-end tokenziation.""" + + def __init__(self, vocab_file, do_lower_case=True): + """Construct a FullTokenizer. + + Args: + vocab_file: vocab file. + do_lower_case: Whether to lower case the input. + """ + self.vocab = load_vocab(vocab_file) + self.inv_vocab = {v: k for k, v in self.vocab.items()} + self.basic_tokenizer = BasicTokenizer(do_lower_case=do_lower_case) + self.wordpiece_tokenizer = WordpieceTokenizer(vocab=self.vocab) + + def tokenize(self, text): + """Tokenize text.""" + split_tokens = [] + for token in self.basic_tokenizer.tokenize(text): + for sub_token in self.wordpiece_tokenizer.tokenize(token): + split_tokens.append(sub_token) + + return split_tokens + + def convert_tokens_to_ids(self, tokens): + """Convert tokens to ids.""" + return convert_by_vocab(self.vocab, tokens) + + def convert_ids_to_tokens(self, ids): + """Convert ids to tokens.""" + return convert_by_vocab(self.inv_vocab, ids) + + +class BasicTokenizer(object): + """Run basic tokenization (punctuation splitting, lower casing, etc.).""" + + def __init__(self, do_lower_case=True): + """Construct a BasicTokenizer. + + Args: + do_lower_case: Whether to lower case the input. + """ + self.do_lower_case = do_lower_case + + def tokenize(self, text): + """Tokenizes a piece of text.""" + text = convert_to_unicode(text) + text = self._clean_text(text) + + # This was added on November 1st, 2018 for the multilingual and Chinese + # models. This is also applied to the English models now, but it doesn't + # matter since the English models were not trained on any Chinese data + # and generally don't have any Chinese data in them (there are Chinese + # characters in the vocabulary because Wikipedia does have some Chinese + # words in the English Wikipedia.). + text = self._tokenize_chinese_chars(text) + + orig_tokens = whitespace_tokenize(text) + split_tokens = [] + for token in orig_tokens: + if self.do_lower_case: + token = token.lower() + token = self._run_strip_accents(token) + split_tokens.extend(self._run_split_on_punc(token)) + + output_tokens = whitespace_tokenize(" ".join(split_tokens)) + return output_tokens + + def _run_strip_accents(self, text): + """Strip accents from a piece of text.""" + text = unicodedata.normalize("NFD", text) + output = [] + for char in text: + cat = unicodedata.category(char) + if cat == "Mn": + continue + output.append(char) + return "".join(output) + + def _run_split_on_punc(self, text): + """Split punctuation on a piece of text.""" + chars = list(text) + i = 0 + start_new_word = True + output = [] + while i < len(chars): + char = chars[i] + if _is_punctuation(char): + output.append([char]) + start_new_word = True + else: + if start_new_word: + output.append([]) + start_new_word = False + output[-1].append(char) + i += 1 + + return ["".join(x) for x in output] + + def _tokenize_chinese_chars(self, text): + """Add whitespace around any CJK character.""" + output = [] + for char in text: + cp = ord(char) + if self._is_chinese_char(cp): # pragma: no cover + output.append(" ") + output.append(char) + output.append(" ") + else: + output.append(char) + return "".join(output) + + def _is_chinese_char(self, cp): + """Check whether CP is the codepoint of a CJK character.""" + # This defines a "chinese character" as anything in the CJK Unicode block: + # https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block) + # + # Note that the CJK Unicode block is NOT all Japanese and Korean characters, + # despite its name. The modern Korean Hangul alphabet is a different block, + # as is Japanese Hiragana and Katakana. Those alphabets are used to write + # space-separated words, so they are not treated specially and handled + # like the all of the other languages. + if ((cp >= 0x4E00 and cp <= 0x9FFF) or # + (cp >= 0x3400 and cp <= 0x4DBF) or # + (cp >= 0x20000 and cp <= 0x2A6DF) or # + (cp >= 0x2A700 and cp <= 0x2B73F) or # + (cp >= 0x2B740 and cp <= 0x2B81F) or # + (cp >= 0x2B820 and cp <= 0x2CEAF) or + (cp >= 0xF900 and cp <= 0xFAFF) or # + (cp >= 0x2F800 and cp <= 0x2FA1F)): # + return True + + return False + + def _clean_text(self, text): + """Perform invalid character removal and whitespace cleanup on text.""" + output = [] + for char in text: + cp = ord(char) + if cp == 0 or cp == 0xfffd or _is_control(char): + continue + if _is_whitespace(char): + output.append(" ") + else: + output.append(char) + return "".join(output) + + +class WordpieceTokenizer(object): + """Run WordPiece tokenziation.""" + + def __init__(self, vocab, unk_token="[UNK]", max_input_chars_per_word=200): + """Construct a WordpieceTokenizer. + + Args: + vocab: the given vocabulary. + unk_token: unknown token. + max_input_chars_per_word: max input chars number in any word. + """ + self.vocab = vocab + self.unk_token = unk_token + self.max_input_chars_per_word = max_input_chars_per_word + + def tokenize(self, text): + """Tokenize a piece of text into its word pieces. + + This uses a greedy longest-match-first algorithm to perform tokenization + using the given vocabulary. + For example: + input = "unaffable" + output = ["un", "##aff", "##able"] + Args: + text: A single token or whitespace separated tokens. This should have + already been passed through `BasicTokenizer. + + Returns: + A list of wordpiece tokens. + """ + text = convert_to_unicode(text) + + output_tokens = [] + for token in whitespace_tokenize(text): + chars = list(token) + if len(chars) > self.max_input_chars_per_word: # pragma: no cover + output_tokens.append(self.unk_token) + continue + + is_bad = False + start = 0 + sub_tokens = [] + while start < len(chars): + end = len(chars) + cur_substr = None + while start < end: + substr = "".join(chars[start:end]) + if start > 0: + substr = "##" + substr + if substr in self.vocab: + cur_substr = substr + break + end -= 1 + if cur_substr is None: + is_bad = True + break + sub_tokens.append(cur_substr) + start = end + + if is_bad: + output_tokens.append(self.unk_token) + else: + output_tokens.extend(sub_tokens) + return output_tokens + +def _is_whitespace(char): + """Check whether `chars` is a whitespace character.""" + # \t, \n, and \r are technically contorl characters but we treat them + # as whitespace since they are generally considered as such. + if char == " " or char == "\t" or char == "\n" or char == "\r": + return True + cat = unicodedata.category(char) + if cat == "Zs": # pragma: no cover + return True + return False + +def _is_control(char): # pragma: no cover + """Check whether `chars` is a control character.""" + # These are technically control characters but we count them as whitespace + # characters. + if char == "\t" or char == "\n" or char == "\r": + return False + cat = unicodedata.category(char) + if cat in ("Cc", "Cf"): + return True + return False + +def _is_punctuation(char): # pragma: no cover + """Check whether `chars` is a punctuation character.""" + cp = ord(char) + # We treat all non-letter/number ASCII as punctuation. + # Characters such as "^", "$", and "`" are not in the Unicode + # Punctuation class but we treat them as punctuation anyways, for + # consistency. + if ((cp >= 33 and cp <= 47) or (cp >= 58 and cp <= 64) or + (cp >= 91 and cp <= 96) or (cp >= 123 and cp <= 126)): + return True + cat = unicodedata.category(char) + if cat.startswith("P"): + return True + return False diff --git a/neural_compressor/data/transforms/transform.py b/neural_compressor/data/transforms/transform.py new file mode 100644 index 00000000000..3c645057f23 --- /dev/null +++ b/neural_compressor/data/transforms/transform.py @@ -0,0 +1,2727 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Neural Compressor built-in Transforms on multiple framework backends.""" + +import numpy as np +import collections +from abc import abstractmethod +from neural_compressor.utils.utility import LazyImport, singleton +from neural_compressor.utils import logger + +torchvision = LazyImport('torchvision') +torch = LazyImport('torch') +tf = LazyImport('tensorflow') +mx = LazyImport('mxnet') +cv2 = LazyImport('cv2') + +class Transforms(object): + """INC supports built-in preprocessing, postprocessing and general methods on different framework backends. + + Transforms base class provides the abstract methods. + Users can also register their own Transforms classes by inheriting this base class. + """ + + def __init__(self, process, concat_general=True): + """Initialize `Transforms` class. + + Args: + process (str): processing type, the value can be preprocess, postprocess or general + concat_general (Boolean): users can use general transform in both preprocess + or postprocess if set True + """ + transform_map = {"preprocess": self._get_preprocess, + "postprocess": self._get_postprocess, + "general": self._get_general, } + self.transforms = transform_map[process]() + if concat_general: + self.transforms.update(transform_map['general']()) + + @abstractmethod + def _get_preprocess(self): + """Abstract method to get preprocessing method.""" + raise NotImplementedError + + @abstractmethod + def _get_postprocess(self): + """Abstract method to get postprocess method.""" + raise NotImplementedError + + @abstractmethod + def _get_general(self): + """Abstract method to get general method.""" + raise NotImplementedError + + +class TensorflowTransforms(Transforms): + """Tensorflow Transforms subclass.""" + + def _get_preprocess(self): + """Tensorflow get preprocess method. + + Returns: + preprocess: a dict including all the registered preprocess methods + """ + preprocess = { + "DecodeImage": TensorflowWrapFunction(tf.io.decode_jpeg), + "EncodeJpeg": TensorflowWrapFunction(tf.io.encode_jpeg), + } + # update the registry transforms + preprocess.update(TENSORFLOW_TRANSFORMS["preprocess"]) + return preprocess + + def _get_postprocess(self): + """Tensorflow get postprocess method. + + Returns: + postprocess: a dict including all the registered postprocess methods + """ + postprocess = {} + postprocess.update(TENSORFLOW_TRANSFORMS["postprocess"]) + return postprocess + + def _get_general(self): + """Tensorflow get general method. + + Returns: + general: a dict including all the registered general methods + """ + general = {} + general.update(TENSORFLOW_TRANSFORMS["general"]) + return general + + +class MXNetTransforms(Transforms): + """Mxnet Transforms subclass.""" + + def _get_preprocess(self): + """Mxnet get preprocess method. + + Returns: + preprocess: a dict including all the registered preprocess methods + """ + preprocess = { + 'ToTensor': PytorchMxnetWrapFunction( + mx.gluon.data.vision.transforms.ToTensor), + 'CenterCrop': PytorchMxnetWrapFunction( + mx.gluon.data.vision.transforms.CenterCrop), + 'RandomHorizontalFlip': PytorchMxnetWrapFunction( + mx.gluon.data.vision.transforms.RandomFlipLeftRight), + 'RandomVerticalFlip': PytorchMxnetWrapFunction( + mx.gluon.data.vision.transforms.RandomFlipTopBottom), + } + preprocess.update(MXNET_TRANSFORMS["preprocess"]) + return preprocess + + def _get_postprocess(self): + """Mxnet get postprocess method. + + Returns: + postprocess: a dict including all the registered postprocess methods + """ + postprocess = {} + postprocess.update(MXNET_TRANSFORMS["postprocess"]) + return postprocess + + def _get_general(self): + """Mxnet get general method. + + Returns: + general: a dict including all the registered general methods + """ + general = { + 'Compose': mx.gluon.data.vision.transforms.Compose, + 'Cast': PytorchMxnetWrapFunction( + mx.gluon.data.vision.transforms.Cast), + } + general.update(MXNET_TRANSFORMS["general"]) + return general + + +class PyTorchTransforms(Transforms): + """Pytorch Transforms subclass.""" + + def _get_preprocess(self): + """Pytorch get preprocessing method. + + Returns: + preprocess: a dict including all the registered preprocess methods + """ + preprocess = { + "ToTensor": PytorchMxnetWrapFunction( + torchvision.transforms.ToTensor), + "ToPILImage": PytorchMxnetWrapFunction( + torchvision.transforms.ToPILImage), + "CenterCrop": PytorchMxnetWrapFunction( + torchvision.transforms.CenterCrop), + "RandomCrop": PytorchMxnetWrapFunction( + torchvision.transforms.RandomCrop), + "RandomHorizontalFlip": PytorchMxnetWrapFunction( + torchvision.transforms.RandomHorizontalFlip), + "RandomVerticalFlip": PytorchMxnetWrapFunction( + torchvision.transforms.RandomVerticalFlip), + "Pad": PytorchMxnetWrapFunction( + torchvision.transforms.Pad), + "ColorJitter": PytorchMxnetWrapFunction( + torchvision.transforms.ColorJitter), + } + preprocess.update(PYTORCH_TRANSFORMS["preprocess"]) + return preprocess + + def _get_postprocess(self): + """Pytorch get postprocess method. + + Returns: + postprocess: a dict including all the registered postprocess methods + """ + postprocess = {} + postprocess.update(PYTORCH_TRANSFORMS["postprocess"]) + return postprocess + + def _get_general(self): + """Pytorch get general method. + + Returns: + general: a dict including all the registered general methods + """ + general = { + "Compose": torchvision.transforms.Compose, + } + general.update(PYTORCH_TRANSFORMS["general"]) + return general + +class ONNXRTQLTransforms(Transforms): + """Onnxrt_qlinearops Transforms subclass.""" + + def _get_preprocess(self): + """Onnxrt_qlinearops get preprocessing method. + + Returns: + preprocess: a dict including all the registered preprocess methods + """ + preprocess = {} + preprocess.update(ONNXRT_QL_TRANSFORMS["preprocess"]) + return preprocess + + def _get_postprocess(self): + """Onnxrt_qlinearops get postprocess method. + + Returns: + postprocess: a dict including all the registered postprocess methods + """ + postprocess = {} + postprocess.update(ONNXRT_QL_TRANSFORMS["postprocess"]) + return postprocess + + def _get_general(self): + """Onnxrt_qlinearops get general method. + + Returns: + general: a dict including all the registered general methods + """ + general = {} + general.update(ONNXRT_QL_TRANSFORMS["general"]) + return general + +class ONNXRTITTransforms(Transforms): + """Onnxrt_integerops Transforms subclass.""" + + def _get_preprocess(self): + """Onnxrt_integerops get preprocessing method. + + Returns: + preprocess: a dict including all the registered preprocess methods + """ + preprocess = {} + preprocess.update(ONNXRT_IT_TRANSFORMS["preprocess"]) + return preprocess + + def _get_postprocess(self): + """Onnxrt_integerops get postprocess method. + + Returns: + postprocess: a dict including all the registered postprocess methods + """ + postprocess = {} + postprocess.update(ONNXRT_IT_TRANSFORMS["postprocess"]) + return postprocess + + def _get_general(self): + """Onnxrt_integerops get general method. + + Returns: + general: a dict including all the registered general methods + """ + general = {} + general.update(ONNXRT_IT_TRANSFORMS["general"]) + return general + + +framework_transforms = {"tensorflow": TensorflowTransforms, + "tensorflow_itex": TensorflowTransforms, + "mxnet": MXNetTransforms, + "pytorch": PyTorchTransforms, + "pytorch_ipex": PyTorchTransforms, + "pytorch_fx": PyTorchTransforms, + "onnxrt_qlinearops": ONNXRTQLTransforms, + "onnxrt_integerops": ONNXRTITTransforms, + "onnxrt_qoperator": ONNXRTQLTransforms, + "onnxrt_qdq": ONNXRTQLTransforms} + +# transform registry will register transforms into these dicts +TENSORFLOW_TRANSFORMS = {"preprocess": {}, "postprocess": {}, "general": {}} +TENSORFLOW_ITEX_TRANSFORMS = {"preprocess": {}, "postprocess": {}, "general": {}} +MXNET_TRANSFORMS = {"preprocess": {}, "postprocess": {}, "general": {}} +PYTORCH_TRANSFORMS = {"preprocess": {}, "postprocess": {}, "general": {}} +ONNXRT_QL_TRANSFORMS = {"preprocess": {}, "postprocess": {}, "general": {}} +ONNXRT_IT_TRANSFORMS = {"preprocess": {}, "postprocess": {}, "general": {}} + +registry_transforms = {"tensorflow": TENSORFLOW_TRANSFORMS, + "tensorflow_itex": TENSORFLOW_ITEX_TRANSFORMS, + "mxnet": MXNET_TRANSFORMS, + "pytorch": PYTORCH_TRANSFORMS, + "pytorch_ipex": PYTORCH_TRANSFORMS, + "pytorch_fx": PYTORCH_TRANSFORMS, + "onnxrt_qlinearops": ONNXRT_QL_TRANSFORMS, + "onnxrt_qdq": ONNXRT_QL_TRANSFORMS, + "onnxrt_qoperator": ONNXRT_QL_TRANSFORMS, + "onnxrt_integerops": ONNXRT_IT_TRANSFORMS, + } + +class TRANSFORMS(object): + """Transforms collection class. + + Provide register method to register new Transforms + and provide __getitem__ method to get Transforms according to Transforms type. + """ + + def __init__(self, framework, process): + """Initialize `TRANSFORMS` class. + + Args: + framework (str): different framework type like tensorflow, pytorch and so on + process (str): process type, the value can be preprocess, postprocess or general + """ + assert framework in ("tensorflow", "tensorflow_itex", "onnxrt_qoperator", \ + "pytorch", "pytorch_ipex", "pytorch_fx", "onnxrt_qdq", \ + "onnxrt_qlinearops", "onnxrt_integerops", "mxnet"), \ + "framework support tensorflow pytorch mxnet onnxrt" + assert process in ("preprocess", "postprocess", + "general"), "process support preprocess postprocess, general" + self.transforms = framework_transforms[framework](process).transforms + self.framework = framework + self.process = process + + def __getitem__(self, transform_type): + """Get Transform according to Transforms type. + + Args: + transform_type (str): the value can be preprocess, postprocess or general + + Returns: + Transforms: the registered Transforms + """ + assert transform_type in self.transforms.keys(), "transform support {}".\ + format(self.transforms.keys()) + return self.transforms[transform_type] + + def register(self, name, transform_cls): + """Register new Transform according to Transforms type. + + Args: + name (str): process name + transform_cls (class): process function wrapper class + """ + assert name not in registry_transforms[self.framework][self.process].keys(), \ + 'register transform name already exists.' + registry_transforms[self.framework][self.process].update({name: transform_cls}) + + +def transform_registry(transform_type, process, framework): + """Class decorator used to register all transform subclasses. + + Args: + transform_type (str): Transform registration name + process (str): support 3 process including 'preprocess', 'postprocess', 'general' + framework (str): support 4 framework including 'tensorflow', 'pytorch', 'mxnet', 'onnxrt' + cls (class): The class of register. + + Returns: + cls: The class of register. + """ + def decorator_transform(cls): + for single_framework in [fwk.strip() for fwk in framework.split(',')]: + assert single_framework in [ + "tensorflow", + "tensorflow_itex", + "mxnet", + "pytorch", + "pytorch_ipex", + "pytorch_fx", + "onnxrt_qlinearops", + "onnxrt_qdq", + "onnxrt_integerops", + "onnxrt_qoperator", + ], "The framework support tensorflow mxnet pytorch onnxrt" + if transform_type in registry_transforms[single_framework][process].keys(): + raise ValueError('Cannot have two transforms with the same name') + registry_transforms[single_framework][process][transform_type] = cls + return cls + return decorator_transform + + +class BaseTransform(object): + """The base class for transform.""" + + @abstractmethod + def __call__(self, *args, **kwargs): + """__call__ method is needed when write user specific transform.""" + raise NotImplementedError + + +class TensorflowWrapFunction(object): + """Tensorflow wrapper function class.""" + + def __init__(self, transform_func): + """Initialize `TensorflowWrapFunction` class. + + Args: + transform_func (function): tensorflow tranform function + """ + self.transform_func = transform_func + + def __call__(self, **kwargs): + """__call__ method. + + Returns: + TensorflowTransform class + """ + return TensorflowTransform(self.transform_func, **kwargs) + +class TensorflowTransform(BaseTransform): + """Tensorflow transform class, the subclass of BaseTransform.""" + + def __init__(self, transform_func, **kwargs): + """Initialize `TensorflowTransform` class. + + Args: + transform_func (function): tensorflow tranform function + """ + self.kwargs = kwargs + self.transform_func = transform_func + + def __call__(self, sample): + """__call__ method. + + Returns: + a tuple of image and lable which get from tensorflow tranform processing + """ + image, label = sample + image = self.transform_func(image, **self.kwargs) + return (image, label) + +class PytorchMxnetWrapFunction(object): + """Pytorch and MXNet wrapper function class.""" + + def __init__(self, transform_func): + """Initialize `PytorchMxnetWrapFunction` class. + + Args: + transform_func (function): pytorch or mxnet tranform function + """ + self.transform_func = transform_func + + def __call__(self, **args): + """__call__ method. + + Returns: + PytorchMxnetTransform class + """ + return PytorchMxnetTransform(self.transform_func(**args)) + +class PytorchMxnetTransform(BaseTransform): + """Pytorch and Mxnet transform class, the subclass of BaseTransform.""" + + def __init__(self, transform_func): + """Initialize `PytorchMxnetTransform` class. + + Args: + transform_func (function): pytorch or mxnet tranform function + """ + self.transform_func = transform_func + + def __call__(self, sample): + """__call__ method. + + Returns: + a tuple of image and lable which get from pytorch or mxnet tranform processing + """ + image, label = sample + image = self.transform_func(image) + return (image, label) + +interpolation_map = { + 'nearest': cv2.INTER_NEAREST, + 'bilinear': cv2.INTER_LINEAR, + 'bicubic': cv2.INTER_CUBIC, +} + +interpolation_pytorch_map = { + 'nearest': 0, + 'bilinear': 2, + 'bicubic': 3, +} + +interpolation_mxnet_map = { + 'nearest': 0, + 'bilinear': 1, + 'bicubic': 2, +} + +def get_torchvision_map(interpolation): + """Get torchvision interpolation map.""" + try: + from torchvision.transforms.functional import InterpolationMode + interpolation_torchvision_map = { + 0: InterpolationMode.NEAREST, + 2: InterpolationMode.BILINEAR, + 3: InterpolationMode.BICUBIC, + } + return interpolation_torchvision_map[interpolation] + except: # pragma: no cover + return interpolation + +@transform_registry(transform_type="Compose", process="general", \ + framework="onnxrt_qlinearops, onnxrt_integerops, tensorflow, \ + tensorflow_itex") +class ComposeTransform(BaseTransform): + """Composes several transforms together. + + Args: + transform_list (list of Transform objects): list of transforms to compose + + Returns: + sample (tuple): tuple of processed image and label + """ + + def __init__(self, transform_list): + """Initialize `ComposeTransform` class.""" + self.transform_list = transform_list + + def __call__(self, sample): + """Call transforms in transform_list.""" + for transform in self.transform_list: + sample = transform(sample) + return sample + +@transform_registry(transform_type="CropToBoundingBox", process="preprocess", \ + framework="pytorch") +class CropToBoundingBox(BaseTransform): + """Crops an image to a specified bounding box. + + Args: + offset_height (int): Vertical coordinate of the top-left corner of the result in the input + offset_width (int): Horizontal coordinate of the top-left corner of the result in the input + target_height (int): Height of the result + target_width (int): Width of the result + + Returns: + tuple of processed image and label + """ + + def __init__(self, offset_height, offset_width, target_height, target_width): + """Initialize `CropToBoundingBox` class.""" + self.offset_height = offset_height + self.offset_width = offset_width + self.target_height = target_height + self.target_width = target_width + + def __call__(self, sample): + """Call torchvision.transforms.functional.crop.""" + image, label = sample + image = torchvision.transforms.functional.crop( + image, + self.offset_height, + self.offset_width, + self.target_height, + self.target_width) + return (image, label) + +@transform_registry(transform_type="CropToBoundingBox", process="preprocess", \ + framework="mxnet") +class MXNetCropToBoundingBox(CropToBoundingBox): + """Crops an image to a specified bounding box. + + Args: + offset_height (int): Vertical coordinate of the top-left corner of the result in the input + offset_width (int): Horizontal coordinate of the top-left corner of the result in the input + target_height (int): Height of the result + target_width (int): Width of the result + + Returns: + tuple of processed image and label + """ + + def __call__(self, sample): + """Call mx.image.fixed_crop.""" + image, label = sample + image = mx.image.fixed_crop( + image, + self.offset_height, + self.offset_width, + self.target_height, + self.target_width) + return (image, label) + +@transform_registry(transform_type="CropToBoundingBox", process="preprocess", \ + framework="onnxrt_qlinearops, onnxrt_integerops") +class ONNXRTCropToBoundingBox(CropToBoundingBox): + """Crops an image to a specified bounding box. + + Args: + offset_height (int): Vertical coordinate of the top-left corner of the result in the input + offset_width (int): Horizontal coordinate of the top-left corner of the result in the input + target_height (int): Height of the result + target_width (int): Width of the result + + Returns: + tuple of processed image and label + """ + + def __call__(self, sample): + """Crop the image in sample.""" + image, label = sample + image = image[self.offset_height : self.offset_height+self.target_height, + self.offset_width : self.offset_width+self.target_width, :] + return (image, label) + +@transform_registry(transform_type="CropToBoundingBox", process="preprocess", \ + framework="tensorflow, tensorflow_itex") +class TensorflowCropToBoundingBox(CropToBoundingBox): + """Crops an image to a specified bounding box. + + Args: + offset_height (int): Vertical coordinate of the top-left corner of the result in the input + offset_width (int): Horizontal coordinate of the top-left corner of the result in the input + target_height (int): Height of the result + target_width (int): Width of the result + + Returns: + tuple of processed image and label + """ + + def __call__(self, sample): + """Crop the image in sample.""" + image, label = sample + if isinstance(image, tf.Tensor): + image = tf.image.crop_to_bounding_box(image, self.offset_height, + self.offset_width, self.target_height, self.target_width) + else: + image = image[self.offset_height : self.offset_height+self.target_height, + self.offset_width : self.offset_width+self.target_width, :] + return (image, label) + +@transform_registry(transform_type="ResizeWithRatio", process="preprocess", \ + framework="onnxrt_qlinearops, onnxrt_integerops, pytorch, mxnet") +class ResizeWithRatio(BaseTransform): + """Resize image with aspect ratio and pad it to max shape(optional). + + If the image is padded, the label will be processed at the same time. + The input image should be np.array. + + Args: + min_dim (int, default=800): + Resizes the image such that its smaller dimension == min_dim + max_dim (int, default=1365): + Ensures that the image longest side doesn't exceed this value + padding (bool, default=False): + If true, pads image with zeros so its size is max_dim x max_dim + + Returns: + tuple of processed image and label + """ + + def __init__(self, min_dim=800, max_dim=1365, padding=False, constant_value=0): + """Initialize `ResizeWithRatio` class.""" + self.min_dim = min_dim + self.max_dim = max_dim + self.padding = padding + self.constant_value = constant_value + + def __call__(self, sample): + """Resize the image with ratio in sample.""" + image, label = sample + height, width = image.shape[:2] + scale = 1 + if self.min_dim: + scale = max(1, self.min_dim / min(height, width)) + if self.max_dim: + image_max = max(height, width) + if round(image_max * scale) > self.max_dim: + scale = self.max_dim / image_max + if scale != 1: + image = cv2.resize(image, (round(height * scale), round(width * scale))) + + bbox, str_label, int_label, image_id = label + + if self.padding: + h, w = image.shape[:2] + pad_param = [[(self.max_dim-h)//2, self.max_dim-h-(self.max_dim-h)//2], + [(self.max_dim-w)//2, self.max_dim-w-(self.max_dim-w)//2], + [0, 0]] + if not isinstance(bbox, np.ndarray): + bbox = np.array(bbox) + resized_box = bbox * [height, width, height, width] * scale + moved_box = (resized_box + [(self.max_dim-h)//2, (self.max_dim-w)//2, \ + (self.max_dim-h)//2, (self.max_dim-w)//2]) + bbox = moved_box / [self.max_dim, self.max_dim, self.max_dim, self.max_dim] + image = np.pad(image, pad_param, mode='constant', constant_values=self.constant_value) + return image, (bbox, str_label, int_label, image_id) + +@transform_registry(transform_type="ResizeWithRatio", process="preprocess", \ + framework="tensorflow, tensorflow_itex") +class TensorflowResizeWithRatio(BaseTransform): + """Resize image with aspect ratio and pad it to max shape(optional). + + If the image is padded, the label will be processed at the same time. + The input image should be np.array or tf.Tensor. + + Args: + min_dim (int, default=800): + Resizes the image such that its smaller dimension == min_dim + max_dim (int, default=1365): + Ensures that the image longest side doesn't exceed this value + padding (bool, default=False): + If true, pads image with zeros so its size is max_dim x max_dim + + Returns: + tuple of processed image and label + """ + + def __init__(self, min_dim=800, max_dim=1365, padding=False, constant_value=0): + """Initialize `TensorflowResizeWithRatio` class.""" + self.min_dim = min_dim + self.max_dim = max_dim + self.padding = padding + self.constant_value = constant_value + + def __call__(self, sample): + """Resize the image with ratio in sample.""" + image, label = sample + if isinstance(image, tf.Tensor): + shape = tf.shape(input=image) + height = tf.cast(shape[0], dtype=tf.float32) + width = tf.cast(shape[1], dtype=tf.float32) + scale = 1 + if self.min_dim: + scale = tf.maximum(1., tf.cast(self.min_dim / tf.math.minimum(height, width),\ + dtype=tf.float32)) + if self.max_dim: + image_max = tf.cast(tf.maximum(height, width), dtype=tf.float32) + scale = tf.cond(pred=tf.greater(tf.math.round(image_max * scale), self.max_dim), \ + true_fn=lambda: self.max_dim / image_max, + false_fn=lambda: scale) + image = tf.image.resize(image, (tf.math.round(height * scale), \ + tf.math.round(width * scale))) + bbox, str_label, int_label, image_id = label + + if self.padding: + shape = tf.shape(input=image) + h = tf.cast(shape[0], dtype=tf.float32) + w = tf.cast(shape[1], dtype=tf.float32) + pad_param = [[(self.max_dim-h)//2, self.max_dim-h-(self.max_dim-h)//2], + [(self.max_dim-w)//2, self.max_dim-w-(self.max_dim-w)//2], + [0, 0]] + resized_box = bbox * [height, width, height, width] * scale + moved_box = (resized_box + [(self.max_dim-h)//2, (self.max_dim-w)//2, \ + (self.max_dim-h)//2, (self.max_dim-w)//2]) + bbox = moved_box / [self.max_dim, self.max_dim, self.max_dim, self.max_dim] + image = tf.pad(image, pad_param, constant_values=self.constant_value) + else: + transform = ResizeWithRatio(self.min_dim, self.max_dim, self.padding) + image, (bbox, str_label, int_label, image_id) = transform(sample) + return image, (bbox, str_label, int_label, image_id) + +@transform_registry(transform_type="Transpose", process="preprocess", \ + framework="onnxrt_qlinearops, onnxrt_integerops") +class Transpose(BaseTransform): + """Transpose image according to perm. + + Args: + perm (list): A permutation of the dimensions of input image + + Returns: + tuple of processed image and label + """ + + def __init__(self, perm): + """Initialize `Transpose` class.""" + self.perm = perm + + def __call__(self, sample): + """Transpose the image according to perm in sample.""" + image, label = sample + assert len(image.shape) == len(self.perm), "Image rank doesn't match Perm rank" + image = np.transpose(image, axes=self.perm) + return (image, label) + +@transform_registry(transform_type="Transpose", process="preprocess", \ + framework="tensorflow, tensorflow_itex") +class TensorflowTranspose(Transpose): + """Transpose image according to perm. + + Args: + perm (list): A permutation of the dimensions of input image + + Returns: + tuple of processed image and label + """ + + def __call__(self, sample): + """Transpose the image according to perm in sample.""" + image, label = sample + assert len(image.shape) == len(self.perm), "Image rank doesn't match Perm rank" + if isinstance(image, tf.Tensor): + image = tf.transpose(image, perm=self.perm) + else: + image = np.transpose(image, axes=self.perm) + return (image, label) + +@transform_registry(transform_type="Transpose", process="preprocess", framework="mxnet") +class MXNetTranspose(Transpose): + """Transpose image according to perm. + + Args: + perm (list): A permutation of the dimensions of input image + + Returns: + tuple of processed image and label + """ + + def __call__(self, sample): + """Transpose the image according to perm in sample.""" + image, label = sample + assert len(image.shape) == len(self.perm), "Image rank doesn't match Perm rank" + image = mx.ndarray.transpose(image, self.perm) + return (image, label) + +@transform_registry(transform_type="Transpose", process="preprocess", framework="pytorch") +class PyTorchTranspose(Transpose): + """Transpose image according to perm. + + Args: + perm (list): A permutation of the dimensions of input image + + Returns: + tuple of processed image and label + """ + + def __call__(self, sample): + """Transpose the image according to perm in sample.""" + image, label = sample + assert len(image.shape) == len(self.perm), "Image rank doesn't match Perm rank" + image = image.permute(self.perm) + return (image, label) + +@transform_registry(transform_type="RandomVerticalFlip", process="preprocess", \ + framework="onnxrt_qlinearops, onnxrt_integerops") +class RandomVerticalFlip(BaseTransform): + """Vertically flip the given image randomly. + + Returns: + tuple of processed image and label + """ + + def __call__(self, sample): + """Vertically flip the image in sample.""" + image, label = sample + if np.random.rand(1)[0] > 0.5: + image = np.flipud(image) + return (image, label) + +@transform_registry(transform_type="RandomVerticalFlip", process="preprocess", \ + framework="tensorflow, tensorflow_itex") +class TensorflowRandomVerticalFlip(BaseTransform): + """Vertically flip the given image randomly. + + Returns: + tuple of processed image and label + """ + + def __call__(self, sample): + """Vertically flip the image in sample.""" + image, label = sample + if isinstance(image, tf.Tensor): + image = tf.image.random_flip_up_down(image) + else: + if np.random.rand(1)[0] > 0.5: + image = np.flipud(image) + return (image, label) + +@transform_registry(transform_type="RandomHorizontalFlip", process="preprocess", \ + framework="onnxrt_qlinearops, onnxrt_integerops") +class RandomHorizontalFlip(BaseTransform): + """Horizontally flip the given image randomly. + + Returns: + tuple of processed image and label + """ + + def __call__(self, sample): + """Horizontally flip the image in sample.""" + image, label = sample + if np.random.rand(1)[0] > 0.5: + image = np.fliplr(image) + return (image, label) + +@transform_registry(transform_type="RandomHorizontalFlip", process="preprocess", \ + framework="tensorflow, tensorflow_itex") +class TensorflowRandomHorizontalFlip(BaseTransform): + """Horizontally flip the given image randomly. + + Returns: + tuple of processed image and label + """ + + def __call__(self, sample): + """Horizontally flip the image in sample.""" + image, label = sample + if isinstance(image, tf.Tensor): + image = tf.image.random_flip_left_right(image) + else: + if np.random.rand(1)[0] > 0.5: + image = np.fliplr(image) + return (image, label) + +@transform_registry(transform_type="ToArray", process="preprocess", \ + framework="onnxrt_qlinearops, onnxrt_integerops, tensorflow, \ + tensorflow_itex, pytorch, mxnet") +class ToArray(BaseTransform): + """Convert PIL Image or NDArray to numpy array. + + Returns: + tuple of processed image and label + """ + + def __call__(self, sample): + """Convert image in sample to numpy array.""" + from PIL import Image + image, label = sample + if isinstance(image, Image.Image): + image = np.array(image) + elif isinstance(image, mx.ndarray.NDArray): # pylint: disable=no-member + image = image.asnumpy() + else: + raise ValueError("Unknown image type!") + return (image, label) + +np_dtype_map = {'int8': np.int8, 'uint8': np.uint8, 'complex64': np.complex64, + 'uint16': np.uint16, 'int32': np.int32, 'uint32': np.uint32, + 'int64': np.int64, 'uint64': np.uint64, 'float32': np.float32, + 'float16': np.float16, 'float64': np.float64, 'bool': bool, + 'string': str, 'complex128': np.complex128, 'int16': np.int16} + +@transform_registry(transform_type="Cast", process="general", \ + framework="tensorflow, tensorflow_itex") +class CastTFTransform(BaseTransform): + """Convert image to given dtype. + + Args: + dtype (str, default='float32'): A dtype to convert image to + + Returns: + tuple of processed image and label + """ + + def __init__(self, dtype='float32'): + """Initialize `CastTFTransform` class.""" + self.tf_dtype_map = {'int16': tf.int16, 'uint8': tf.uint8, 'uint16': tf.uint16, + 'uint32':tf.uint32, 'uint64': tf.uint64, 'complex64': tf.complex64, + 'int32': tf.int32, 'int64':tf.int64, 'float32': tf.float32, + 'float16': tf.float16, 'float64':tf.float64, 'bool': tf.bool, + 'string': tf.string, 'int8': tf.int8, 'complex128': tf.complex128} + + assert dtype in self.tf_dtype_map.keys(), 'Unknown dtype' + self.dtype = dtype + + def __call__(self, sample): + """Convert image in sample to given dtype.""" + image, label = sample + if isinstance(image, tf.Tensor): + image = tf.image.convert_image_dtype(image, dtype=self.tf_dtype_map[self.dtype]) + else: + image = image.astype(np_dtype_map[self.dtype]) + return (image, label) + +@transform_registry(transform_type="Cast", process="general", + framework="onnxrt_qlinearops, onnxrt_integerops") +class CastONNXTransform(BaseTransform): + """Convert image to given dtype. + + Args: + dtype (str, default='float32'): A dtype to convert image to + + Returns: + tuple of processed image and label + """ + + def __init__(self, dtype='float32'): + """Initialize `CastONNXTransform` class.""" + assert dtype in np_dtype_map.keys(), 'Unknown dtype' + self.dtype = dtype + + def __call__(self, sample): + """Convert image in sample to given dtype.""" + image, label = sample + image = image.astype(np_dtype_map[self.dtype]) + return (image, label) + +@transform_registry(transform_type="Cast", process="general", framework="pytorch") +class CastPyTorchTransform(BaseTransform): + """Convert image to given dtype. + + Args: + dtype (str, default='float32'): A dtype to convert image to + + Returns: + tuple of processed image and label + """ + + def __init__(self, dtype='float32'): + """Initialize `CastPyTorchTransform` class.""" + dtype_map = {'int8': torch.int8, 'uint8': torch.uint8, 'complex128': torch.complex128, + 'int32':torch.int32, 'int64':torch.int64, 'complex64': torch.complex64, + 'bfloat16':torch.bfloat16, 'float64':torch.float64, 'bool': torch.bool, + 'float16':torch.float16, 'int16':torch.int16, 'float32': torch.float32} + assert dtype in dtype_map.keys(), 'Unknown dtype' + self.dtype = dtype_map[dtype] + + def __call__(self, sample): + """Convert image in sample to given dtype.""" + image, label = sample + image = image.type(self.dtype) + return (image, label) + +@transform_registry(transform_type="CenterCrop", process="preprocess", \ + framework="tensorflow, tensorflow_itex") +class CenterCropTFTransform(BaseTransform): + """Crops the given image at the center to the given size. + + Args: + size (list or int): Size of the result + + Returns: + tuple of processed image and label + """ + + def __init__(self, size): + """Initialize `CenterCropTFTransform` class.""" + if isinstance(size, int): + self.size = size, size + elif isinstance(size, list): + if len(size) == 1: + self.size = size[0], size[0] + elif len(size) == 2: + self.size = size[0], size[1] + + def __call__(self, sample): + """Crops image in sample to the given size.""" + image, label = sample + if isinstance(image, tf.Tensor): + if len(image.shape) == 3: + height, width = image.shape[0:2] + elif len(image.shape) == 4: + height, width = image.shape[1:3] + else: + raise ValueError("Unknown image shape") + if height < self.size[0] or width < self.size[1]: + raise ValueError("Target size shouldn't be lager than image size") + y0 = (height - self.size[0]) // 2 + x0 = (width - self.size[1]) // 2 + image = tf.image.crop_to_bounding_box(image, y0, x0, self.size[0], self.size[1]) + else: + transform = CenterCropTransform(self.size) + image, label = transform(sample) + return (image, label) + +@transform_registry(transform_type="PaddedCenterCrop", process="preprocess", \ + framework="tensorflow, tensorflow_itex") +class PaddedCenterCropTransform(BaseTransform): + """Crops the given image at the center to the given size with padding. + + Args: + size (list or int): Size of the result + crop_padding (int): crop padding number + + Returns: + tuple of processed image and label + """ + + def __init__(self, size, crop_padding=0): + """Initialize `PaddedCenterCropTransform` class.""" + if isinstance(size, int): + self.image_size = size + elif isinstance(size, list): + if len(size) == 1: + self.image_size = size[0] + elif len(size) == 2: + if size[0] != size[1]: + raise ValueError("'crop height must eaqual to crop width'") + self.image_size = size[0] + self.crop_padding = crop_padding + + def __call__(self, sample): + """Crops image in sample to the given size with padding.""" + image, label = sample + h, w = image.shape[0], image.shape[1] + + padded_center_crop_size = \ + int((self.image_size / (self.image_size + self.crop_padding)) * min(h, w)) + + y0 = (h - padded_center_crop_size + 1) // 2 + x0 = (w - padded_center_crop_size + 1) // 2 + image = image[y0:y0 + padded_center_crop_size, x0:x0 + padded_center_crop_size, :] + return (image, label) + +@transform_registry(transform_type="Resize", process="preprocess", \ + framework="tensorflow, tensorflow_itex") +class ResizeTFTransform(BaseTransform): + """Resize the input image to the given size. + + Args: + size (list or int): Size of the result + interpolation (str, default='bilinear'):Desired interpolation type, + support 'bilinear', 'nearest', 'bicubic' + + Returns: + tuple of processed image and label + """ + + def __init__(self, size, interpolation='bilinear'): + """Initialize `ResizeTFTransform` class.""" + if isinstance(size, int): + self.size = size, size + elif isinstance(size, list): + if len(size) == 1: + self.size = size[0], size[0] + elif len(size) == 2: + self.size = size[0], size[1] + self.interpolation = interpolation + + if self.interpolation not in ['bilinear', 'nearest', 'bicubic']: + raise ValueError('Unsupported interpolation type!') + + def __call__(self, sample): + """Resize the input image in sample to the given size.""" + image, label = sample + if isinstance(image, tf.Tensor): + image = tf.image.resize(image, self.size, method=self.interpolation) + else: + image = cv2.resize(image, self.size, + interpolation=interpolation_map[self.interpolation]) + return (image, label) + +@transform_registry(transform_type="Resize", process="preprocess", \ + framework="pytorch") +class ResizePytorchTransform(BaseTransform): + """Resize the input image to the given size. + + Args: + size (list or int): Size of the result + interpolation (str, default='bilinear'):Desired interpolation type, + support 'bilinear', 'nearest', 'bicubic' + + Returns: + tuple of processed image and label + """ + + def __init__(self, size, interpolation='bilinear'): + """Initialize `ResizePytorchTransform` class.""" + self.size = size + if interpolation in interpolation_pytorch_map.keys(): + self.interpolation = get_torchvision_map(interpolation_pytorch_map[interpolation]) + else: + raise ValueError("Undefined interpolation type") + + def __call__(self, sample): + """Resize the input image in sample to the given size.""" + image, label = sample + transformer = torchvision.transforms.Resize(size=self.size, + interpolation=self.interpolation) + return (transformer(image), label) + +@transform_registry(transform_type="RandomCrop", process="preprocess", \ + framework="tensorflow, tensorflow_itex") +class RandomCropTFTransform(BaseTransform): + """Crop the image at a random location to the given size. + + Args: + size (list or tuple or int): Size of the result + + Returns: + tuple of processed image and label + """ + + def __init__(self, size): + """Initialize `RandomCropTFTransform` class.""" + if isinstance(size, int): + self.size = size, size + elif isinstance(size, list) or isinstance(size, tuple): + if len(size) == 1: + self.size = size[0], size[0] + elif len(size) == 2: + self.size = size[0], size[1] + + def __call__(self, sample): + """Crop the image in sample to the given size.""" + image, label = sample + if isinstance(image, tf.Tensor): + if len(image.shape) == 3: + height, width = image.shape[0:2] + elif len(image.shape) == 4: + height, width = image.shape[1:3] + + if self.size[0] > height or self.size[1] > width: + raise ValueError('Crop size must be smaller than image size') + + if self.size[0] == height and self.size[1] == width: + return (image, label) + + height = tf.cast(height, dtype=tf.float32) + width = tf.cast(width, dtype=tf.float32) + offset_height = (height - self.size[0]) / 2 + offset_width = (width - self.size[1]) / 2 + offset_height = tf.cast(offset_height, dtype=tf.int32) + offset_width = tf.cast(offset_width, dtype=tf.int32) + + image = tf.image.crop_to_bounding_box(image, offset_height, + offset_width, self.size[0], self.size[1]) + else: + transform = RandomCropTransform(self.size) + image, label = transform(sample) + return (image, label) + +@transform_registry(transform_type="RandomResizedCrop", process="preprocess", \ + framework="pytorch") +class RandomResizedCropPytorchTransform(BaseTransform): + """Crop the given image to random size and aspect ratio. + + Args: + size (list or int): + Size of the result + scale (tuple or list, default=(0.08, 1.0)): + range of size of the origin size cropped + ratio (tuple or list, default=(3. / 4., 4. / 3.)): + range of aspect ratio of the origin aspect ratio cropped + interpolation (str, default='bilinear'): + Desired interpolation type, support 'bilinear', 'nearest', 'bicubic' + + Returns: + tuple of processed image and label + """ + + def __init__(self, size, scale=(0.08, 1.0), ratio=(3. / 4., 4. / 3.), + interpolation='bilinear'): + """Initialize `RandomResizedCropPytorchTransform` class.""" + self.size = size + self.scale = scale + self.ratio = ratio + + if interpolation in interpolation_pytorch_map.keys(): + self.interpolation = get_torchvision_map(interpolation_pytorch_map[interpolation]) + else: + raise ValueError("Undefined interpolation type") + + if scale[0] > scale[1] or ratio[0] > ratio[1]: + raise ValueError("Scale and ratio should be of kind (min, max)") + + def __call__(self, sample): + """Crop the image in sample to the random size.""" + image, label = sample + transformer = torchvision.transforms.RandomResizedCrop(size=self.size, + scale=self.scale, ratio=self.ratio, interpolation=self.interpolation) + return (transformer(image), label) + +@transform_registry(transform_type="RandomResizedCrop", process="preprocess", \ + framework="mxnet") +class RandomResizedCropMXNetTransform(BaseTransform): + """Crop the given image to random size and aspect ratio. + + Args: + size (list or int): + Size of the result + scale (tuple or list, default=(0.08, 1.0)): + range of size of the origin size cropped + ratio (tuple or list, default=(3. / 4., 4. / 3.)): + range of aspect ratio of the origin aspect ratio cropped + interpolation (str, default='bilinear'): + Desired interpolation type, support 'bilinear', 'nearest', 'bicubic' + + Returns: + tuple of processed image and label + """ + + def __init__(self, size, scale=(0.08, 1.0), ratio=(3. / 4., 4. / 3.), + interpolation='bilinear'): + """Initialize `RandomResizedCropMXNetTransform` class.""" + if isinstance(size, int): + self.size = size, size + elif isinstance(size, list): + if len(size) == 1: + self.size = size[0], size[0] + elif len(size) == 2: + self.size = size[1], size[0] + self.scale = scale + self.ratio = ratio + + if interpolation in interpolation_mxnet_map.keys(): + self.interpolation = interpolation_mxnet_map[interpolation] + else: + raise ValueError("Undefined interpolation type") + + if scale[0] > scale[1] or ratio[0] > ratio[1]: + raise ValueError("Scale and ratio should be of kind (min, max)") + + def __call__(self, sample): + """Crop the image in sample to the random size.""" + image, label = sample + transformer = mx.gluon.data.vision.transforms.RandomResizedCrop(size=self.size, + scale=self.scale, ratio=self.ratio, interpolation=self.interpolation) + return (transformer(image), label) + + +@transform_registry(transform_type="RandomResizedCrop", process="preprocess", \ + framework="tensorflow, tensorflow_itex") +class RandomResizedCropTFTransform(BaseTransform): + """Crop the given image to random size and aspect ratio. + + Args: + size (list or int): + Size of the result + scale (tuple or list, default=(0.08, 1.0)): + range of size of the origin size cropped + ratio (tuple or list, default=(3. / 4., 4. / 3.)): + range of aspect ratio of the origin aspect ratio cropped + interpolation (str, default='bilinear'): + Desired interpolation type, support 'bilinear', 'nearest' + + Returns: + tuple of processed image and label + """ + + def __init__(self, size, scale=(0.08, 1.0), ratio=( + 3. / 4., 4. / 3.), interpolation='bilinear'): + """Initialize `RandomResizedCropTFTransform` class.""" + if isinstance(size, int): + self.size = size, size + elif isinstance(size, list): + if len(size) == 1: + self.size = size[0], size[0] + elif len(size) == 2: + self.size = size[0], size[1] + + self.scale = scale + self.ratio = ratio + self.interpolation = interpolation + if self.interpolation not in ['bilinear', 'nearest']: + raise ValueError('Unsupported interpolation type!') + if scale[0] > scale[1] or ratio[0] > ratio[1]: + raise ValueError("Scale and ratio should be of kind (min, max)") + + def get_params(self, image, scale, ratio): + """Get the image prameters: position, height and width.""" + shape = image.shape + height = tf.cast(shape[0], dtype=tf.float32) + width = tf.cast(shape[1], dtype=tf.float32) + src_area = height * width + + for _ in range(10): + target_area = np.random.uniform(scale[0], scale[1]) * src_area + log_ratio = (np.log(ratio[0]), np.log(ratio[1])) + new_ratio = np.exp(np.random.uniform(log_ratio[0], log_ratio[1])) + + new_w = tf.math.round( + tf.math.sqrt(tf.math.multiply(target_area, new_ratio))) + new_h = tf.math.round( + tf.math.sqrt(tf.math.divide(target_area, new_ratio))) + + x0, y0 = tf.case( + [(tf.math.logical_and( + tf.math.greater(width, new_w), tf.math.greater(height, new_h)), + lambda: (tf.random.uniform( + shape=[], maxval=tf.math.subtract(width, new_w)), + tf.random.uniform( + shape=[], maxval=tf.math.subtract(height, new_h))) + )], + default=lambda: (-1.0, -1.0)) + if x0 != -1.0 and y0 != -1.0: + return y0, x0, new_h, new_w + + in_ratio = width / height + new_w, new_h = tf.case([(tf.math.greater(min(ratio), in_ratio), + lambda: (width, tf.math.round(width / min(ratio)))), + (tf.math.greater(in_ratio, max(ratio)), + lambda: (height, tf.math.round(height * max(ratio))))], + default=lambda: (width, height)) + + y0 = (height - new_h) / 2 + x0 = (width - new_w) / 2 + return y0, x0, new_h, new_w + + def __call__(self, sample): + """Crop the image in sample to the random size.""" + image, label = sample + if isinstance(image, tf.Tensor): + y0, x0, h, w = self.get_params(image, self.scale, self.ratio) + squeeze = False + if len(image.shape) == 3: + squeeze = True + image = tf.expand_dims(image, axis=0) + height, width = image.shape[1:3] + height = tf.cast(height, dtype=tf.float32) + width = tf.cast(width, dtype=tf.float32) + box_indices = tf.range(0, image.shape[0], dtype=tf.int32) + boxes = [y0/height, x0/width, (y0+h)/height, (x0+w)/width] + boxes = tf.broadcast_to(boxes, [image.shape[0], 4]) + image = tf.image.crop_and_resize(image, boxes, box_indices, + self.size, self.interpolation) + if squeeze: + image = tf.squeeze(image, axis=0) + else: + transform = RandomResizedCropTransform(self.size, self.scale, + self.ratio, self.interpolation) + image, label = transform(sample) + return (image, label) + +@transform_registry(transform_type="Normalize", process="preprocess", \ + framework="tensorflow, tensorflow_itex") +class NormalizeTFTransform(BaseTransform): + """Normalize a image with mean and standard deviation. + + Args: + mean (list, default=[0.0]): + means for each channel, if len(mean)=1, mean will be broadcasted to each channel, + otherwise its length should be same with the length of image shape + std (list, default=[1.0]): + stds for each channel, if len(std)=1, std will be broadcasted to each channel, + otherwise its length should be same with the length of image shape + + Returns: + tuple of processed image and label + """ + + def __init__(self, mean=[0.0], std=[1.0], rescale=None): + """Initialize `NormalizeTFTransform` class.""" + self.mean = mean + self.std = std + self.rescale = rescale + for item in self.std: + if item < 10**-6: + raise ValueError("Std should be greater than 0") + + def __call__(self, sample): + """Normalize the image in sample.""" + image, label = sample + if isinstance(image, tf.Tensor): + orig_dtype = image.dtype + mean = tf.broadcast_to(self.mean, tf.shape(input=image)) + mean = tf.cast(mean, dtype=image.dtype) + std = tf.broadcast_to(self.std, tf.shape(input=image)) + std = tf.cast(std, dtype=image.dtype) + image = (image - mean) / std + image = tf.cast(image, dtype=orig_dtype) + else: + transform = NormalizeTransform(self.mean, self.std) + image, label = transform(sample) + if self.rescale: + image /= self.rescale[0] + image -= self.rescale[1] + return (image, label) + +@transform_registry(transform_type='KerasRescale', process="preprocess", \ + framework='tensorflow, tensorflow_itex') +class RescaleKerasPretrainTransform(BaseTransform): + """Scale the values of image to [0,1]. + + Returns: + tuple of processed image and label + """ + + def __init__(self, rescale=None): + """Initialize `RescaleKerasPretrainTransform` class.""" + self.rescale = rescale + + def __call__(self, sample): + """Scale the values of the image in sample.""" + image, label = sample + if self.rescale: + image /= self.rescale[0] + image -= self.rescale[1] + return (image, label) + +@transform_registry(transform_type='Rescale', process="preprocess", \ + framework='tensorflow, tensorflow_itex') +class RescaleTFTransform(BaseTransform): + """Scale the values of image to [0,1]. + + Returns: + tuple of processed image and label + """ + + def __call__(self, sample): + """Scale the values of the image in sample.""" + image, label = sample + if isinstance(image, tf.Tensor): + image = tf.cast(image, tf.float32) / 255. + else: + image = image.astype('float32') / 255. + return (image, label) + +@transform_registry(transform_type='Rescale', process="preprocess", \ + framework='onnxrt_qlinearops, onnxrt_integerops') +class RescaleTransform(BaseTransform): + """Scale the values of image to [0,1]. + + Returns: + tuple of processed image and label + """ + + def __call__(self, sample): + """Scale the values of the image in sample.""" + image, label = sample + if isinstance(image, np.ndarray): + image = image.astype('float32') / 255. + return (image, label) + +@transform_registry(transform_type='AlignImageChannel', process="preprocess", \ + framework='tensorflow, tensorflow_itex, \ + onnxrt_qlinearops, onnxrt_integerops, mxnet') +class AlignImageChannelTransform(BaseTransform): + """Align image channel, now just support [H,W]->[H,W,dim], [H,W,4]->[H,W,3] and [H,W,3]->[H,W]. + + Input image must be np.ndarray. + + Returns: + tuple of processed image and label + """ + + def __init__(self, dim=3): + """Initialize `AlignImageChannelTransform` class.""" + logger.warning("This transform is going to be deprecated") + if dim < 1 or dim > 4: + raise ValueError('Unsupport image dim!') + self.dim = dim + + def __call__(self, sample): + """Align channel of the image in sample.""" + image, label = sample + if len(image.shape) == 2: + image = np.dstack([image]*self.dim) + if isinstance(image, np.ndarray) and image.shape[-1] != self.dim: + if image.shape[-1] == 4 and self.dim == 3: + image = cv2.cvtColor(image, cv2.COLOR_RGBA2RGB) + elif image.shape[-1] == 3 and self.dim == 1: + image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) + image = np.expand_dims(image, axis=-1) + else: + raise ValueError('Unsupport conversion!') + return (image, label) + +@transform_registry(transform_type='AlignImageChannel', process="preprocess", \ + framework='pytorch') +class PyTorchAlignImageChannel(BaseTransform): + """Align image channel, now just support [H,W,4]->[H,W,3] and [H,W,3]->[H,W]. + + Input image must be PIL Image. + + Returns: + tuple of processed image and label + """ + + def __init__(self, dim=3): + """Initialize `PyTorchAlignImageChannel` class.""" + logger.warning("This transform is going to be deprecated") + if dim != 1 and dim != 3: + raise ValueError('Unsupport image dim!') + self.dim = dim + + def __call__(self, sample): + """Align channel of the image in sample.""" + from PIL import Image + image, label = sample + assert isinstance(image, Image.Image), 'Input image must be PIL Image' + if self.dim == 3: + image = image.convert('RGB') + elif self.dim == 1: + image = image.convert('L') + else: + raise ValueError('Unsupport conversion!') + return (image, label) + +@transform_registry(transform_type="ToNDArray", process="preprocess", \ + framework="mxnet") +class ToNDArrayTransform(BaseTransform): + """Convert np.array to NDArray. + + Returns: + tuple of processed image and label + """ + + def __call__(self, sample): + """Convert np.array of the image in sample.""" + image, label = sample + image = mx.nd.array(image) + return image, label + +@transform_registry(transform_type="Resize", process="preprocess", framework="mxnet") +class ResizeMXNetTransform(BaseTransform): + """Resize the input image to the given size. + + Args: + size (list or int): Size of the result + interpolation (str, default='bilinear'):Desired interpolation type, + support 'bilinear', 'nearest', 'bicubic' + + Returns: + tuple of processed image and label + """ + + def __init__(self, size, interpolation='bilinear'): + """Initialize `ResizeMXNetTransform` class.""" + if isinstance(size, int): + self.size = size, size + elif isinstance(size, list): + if len(size) == 1: + self.size = size[0], size[0] + elif len(size) == 2: + self.size = size[1], size[0] + + if interpolation in interpolation_mxnet_map.keys(): + self.interpolation = interpolation_mxnet_map[interpolation] + else: + raise ValueError("Undefined interpolation type") + + def __call__(self, sample): + """Resize the input image in sample to the given size.""" + image, label = sample + transformer = mx.gluon.data.vision.transforms.Resize(size=self.size, + interpolation=self.interpolation) + return (transformer(image), label) + + +@transform_registry(transform_type="Resize", process="preprocess", \ + framework="onnxrt_qlinearops, onnxrt_integerops") +class ResizeTransform(BaseTransform): + """Resize the input image to the given size. + + Args: + size (list or int): Size of the result + interpolation (str, default='bilinear'):Desired interpolation type, + support 'bilinear', 'nearest', 'bicubic'. + + Returns: + tuple of processed image and label + """ + + def __init__(self, size, interpolation='bilinear'): + """Initialize `ResizeTransform` class.""" + if isinstance(size, int): + self.size = size, size + elif isinstance(size, list): + if len(size) == 1: + self.size = size[0], size[0] + elif len(size) == 2: + self.size = size[0], size[1] + + if interpolation in interpolation_map.keys(): + self.interpolation = interpolation_map[interpolation] + else: + raise ValueError("Undefined interpolation type") + + def __call__(self, sample): + """Resize the input image in sample to the given size.""" + image, label = sample + image = cv2.resize(image, self.size, interpolation=self.interpolation) + if len(image.shape) == 2: + image = np.expand_dims(image, -1) + return (image, label) + +@transform_registry(transform_type="CropResize", process="preprocess", \ + framework="tensorflow, tensorflow_itex") +class CropResizeTFTransform(BaseTransform): + """Crop the input image with given location and resize it. + + Args: + x (int):Left boundary of the cropping area + y (int):Top boundary of the cropping area + width (int):Width of the cropping area + height (int):Height of the cropping area + size (list or int): resize to new size after cropping + interpolation (str, default='bilinear'):Desired interpolation type, + support 'bilinear', 'nearest', 'bicubic' + + Returns: + tuple of processed image and label + """ + + def __init__(self, x, y, width, height, size, interpolation='bilinear'): + """Initialize `CropResizeTFTransform` class.""" + if interpolation not in ['bilinear', 'nearest', 'bicubic']: + raise ValueError('Unsupported interpolation type!') + self.interpolation = interpolation + self.x = x + self.y = y + self.width = width + self.height = height + if isinstance(size, int): + self.size = size, size + elif isinstance(size, list): + if len(size) == 1: + self.size = size[0], size[0] + elif len(size) == 2: + self.size = size[0], size[1] + + def __call__(self, sample): + """Resize the input image in sample with given location.""" + image, label = sample + if isinstance(image, tf.Tensor): + image = tf.image.crop_to_bounding_box( + image, self.y, self.x, self.height, self.width) + image = tf.image.resize(image, self.size, method=self.interpolation) + else: + transform = CropResizeTransform(self.x, self.y, self.width, + self.height, self.size, self.interpolation) + image, label = transform(sample) + return (image, label) + +@transform_registry(transform_type="CropResize", process="preprocess", framework="pytorch") +class PyTorchCropResizeTransform(BaseTransform): + """Crop the input image with given location and resize it. + + Args: + x (int):Left boundary of the cropping area + y (int):Top boundary of the cropping area + width (int):Width of the cropping area + height (int):Height of the cropping area + size (list or int): resize to new size after cropping + interpolation (str, default='bilinear'):Desired interpolation type, + support 'bilinear', 'nearest', 'bicubic' + + Returns: + tuple of processed image and label + """ + + def __init__(self, x, y, width, height, size, interpolation='bilinear'): + """Initialize `PyTorchCropResizeTransform` class.""" + if interpolation in interpolation_pytorch_map.keys(): + self.interpolation = get_torchvision_map(interpolation_pytorch_map[interpolation]) + else: + raise ValueError("Undefined interpolation type") + self.x = x + self.y = y + self.width = width + self.height = height + self.size = size + + def __call__(self, sample): + """Resize the input image in sample with given location.""" + image, label = sample + image = image.crop((self.x, self.y, self.x + self.width, self.y + self.height)) + transformer = torchvision.transforms.Resize(size=self.size, + interpolation=self.interpolation) + return (transformer(image), label) + +@transform_registry(transform_type="CropResize", process="preprocess", framework="mxnet") +class MXNetCropResizeTransform(BaseTransform): + """Crop the input image with given location and resize it. + + Args: + x (int):Left boundary of the cropping area + y (int):Top boundary of the cropping area + width (int):Width of the cropping area + height (int):Height of the cropping area + size (list or int): resize to new size after cropping + interpolation (str, default='bilinear'):Desired interpolation type, + support 'bilinear', 'nearest', 'bicubic' + + Returns: + tuple of processed image and label + """ + + def __init__(self, x, y, width, height, size, interpolation='bilinear'): + """Initialize `MXNetCropResizeTransform` class.""" + if interpolation in interpolation_mxnet_map.keys(): + self.interpolation = interpolation_mxnet_map[interpolation] + else: + raise ValueError("Undefined interpolation type") + self.x = x + self.y = y + self.width = width + self.height = height + self.size = size + + def __call__(self, sample): + """Resize the input image in sample with given location.""" + image, label = sample + transformer = mx.gluon.data.vision.transforms.CropResize(self.x, self.y, self.width, + self.height, self.size, self.interpolation) + return (transformer(image), label) + +@transform_registry(transform_type="CropResize", process="preprocess", \ + framework="onnxrt_qlinearops, onnxrt_integerops") +class CropResizeTransform(BaseTransform): + """Crop the input image with given location and resize it. + + Args: + x (int):Left boundary of the cropping area + y (int):Top boundary of the cropping area + width (int):Width of the cropping area + height (int):Height of the cropping area + size (list or int): resize to new size after cropping + interpolation (str, default='bilinear'):Desired interpolation type, + support 'bilinear', 'nearest', 'bicubic' + + Returns: + tuple of processed image and label + """ + + def __init__(self, x, y, width, height, size, interpolation='bilinear'): + """Initialize `CropResizeTransform` class.""" + if interpolation in interpolation_map.keys(): + self.interpolation = interpolation_map[interpolation] + else: + raise ValueError("Undefined interpolation type") + self.x = x + self.y = y + self.width = width + self.height = height + if isinstance(size, int): + self.size = size, size + elif isinstance(size, list) or isinstance(size, tuple): + if len(size) == 1: + self.size = size[0], size[0] + elif len(size) == 2: + self.size = size[0], size[1] + + def __call__(self, sample): + """Crop the input image in sample with given location.""" + image, label = sample + image = image[self.y:self.y+self.height, self.x:self.x+self.width, :] + image = cv2.resize(image, self.size, interpolation=self.interpolation) + return (image, label) + +@transform_registry(transform_type="CenterCrop", process="preprocess", \ + framework="onnxrt_qlinearops, onnxrt_integerops") +class CenterCropTransform(BaseTransform): + """Crops the given image at the center to the given size. + + Args: + size (list or int): Size of the result + + Returns: + tuple of processed image and label + """ + + def __init__(self, size): + """Initialize `CenterCropTransform` class.""" + if isinstance(size, int): + self.height, self.width = size, size + elif isinstance(size, list) or isinstance(size, tuple): + if len(size) == 1: + self.height, self.width = size[0], size[0] + elif len(size) == 2: + self.height, self.width = size[0], size[1] + + def __call__(self, sample): + """Crop the input image in sample at the center to the given size.""" + image, label = sample + h, w = image.shape[0], image.shape[1] + if h + 1 < self.height or w + 1 < self.width: + raise ValueError( + "Required crop size {} is larger then input image size {}".format( + (self.height, self.width), (h, w))) + + if self.height == h and self.width == w: + return (image, label) + + y0 = (h - self.height) // 2 + x0 = (w - self.width) // 2 + image = image[y0:y0 + self.height, x0:x0 + self.width, :] + return (image, label) + +@transform_registry(transform_type="Normalize", process="preprocess", framework="mxnet") +class MXNetNormalizeTransform(BaseTransform): + """Normalize a image with mean and standard deviation. + + Args: + mean (list, default=[0.0]): + means for each channel, if len(mean)=1, mean will be broadcasted to each channel, + otherwise its length should be same with the length of image shape + std (list, default=[1.0]): + stds for each channel, if len(std)=1, std will be broadcasted to each channel, + otherwise its length should be same with the length of image shape + + Returns: + tuple of processed image and label + """ + + def __init__(self, mean=[0.0], std=[1.0]): + """Initialize `MXNetNormalizeTransform` class.""" + self.mean = mean + self.std = std + for item in self.std: + if item < 10**-6: + raise ValueError("Std should be greater than 0") + + def __call__(self, sample): + """Normalize the image in sample.""" + image, label = sample + axes = [len(image.shape) - 1] + axes.extend(list(np.arange(len(image.shape)-1))) + image = mx.ndarray.transpose(image, axes) + assert len(self.mean) == image.shape[0], 'Mean channel must match image channel' + transformer = mx.gluon.data.vision.transforms.Normalize(self.mean, self.std) + image = transformer(image) + axes = list(np.arange(1, len(image.shape))) + axes.extend([0]) + image = mx.ndarray.transpose(image, axes) + return (image, label) + +@transform_registry(transform_type="Normalize", process="preprocess", framework="pytorch") +class PyTorchNormalizeTransform(MXNetNormalizeTransform): + """Normalize a image with mean and standard deviation. + + Args: + mean (list, default=[0.0]): + means for each channel, if len(mean)=1, mean will be broadcasted to each channel, + otherwise its length should be same with the length of image shape + std (list, default=[1.0]): + stds for each channel, if len(std)=1, std will be broadcasted to each channel, + otherwise its length should be same with the length of image shape + + Returns: + tuple of processed image and label + """ + + def __call__(self, sample): + """Normalize the image in sample.""" + image, label = sample + transformer = torchvision.transforms.Normalize(self.mean, self.std) + image = transformer(image) + return (image, label) + +@transform_registry(transform_type="Normalize", process="preprocess", \ + framework="onnxrt_qlinearops, onnxrt_integerops") +class NormalizeTransform(BaseTransform): + """Normalize a image with mean and standard deviation. + + Args: + mean (list, default=[0.0]): + means for each channel, if len(mean)=1, mean will be broadcasted to each channel, + otherwise its length should be same with the length of image shape + std (list, default=[1.0]): + stds for each channel, if len(std)=1, std will be broadcasted to each channel, + otherwise its length should be same with the length of image shape + + Returns: + tuple of processed image and label + """ + + def __init__(self, mean=[0.0], std=[1.0]): + """Initialize `NormalizeTransform` class.""" + self.mean = mean + self.std = std + for item in self.std: + if item < 10**-6: + raise ValueError("Std should be greater than 0") + + def __call__(self, sample): + """Normalize the image in sample.""" + image, label = sample + assert len(self.mean) == image.shape[-1], 'Mean channel must match image channel' + image = (image - self.mean) / self.std + return (image, label) + +@transform_registry(transform_type="RandomCrop", process="preprocess", \ + framework="mxnet, onnxrt_qlinearops, onnxrt_integerops") +class RandomCropTransform(BaseTransform): + """Crop the image at a random location to the given size. + + Args: + size (list or tuple or int): Size of the result + + Returns: + tuple of processed image and label + """ + + def __init__(self, size): + """Initialize `RandomCropTransform` class.""" + if isinstance(size, int): + self.height, self.width = size, size + elif isinstance(size, list) or isinstance(size, tuple): + if len(size) == 1: + self.height, self.width = size[0], size[0] + elif len(size) == 2: + self.height, self.width = size[0], size[1] + + def __call__(self, sample): + """Crop the image in sample to the given size.""" + image, label = sample + h, w = image.shape[0], image.shape[1] + if h + 1 < self.height or w + 1 < self.width: + raise ValueError( + "Required crop size {} is larger then input image size {}".format( + (self.height, self.width), (h, w))) + + if self.height == h and self.width == w: + return (image, label) + + rand_h = np.random.randint(0, h - self.height + 1) + rand_w = np.random.randint(0, w - self.width + 1) + if len(image.shape) == 2: + image = image[rand_h:rand_h + self.height, rand_w:rand_w + self.width] + else: + image = image[rand_h:rand_h + self.height, rand_w:rand_w + self.width, :] + return (image, label) + +@transform_registry(transform_type="RandomResizedCrop", process="preprocess", \ + framework="onnxrt_qlinearops, onnxrt_integerops") +class RandomResizedCropTransform(BaseTransform): + """Crop the given image to random size and aspect ratio. + + Args: + size (list or int): + Size of the result + scale (tuple or list, default=(0.08, 1.0)): + range of size of the origin size cropped + ratio (tuple or list, default=(3. / 4., 4. / 3.)): + range of aspect ratio of the origin aspect ratio cropped + interpolation (str, default='bilinear'): + Desired interpolation type, support 'bilinear', 'nearest' + + Returns: + tuple of processed image and label + """ + + def __init__(self, size, scale=(0.08, 1.0), ratio=( + 3. / 4., 4. / 3.), interpolation='bilinear'): + """Initialize `RandomResizedCropTransform` class.""" + if isinstance(size, int): + self.size = size, size + elif isinstance(size, list) or isinstance(size, tuple): + if len(size) == 1: + self.size = size[0], size[0] + elif len(size) == 2: + self.size = size[0], size[1] + + self.scale = scale + self.ratio = ratio + + if interpolation in interpolation_map.keys(): + self.interpolation = interpolation_map[interpolation] + else: + raise ValueError("Undefined interpolation type") + + if scale[0] > scale[1] or ratio[0] > ratio[1]: + raise ValueError("Scale and ratio should be of kind (min, max)") + + def get_params(self, image, scale, ratio): + """Get the image prameters: position, height and width.""" + h, w = image.shape[0], image.shape[1] + src_area = h * w + + for _ in range(10): + target_area = np.random.uniform(scale[0], scale[1]) * src_area + log_ratio = (np.log(ratio[0]), np.log(ratio[1])) + new_ratio = np.exp(np.random.uniform(log_ratio[0], log_ratio[1])) + + new_w = int(np.round(np.sqrt(target_area * new_ratio))) + new_h = int(np.round(np.sqrt(target_area / new_ratio))) + + if new_w < w and new_h < h: + x0 = np.random.randint(0, w - new_w) + y0 = np.random.randint(0, h - new_h) + return y0, x0, new_h, new_w + + in_ratio = float(w) / float(h) + if in_ratio < min(ratio): + new_w = w + new_h = int(round(new_w / min(ratio))) + elif in_ratio > max(ratio): + new_h = h + new_w = int(round(new_h * max(ratio))) + else: + new_w = w + new_h = h + y0 = (h - new_h) // 2 + x0 = (w - new_w) // 2 + return y0, x0, new_h, new_w + + def __call__(self, sample): + """Crop the image in sample to random size.""" + image, label = sample + y0, x0, h, w = self.get_params(image, self.scale, self.ratio) + crop_img = image[y0:y0 + h, x0:x0 + w, :] + image = cv2.resize(crop_img, self.size, interpolation=self.interpolation) + return (image, label) + +def _compute_softmax(scores): + """Compute softmax probability over raw logits.""" + import math + if not scores: + return [] + + max_score = None + for score in scores: + if max_score is None or score > max_score: + max_score = score + + exp_scores = [] + total_sum = 0.0 + for score in scores: + x = math.exp(score - max_score) + exp_scores.append(x) + total_sum += x + + probs = [] + for score in exp_scores: + probs.append(score / total_sum) + return probs + +def _get_best_indexes(logits, n_best_size): + """Get the n-best logits from a list.""" + index_and_score = sorted(enumerate(logits), key=lambda x: x[1], reverse=True) + + best_indexes = [] + for i in range(len(index_and_score)): + if i >= n_best_size: + break + best_indexes.append(index_and_score[i][0]) + return best_indexes + +def get_final_text(pred_text, orig_text, do_lower_case): + """Project the tokenized prediction back to the original text.""" + import six + from . import tokenization + def _strip_spaces(text): + ns_chars = [] + ns_to_s_map = collections.OrderedDict() + for (i, c) in enumerate(text): + if c == " ": + continue + ns_to_s_map[len(ns_chars)] = i + ns_chars.append(c) + ns_text = "".join(ns_chars) + return (ns_text, ns_to_s_map) + + tokenizer = tokenization.BasicTokenizer(do_lower_case=do_lower_case) + tok_text = " ".join(tokenizer.tokenize(orig_text)) + start_position = tok_text.find(pred_text) + if start_position == -1: + return orig_text + end_position = start_position + len(pred_text) - 1 + + (orig_ns_text, orig_ns_to_s_map) = _strip_spaces(orig_text) + (tok_ns_text, tok_ns_to_s_map) = _strip_spaces(tok_text) + + if len(orig_ns_text) != len(tok_ns_text): + return orig_text + + tok_s_to_ns_map = {} + for (i, tok_index) in six.iteritems(tok_ns_to_s_map): + tok_s_to_ns_map[tok_index] = i + + orig_start_position = None + if start_position in tok_s_to_ns_map: + ns_start_position = tok_s_to_ns_map[start_position] + if ns_start_position in orig_ns_to_s_map: + orig_start_position = orig_ns_to_s_map[ns_start_position] + + if orig_start_position is None: + return orig_text + + orig_end_position = None + if end_position in tok_s_to_ns_map: + ns_end_position = tok_s_to_ns_map[end_position] + if ns_end_position in orig_ns_to_s_map: + orig_end_position = orig_ns_to_s_map[ns_end_position] + + if orig_end_position is None: + return orig_text + + output_text = orig_text[orig_start_position:(orig_end_position + 1)] + return output_text + +class SquadExample(object): + """A single training/test example for simple sequence classification. + + For examples without an answer, the start and end position are -1. + """ + + def __init__(self, + qas_id, + question_text, + doc_tokens, + orig_answer_text=None, + start_position=None, + end_position=None, + is_impossible=False): + """Initialize `SquadExample` class.""" + self.qas_id = qas_id + self.question_text = question_text + self.doc_tokens = doc_tokens + self.orig_answer_text = orig_answer_text + self.start_position = start_position + self.end_position = end_position + self.is_impossible = is_impossible + +class InputFeatures(object): + """A single set of features of data.""" + + def __init__(self, + unique_id, + example_index, + doc_span_index, + tokens, + token_to_orig_map, + token_is_max_context, + input_ids, + input_mask, + segment_ids, + start_position=None, + end_position=None, + is_impossible=None): + """Initialize `InputFeatures` class.""" + self.unique_id = unique_id + self.example_index = example_index + self.doc_span_index = doc_span_index + self.tokens = tokens + self.token_to_orig_map = token_to_orig_map + self.token_is_max_context = token_is_max_context + self.input_ids = input_ids + self.input_mask = input_mask + self.segment_ids = segment_ids + self.start_position = start_position + self.end_position = end_position + self.is_impossible = is_impossible + +def read_squad_examples(input_file): + """Read a SQuAD json file into a list of SquadExample.""" + import json + with tf.io.gfile.GFile(input_file, "r") as reader: + input_data = json.load(reader)["data"] + + def is_whitespace(c): + if c == " " or c == "\t" or c == "\r" or c == "\n" or ord(c) == 0x202F: + return True + return False + + examples = [] + for entry in input_data: + for paragraph in entry["paragraphs"]: + paragraph_text = paragraph["context"] + doc_tokens = [] + char_to_word_offset = [] + prev_is_whitespace = True + for c in paragraph_text: + if is_whitespace(c): + prev_is_whitespace = True + else: + if prev_is_whitespace: + doc_tokens.append(c) + else: + doc_tokens[-1] += c + prev_is_whitespace = False + char_to_word_offset.append(len(doc_tokens) - 1) + + for qa in paragraph["qas"]: + qas_id = qa["id"] + question_text = qa["question"] + start_position = None + end_position = None + orig_answer_text = None + is_impossible = False + example = SquadExample( + qas_id=qas_id, + question_text=question_text, + doc_tokens=doc_tokens, + orig_answer_text=orig_answer_text, + start_position=start_position, + end_position=end_position, + is_impossible=is_impossible) + examples.append(example) + return examples + +def _check_is_max_context(doc_spans, cur_span_index, position): + """Check if this is the 'max context' doc span for the token.""" + best_score = None + best_span_index = None + for (span_index, doc_span) in enumerate(doc_spans): + end = doc_span.start + doc_span.length - 1 + if position < doc_span.start: + continue + if position > end: + continue + num_left_context = position - doc_span.start + num_right_context = end - position + score = min(num_left_context, num_right_context) + 0.01 * doc_span.length + if best_score is None or score > best_score: + best_score = score + best_span_index = span_index + + return cur_span_index == best_span_index + +def convert_examples_to_features(examples, tokenizer, max_seq_length, + doc_stride, max_query_length, output_fn): + """Load a data file into a list of `InputBatch`s.""" + unique_id = 1000000000 + for (example_index, example) in enumerate(examples): + query_tokens = tokenizer.tokenize(example.question_text) + if len(query_tokens) > max_query_length: + query_tokens = query_tokens[0:max_query_length] + + tok_to_orig_index = [] + orig_to_tok_index = [] + all_doc_tokens = [] + for (i, token) in enumerate(example.doc_tokens): + orig_to_tok_index.append(len(all_doc_tokens)) + sub_tokens = tokenizer.tokenize(token) + for sub_token in sub_tokens: + tok_to_orig_index.append(i) + all_doc_tokens.append(sub_token) + + tok_start_position = None + tok_end_position = None + + # The -3 accounts for [CLS], [SEP] and [SEP] + max_tokens_for_doc = max_seq_length - len(query_tokens) - 3 + + # We can have documents that are longer than the maximum sequence length. + # To deal with this we do a sliding window approach, where we take chunks + # of the up to our max length with a stride of `doc_stride`. + _DocSpan = collections.namedtuple( # pylint: disable=invalid-name + "DocSpan", ["start", "length"]) + doc_spans = [] + start_offset = 0 + while start_offset < len(all_doc_tokens): + length = len(all_doc_tokens) - start_offset + if length > max_tokens_for_doc: + length = max_tokens_for_doc + doc_spans.append(_DocSpan(start=start_offset, length=length)) + if start_offset + length == len(all_doc_tokens): + break + start_offset += min(length, doc_stride) + for (doc_span_index, doc_span) in enumerate(doc_spans): + tokens = [] + token_to_orig_map = {} + token_is_max_context = {} + segment_ids = [] + tokens.append("[CLS]") + segment_ids.append(0) + for token in query_tokens: + tokens.append(token) + segment_ids.append(0) + tokens.append("[SEP]") + segment_ids.append(0) + + for i in range(doc_span.length): + split_token_index = doc_span.start + i + token_to_orig_map[len(tokens)] = tok_to_orig_index[split_token_index] + + is_max_context = _check_is_max_context(doc_spans, doc_span_index, + split_token_index) + token_is_max_context[len(tokens)] = is_max_context + tokens.append(all_doc_tokens[split_token_index]) + segment_ids.append(1) + tokens.append("[SEP]") + segment_ids.append(1) + + input_ids = tokenizer.convert_tokens_to_ids(tokens) + + # The mask has 1 for real tokens and 0 for padding tokens. Only real + # tokens are attended to. + input_mask = [1] * len(input_ids) + + # Zero-pad up to the sequence length. + while len(input_ids) < max_seq_length: + input_ids.append(0) + input_mask.append(0) + segment_ids.append(0) + + assert len(input_ids) == max_seq_length + assert len(input_mask) == max_seq_length + assert len(segment_ids) == max_seq_length + + start_position = None + end_position = None + + feature = InputFeatures( + unique_id=unique_id, + example_index=example_index, + doc_span_index=doc_span_index, + tokens=tokens, + token_to_orig_map=token_to_orig_map, + token_is_max_context=token_is_max_context, + input_ids=input_ids, + input_mask=input_mask, + segment_ids=segment_ids, + start_position=start_position, + end_position=end_position, + is_impossible=example.is_impossible) + # Run callback + output_fn(feature) + unique_id += 1 + +@transform_registry(transform_type="Collect", \ + process="postprocess", framework="tensorflow") +class CollectTransform(BaseTransform): + """Postprocess the predictions, collect data.""" + + def __init__(self, length=10833): + """Initialize `CollectTransform` class.""" + self.length = length + self.unique_id = [] + self.start_logits = [] + self.end_logits = [] + self.all_sample = (None, None) + self.idx = 1000000000 + + def __call__(self, sample): + """Collect postprocess data.""" + all_results, label = sample + result_list = [np.expand_dims(result, 0) for result in all_results] + for result in result_list: + if len(self.unique_id) < self.length: + result = result.transpose(2,0,1) + self.unique_id.append(self.idx) + self.start_logits.append(result[0]) + self.end_logits.append(result[1]) + self.idx += 1 + if len(self.unique_id) == self.length: + self.all_sample = ([self.unique_id, self.start_logits, self.end_logits], label) + return self.all_sample + +@transform_registry(transform_type="SquadV1", process="postprocess", \ + framework="tensorflow, tensorflow_itex") +class TFSquadV1PostTransform(BaseTransform): + """Postprocess the predictions of bert on SQuAD. + + Args: + label_file (str): path of label file + vocab_file(str): path of vocabulary file + n_best_size (int, default=20): + The total number of n-best predictions to generate in nbest_predictions.json + max_seq_length (int, default=384): + The maximum total input sequence length after WordPiece tokenization. + Sequences longer than this will be truncated, shorter than this will be padded + max_query_length (int, default=64): + The maximum number of tokens for the question. + Questions longer than this will be truncated to this length + max_answer_length (int, default=30): + The maximum length of an answer that can be generated. This is needed because + the start and end predictions are not conditioned on one another + do_lower_case (bool, default=True): + Whether to lower case the input text. + Should be True for uncased models and False for cased models + doc_stride (int, default=128): + When splitting up a long document into chunks, + how much stride to take between chunks + + Returns: + tuple of processed prediction and label + """ + + def __init__(self, label_file, vocab_file, n_best_size=20, max_seq_length=384, \ + max_query_length=64, max_answer_length=30, do_lower_case=True, doc_stride=128): + """Initialize `TFSquadV1PostTransform` class.""" + from . import tokenization + self.eval_examples = read_squad_examples(label_file) + tokenizer = tokenization.FullTokenizer( + vocab_file=vocab_file, do_lower_case=do_lower_case) + + self.eval_features = [] + def append_feature(feature): + self.eval_features.append(feature) + + convert_examples_to_features( + examples=self.eval_examples, + tokenizer=tokenizer, + max_seq_length=max_seq_length, + doc_stride=doc_stride, + max_query_length=max_query_length, + output_fn=append_feature) + + self.n_best_size = n_best_size + self.max_answer_length = max_answer_length + self.do_lower_case = do_lower_case + self.RawResult = collections.namedtuple("RawResult", + ["unique_id", "start_logits", "end_logits"]) + + def process_result(self, results): + """Get the processed results.""" + processed_results = [] + # notice the result list sequence + for unique_id, start_logits, end_logits in zip(*results): + processed_results.append( + self.RawResult( + unique_id=int(unique_id), + start_logits=[float(x) for x in start_logits.flat], + end_logits=[float(x) for x in end_logits.flat])) + + return processed_results + + def get_postprocess_result(self, sample): + """Get the post processed results.""" + if sample == (None, None): + return (None, None) + all_results, label = sample + all_results = self.process_result(all_results) + example_index_to_features = collections.defaultdict(list) + for feature in self.eval_features: + example_index_to_features[feature.example_index].append(feature) + + unique_id_to_result = {} + for result in all_results: + unique_id_to_result[result.unique_id] = result + + _PrelimPrediction = collections.namedtuple( # pylint: disable=invalid-name + "PrelimPrediction", + ["feature_index", "start_index", "end_index", "start_logit", "end_logit"]) + + all_predictions = collections.OrderedDict() + for (example_index, example) in enumerate(self.eval_examples): + features = example_index_to_features[example_index] + + prelim_predictions = [] + # keep track of the minimum score of null start+end of position 0 + score_null = 1000000 # large and positive + min_null_feature_index = 0 # the paragraph slice with min mull score + null_start_logit = 0 # the start logit at the slice with min null score + null_end_logit = 0 # the end logit at the slice with min null score + for (feature_index, feature) in enumerate(features): + # skip the case that is not predicted + if not feature.unique_id in unique_id_to_result: + all_predictions[example.qas_id] = "*#skip this example#*" + continue + result = unique_id_to_result[feature.unique_id] + start_indexes = _get_best_indexes(result.start_logits, self.n_best_size) + end_indexes = _get_best_indexes(result.end_logits, self.n_best_size) + + for start_index in start_indexes: + for end_index in end_indexes: + # We could hypothetically create invalid predictions, e.g., predict + # that the start of the span is in the question. We throw out all + # invalid predictions. + if start_index >= len(feature.tokens): + continue + if end_index >= len(feature.tokens): + continue + if start_index not in feature.token_to_orig_map: + continue + if end_index not in feature.token_to_orig_map: + continue + if not feature.token_is_max_context.get(start_index, False): + continue + if end_index < start_index: + continue + length = end_index - start_index + 1 + if length > self.max_answer_length: + continue + prelim_predictions.append( + _PrelimPrediction( + feature_index=feature_index, + start_index=start_index, + end_index=end_index, + start_logit=result.start_logits[start_index], + end_logit=result.end_logits[end_index])) + + prelim_predictions = sorted( + prelim_predictions, + key=lambda x: (x.start_logit + x.end_logit), + reverse=True) + _NbestPrediction = collections.namedtuple( # pylint: disable=invalid-name + "NbestPrediction", ["text", "start_logit", "end_logit"]) + + seen_predictions = {} + nbest = [] + for pred in prelim_predictions: + if len(nbest) >= self.n_best_size: + break + feature = features[pred.feature_index] + if pred.start_index > 0: # this is a non-null prediction + tok_tokens = feature.tokens[pred.start_index:(pred.end_index + 1)] + orig_doc_start = feature.token_to_orig_map[pred.start_index] + orig_doc_end = feature.token_to_orig_map[pred.end_index] + orig_tokens = example.doc_tokens[orig_doc_start:(orig_doc_end + 1)] + tok_text = " ".join(tok_tokens) + + # De-tokenize WordPieces that have been split off. + tok_text = tok_text.replace(" ##", "") + tok_text = tok_text.replace("##", "") + + # Clean whitespace + tok_text = tok_text.strip() + tok_text = " ".join(tok_text.split()) + orig_text = " ".join(orig_tokens) + + final_text = get_final_text(tok_text, orig_text, self.do_lower_case) + if final_text in seen_predictions: + continue + + seen_predictions[final_text] = True + else: + final_text = "" + seen_predictions[final_text] = True + + nbest.append( + _NbestPrediction( + text=final_text, + start_logit=pred.start_logit, + end_logit=pred.end_logit)) + + # In very rare edge cases we could have no valid predictions. So we + # just create a nonce prediction in this case to avoid failure. + if not nbest: + nbest.append( + _NbestPrediction(text="empty", start_logit=0.0, end_logit=0.0)) + + assert len(nbest) >= 1 + + total_scores = [] + best_non_null_entry = None + for entry in nbest: + total_scores.append(entry.start_logit + entry.end_logit) + if not best_non_null_entry: + if entry.text: + best_non_null_entry = entry + probs = _compute_softmax(total_scores) + + nbest_json = [] + for (i, entry) in enumerate(nbest): + output = collections.OrderedDict() + output["text"] = entry.text + output["probability"] = probs[i] + output["start_logit"] = entry.start_logit + output["end_logit"] = entry.end_logit + nbest_json.append(output) + + assert len(nbest_json) >= 1 + all_predictions[example.qas_id] = nbest_json[0]["text"] + return (all_predictions, label) + + def __call__(self, sample): + """Call the get_postprocess_result.""" + return self.get_postprocess_result(sample) + + +@transform_registry(transform_type="ModelZooCollect", \ + process="postprocess", framework="tensorflow, tensorflow_itex") +class TFModelZooCollectTransform(CollectTransform): + """Postprocess the predictions of model zoo, collect data.""" + + def __call__(self, sample): + """Collect postprocess data.""" + all_results, label = sample + all_results = zip(all_results[0], all_results[1]) + for start_logits, end_logits in all_results: + if len(self.unique_id) < self.length: + self.unique_id.append(self.idx) + self.start_logits.append(start_logits) + self.end_logits.append(end_logits) + self.idx += 1 + if len(self.unique_id) == self.length: + self.all_sample = ([self.unique_id, self.start_logits, self.end_logits], label) + return self.all_sample + +@transform_registry(transform_type="SquadV1ModelZoo", \ + process="postprocess", framework="tensorflow, \ + tensorflow_itex") +class TFSquadV1ModelZooPostTransform(TFSquadV1PostTransform): + """Postprocess the predictions of bert on SQuADV1.1. + + See class TFSquadV1PostTransform for more details + """ + + def __init__(self, label_file, vocab_file, n_best_size=20, max_seq_length=384, \ + max_query_length=64, max_answer_length=30, do_lower_case=True, doc_stride=128): + """Initialize `TFSquadV1ModelZooPostTransform` class.""" + super().__init__(label_file, vocab_file, n_best_size, max_seq_length, \ + max_query_length, max_answer_length, do_lower_case, doc_stride) + self.length = len(self.eval_features) + self.collect_data = TFModelZooCollectTransform(length=self.length) + + def __call__(self, sample): + """Collect data and get postprocess results.""" + sample = self.collect_data(sample) + return self.get_postprocess_result(sample) + +@transform_registry(transform_type="ParseDecodeVoc", process="preprocess", \ + framework="tensorflow, tensorflow_itex") +class ParseDecodeVocTransform(BaseTransform): + """Parse features in Example proto. + + Returns: + tuple of parsed image and labels + """ + + def __call__(self, sample): + """Parse decode voc.""" + # Currently only supports jpeg and png. + # Need to use this logic because the shape is not known for + # tf.image.decode_image and we rely on this info to + # extend label if necessary. + def _decode_image(content, channels): + """Decode the image with content.""" + return tf.cond( + tf.image.is_jpeg(content), + lambda: tf.image.decode_jpeg(content, channels), + lambda: tf.image.decode_png(content, channels)) + + features = { + 'image/encoded': + tf.compat.v1.FixedLenFeature((), tf.string, default_value=''), + 'image/filename': + tf.compat.v1.FixedLenFeature((), tf.string, default_value=''), + 'image/format': + tf.compat.v1.FixedLenFeature((), tf.string, default_value='jpeg'), + 'image/height': + tf.compat.v1.FixedLenFeature((), tf.int64, default_value=0), + 'image/width': + tf.compat.v1.FixedLenFeature((), tf.int64, default_value=0), + 'image/segmentation/class/encoded': + tf.compat.v1.FixedLenFeature((), tf.string, default_value=''), + 'image/segmentation/class/format': + tf.compat.v1.FixedLenFeature((), tf.string, default_value='png'), + } + + parsed_features = tf.compat.v1.parse_single_example(sample, features) + + image = _decode_image(parsed_features['image/encoded'], channels=3) + + label = None + label = _decode_image( + parsed_features['image/segmentation/class/encoded'], channels=1) + + sample = { + 'image': image, + } + + label.set_shape([None, None, 1]) + + sample['labels_class'] = label + + return sample['image'], sample['labels_class'] diff --git a/neural_compressor/experimental/benchmark.py b/neural_compressor/experimental/benchmark.py index 7aff71cc14c..b822359d7e7 100644 --- a/neural_compressor/experimental/benchmark.py +++ b/neural_compressor/experimental/benchmark.py @@ -482,7 +482,7 @@ def model(self, user_model): "Please pass an original framework model but not neural compressor model!" self.framework = get_model_fwk_name(user_model) if self.framework == "tensorflow": - from ..model.model import get_model_type + from ..model.tensorflow_model import get_model_type if get_model_type(user_model) == 'keras' and cfg.model.backend == 'itex': self.framework = 'keras' if self.framework == "pytorch": diff --git a/neural_compressor/experimental/common/criterion.py b/neural_compressor/experimental/common/criterion.py index abcba32299c..0273695287a 100644 --- a/neural_compressor/experimental/common/criterion.py +++ b/neural_compressor/experimental/common/criterion.py @@ -130,7 +130,7 @@ class TensorFlowCrossEntropyLoss(object): """TensorFlow CrossEntropyLoss criterion.""" def __init__(self, param_dict): - """Initialize the DATASETS class. + """Initialize the Datasets class. Args: param_dict (dict): The dict of parameters setting by user for CrossEntropyLoss criterion. @@ -164,7 +164,7 @@ class TensorFlowSparseCategoricalCrossentropy(object): """TensorFlow SparseCategoricalCrossentropyLoss criterion.""" def __init__(self, param_dict): - """Initialize the DATASETS class. + """Initialize the Datasets class. Args: param_dict (string): param_dict. @@ -1548,4 +1548,4 @@ def __call__(self, **kwargs): class: PyTorchSelfKnowledgeDistillationLoss param dict (dict): param dict """ - return PyTorchSelfKnowledgeDistillationLoss, self._param_check() \ No newline at end of file + return PyTorchSelfKnowledgeDistillationLoss, self._param_check() diff --git a/neural_compressor/experimental/common/model.py b/neural_compressor/experimental/common/model.py index 6fec668f9e8..07f33501e84 100644 --- a/neural_compressor/experimental/common/model.py +++ b/neural_compressor/experimental/common/model.py @@ -17,7 +17,9 @@ """common Model just collects the information to construct a Model.""" -from neural_compressor.model.model import get_model_fwk_name, MODELS, get_model_type +from neural_compressor.model.model import get_model_fwk_name, MODELS +from neural_compressor.model.tensorflow_model import get_model_type +from neural_compressor.utils import logger class Model(object): """A wrapper of the information needed to construct a Model.""" diff --git a/neural_compressor/experimental/component.py b/neural_compressor/experimental/component.py index 63f9060592e..dec0cfc3cde 100644 --- a/neural_compressor/experimental/component.py +++ b/neural_compressor/experimental/component.py @@ -475,7 +475,7 @@ def model(self, user_model): "Please pass an original framework model but not neural compressor model!" self.framework = get_model_fwk_name(user_model) if self.framework == "tensorflow": - from ..model.model import get_model_type + from ..model.tensorflow_model import get_model_type if get_model_type(user_model) == 'keras' and self.cfg.model.backend == 'itex': self.framework = 'keras' if self.framework == "pytorch": diff --git a/neural_compressor/experimental/data/__init__.py b/neural_compressor/experimental/data/__init__.py index e78431a0c48..bdc10fbbff9 100644 --- a/neural_compressor/experimental/data/__init__.py +++ b/neural_compressor/experimental/data/__init__.py @@ -18,14 +18,14 @@ """Built-in dataloaders, datasets, transforms, filters for multiple framework backends.""" -from .datasets import DATASETS, Dataset, IterableDataset, dataset_registry +from .datasets import Datasets, Dataset, IterableDataset, dataset_registry from .transforms import TRANSFORMS, BaseTransform, transform_registry from .dataloaders import DATALOADERS from .filters import FILTERS, Filter, filter_registry __all__ = [ "DATALOADERS", - "DATASETS", + "Datasets", "Dataset", "IterableDataset", "dataset_registry", diff --git a/neural_compressor/experimental/data/dataloaders/pytorch_dataloader.py b/neural_compressor/experimental/data/dataloaders/pytorch_dataloader.py index fbbf80b3ee1..dcc462ae616 100644 --- a/neural_compressor/experimental/data/dataloaders/pytorch_dataloader.py +++ b/neural_compressor/experimental/data/dataloaders/pytorch_dataloader.py @@ -15,7 +15,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Initialize the DATASETS class.""" +"""Initialize the Datasets class.""" import numpy as np from neural_compressor.utils.utility import LazyImport @@ -66,4 +66,4 @@ def _generate_dataloader(self, dataset, batch_size, last_batch, collate_fn, pin_memory=pin_memory, sampler=sampler, batch_sampler=batch_sampler) - \ No newline at end of file + diff --git a/neural_compressor/experimental/data/datasets/__init__.py b/neural_compressor/experimental/data/datasets/__init__.py index 9d74d15bcec..c2460d737ed 100644 --- a/neural_compressor/experimental/data/datasets/__init__.py +++ b/neural_compressor/experimental/data/datasets/__init__.py @@ -17,7 +17,7 @@ """Built-in datasets class for multiple framework backends.""" -from .dataset import DATASETS, Dataset, IterableDataset, dataset_registry +from .dataset import Datasets, Dataset, IterableDataset, dataset_registry from os.path import dirname, basename, isfile, join import glob @@ -28,4 +28,4 @@ __import__(basename(f)[:-3], globals(), locals(), level=1) -__all__ = ["DATASETS", "Dataset", "IterableDataset", "dataset_registry"] +__all__ = ["Datasets", "Dataset", "IterableDataset", "dataset_registry"] diff --git a/neural_compressor/experimental/data/datasets/dataset.py b/neural_compressor/experimental/data/datasets/dataset.py index 2c7d602f4a6..b591ebb074f 100644 --- a/neural_compressor/experimental/data/datasets/dataset.py +++ b/neural_compressor/experimental/data/datasets/dataset.py @@ -153,12 +153,12 @@ def __getitem__(self, index): The naming convention of new dataset subclass should be something like ImageClassifier, user could choose this dataset by setting "imageclassifier" string in tuning.strategy field of yaml. - DATASETS variable is used to store all implemented Dataset subclasses to support + Datasets variable is used to store all implemented Dataset subclasses to support model specific dataset. """ -class DATASETS(object): +class Datasets(object): """A base class for all framework datasets. Args: diff --git a/neural_compressor/experimental/data/datasets/imagenet_dataset.py b/neural_compressor/experimental/data/datasets/imagenet_dataset.py index 350408437e8..8d5c52ee528 100644 --- a/neural_compressor/experimental/data/datasets/imagenet_dataset.py +++ b/neural_compressor/experimental/data/datasets/imagenet_dataset.py @@ -36,6 +36,7 @@ import numpy as np from PIL import Image from neural_compressor.utils.utility import LazyImport +from neural_compressor.utils import logger from .dataset import dataset_registry, IterableDataset, Dataset tf = LazyImport('tensorflow') mx = LazyImport('mxnet') @@ -146,3 +147,73 @@ def __getitem__(self, index): elif type(image).__name__ == 'EagerTensor': image = image.numpy() return (image, label) + +@dataset_registry(dataset_type="Imagenet", framework="tensorflow", dataset_format='') +class TensorflowImagenetDataset(IterableDataset): + """Configuration for Imagenet dataset.""" + + def __new__(cls, root, subset='validation', num_cores=28, transform=None, filter=None): + """New a imagenet dataset for tensorflow.""" + assert subset in ('validation', 'train'), \ + 'only support subset (validation, train)' + logger.warning("This api is going to be deprecated, " + "please use ImageRecord instead.") + + from tensorflow.python.platform import gfile + glob_pattern = os.path.join(root, '%s-*-of-*' % subset) + file_names = gfile.Glob(glob_pattern) + if not file_names: + raise ValueError('Found no files in --root matching: {}'.format(glob_pattern)) + + from tensorflow.python.data.experimental import parallel_interleave + from neural_compressor.data.transforms.imagenet_transform import ParseDecodeImagenet + ds = tf.data.TFRecordDataset.list_files(file_names, shuffle=False) + ds = ds.apply( + parallel_interleave( + tf.data.TFRecordDataset, cycle_length=num_cores)) + + if transform is not None: + transform.transform_list.insert(0, ParseDecodeImagenet()) + else: + transform = ParseDecodeImagenet() + ds = ds.map(transform, num_parallel_calls=None) + + ds = ds.prefetch(buffer_size=tf.data.experimental.AUTOTUNE) # this number can be tuned + return ds + +@dataset_registry(dataset_type="Imagenet", framework="onnxrt_qlinearops, \ + onnxrt_integerops", dataset_format='') +class ONNXRTImagenetDataset(Dataset): + """Configuration for Imagenet dataset.""" + + def __init__(self, root, subset='val', num_cores=28, transform=None, filter=None): + """Initialize `ONNXRTImagenetDataset` class.""" + self.val_dir = os.path.join(root, subset) + assert os.path.exists(self.val_dir), "find no val dir in {}".format(root) + \ + "please make sure there are train/val subfolders" + import glob + logger.warning("This api is going to be deprecated, " + "please use ImageRecord instead.") + + self.transform = transform + self.image_list = [] + files = glob.glob(os.path.join(self.val_dir, '*')) + files.sort() + for idx, file in enumerate(files): + imgs = glob.glob(os.path.join(file, '*')) + for img in imgs: + self.image_list.append((img, idx)) + + def __len__(self): + """Return the number of images.""" + return len(self.image_list) + + def __getitem__(self, index): + """Return the item of dataset according to the given index.""" + from PIL import Image + sample = self.image_list[index] + image = Image.open(sample[0]) + if self.transform is not None: + image, label = self.transform((image, sample[1])) + return (image, label) + diff --git a/neural_compressor/experimental/data/transforms/imagenet_transform.py b/neural_compressor/experimental/data/transforms/imagenet_transform.py index bb7bfc4a3f8..5afe6b24c06 100644 --- a/neural_compressor/experimental/data/transforms/imagenet_transform.py +++ b/neural_compressor/experimental/data/transforms/imagenet_transform.py @@ -33,6 +33,7 @@ import numpy as np from neural_compressor.utils.utility import LazyImport +from neural_compressor.utils import logger from .transform import transform_registry, BaseTransform tf = LazyImport('tensorflow') cv2 = LazyImport('cv2') @@ -132,5 +133,289 @@ def __call__(self, sample): image = features['image/encoded'] image = tf.image.decode_jpeg( image, channels=3, fancy_upscaling=False, dct_method='INTEGER_FAST') + return (image, label) + +@transform_registry(transform_type="ParseDecodeImagenet", \ + process="preprocess", framework="tensorflow") +class ParseDecodeImagenetTransform(BaseTransform): + """Imagenet decoding will be performed automatically from Neural Compressor v1.4. + + Returns: + sample + """ + + def __call__(self, sample): + """Convert `ParseDecodeImagenetTransform` feature.""" + logger.warning("This transform is going to be deprecated, " \ + "imagenet decoding will be performed automatically from Neural Compressor v1.4.") + return sample + +@transform_registry(transform_type="ResizeCropImagenet", \ + process="preprocess", framework="tensorflow") +class TensorflowResizeCropImagenetTransform(BaseTransform): + """Combination of a series of transforms which is applicable to images in Imagenet. + + Args: + height (int): Height of the result + width (int): Width of the result + random_crop (bool, default=False): whether to random crop + resize_side (int, default=256):desired shape after resize operation + random_flip_left_right (bool, default=False): whether to random flip left and right + mean_value (list, default=[0.0,0.0,0.0]):means for each channel + scale (float, default=1.0):std value + + Returns: + tuple of processed image and label + """ + + def __init__(self, height, width, random_crop=False, resize_side=256, \ + resize_method='bilinear', random_flip_left_right=False, \ + mean_value=[0.0,0.0,0.0], scale=1.0, \ + data_format='channels_last', subpixels='RGB'): + """Initialize `TensorflowResizeCropImagenetTransform` class.""" + self.height = height + self.width = width + self.mean_value = mean_value + self.scale = scale + self.random_crop = random_crop + self.random_flip_left_right = random_flip_left_right + self.resize_side = resize_side + self.resize_method = resize_method + self.data_format = data_format + self.subpixels = subpixels + + # sample is (images, labels) + def __call__(self, sample): + """Convert `TensorflowResizeCropImagenetTransform` feature.""" + image, label = sample + shape = tf.shape(input=image) + + height = tf.cast(shape[0], dtype=tf.float32) \ + if self.data_format=="channels_last" else tf.cast(shape[1], dtype=tf.float32) + width = tf.cast(shape[1], dtype=tf.float32) \ + if self.data_format=="channels_last" else tf.cast(shape[2], dtype=tf.float32) + scale = tf.cond(pred=tf.greater(height, width), \ + true_fn=lambda: self.resize_side / width, + false_fn=lambda: self.resize_side / height,) + + scale = tf.cast(scale, dtype=tf.float32) + new_height = tf.cast(tf.math.rint(height*scale), dtype=tf.int32) + new_width = tf.cast(tf.math.rint(width*scale), dtype=tf.int32) + if self.subpixels=='BGR' and self.data_format=='channels_first': + # 'RGB'->'BGR' + image = tf.cond(tf.equal(tf.rank(image), 3), + lambda: tf.experimental.numpy.moveaxis(image[::-1, ...], 0, -1), + lambda: tf.experimental.numpy.moveaxis(image[:, ::-1, ...], 1, -1)) + elif self.subpixels=='BGR': + # 'RGB'->'BGR' + image = image[..., ::-1] + image = tf.expand_dims(image, 0) + image = tf.image.resize(image, [new_height, new_width], + method=self.resize_method) + image = tf.squeeze(image) + shape = tf.shape(input=image) + if self.random_crop: + y0 = tf.random.uniform(shape=[], minval=0, maxval=(shape[0] - self.height +1), + dtype=tf.dtypes.int32) + x0 = tf.random.uniform(shape=[], minval=0, maxval=(shape[1] - self.width +1), + dtype=tf.dtypes.int32) + else: + y0 = (shape[0] - self.height) // 2 + x0 = (shape[1] - self.width) // 2 + + image = tf.image.crop_to_bounding_box(image, y0, x0, self.height, self.width) + image.set_shape([self.height, self.width, 3]) + if self.random_flip_left_right: + image = tf.image.random_flip_left_right(image) + means = tf.broadcast_to(self.mean_value, tf.shape(input=image)) + image = (image - means) * self.scale + return (image, label) + +@transform_registry(transform_type="BilinearImagenet", \ + process="preprocess", framework="tensorflow") +class BilinearImagenetTransform(BaseTransform): + """Combination of a series of transforms which is applicable to images in Imagenet. + + Args: + height: Height of the result + width:Width of the result + central_fraction(float, default=0.875):fraction of size to crop + mean_value(list, default=[0.0,0.0,0.0]):means for each channel + scale(float, default=1.0):std value + + Returns: + tuple of processed image and label + """ + + def __init__(self, height, width, central_fraction=0.875, + mean_value=[0.0,0.0,0.0], scale=1.0): + """Initialize `BilinearImagenetTransform` class.""" + self.height = height + self.width = width + self.mean_value = mean_value + self.scale = scale + self.central_fraction = central_fraction + + # sample is (images, labels) + def __call__(self, sample): + """Convert `BilinearImagenetTransform` feature.""" + image, label = sample + if image.dtype is not tf.float32: + image = tf.image.convert_image_dtype(image, dtype=tf.float32) + # Crop the central region of the image containing 87.5% area of the original image. + if self.central_fraction: + image = tf.image.central_crop(image, central_fraction=self.central_fraction) + + if self.height and self.width: + # Resize the image to the specified height and width. + image = tf.expand_dims(image, 0) + image = tf.image.resize(image, [self.height, self.width], \ + method=tf.image.ResizeMethod.BILINEAR) + image = tf.squeeze(image, [0]) + + image = tf.subtract(image, 0.5) + image = tf.multiply(image, 2.0) + means = tf.broadcast_to(self.mean_value, tf.shape(input=image)) + image = (image - means) * self.scale return (image, label) + +@transform_registry(transform_type="BilinearImagenet", process="preprocess", \ + framework="onnxrt_qlinearops, onnxrt_integerops") +class OnnxBilinearImagenetTransform(BaseTransform): + """Combination of a series of transforms which is applicable to images in Imagenet. + + Args: + height: Height of the result + width:Width of the result + central_fraction(float, default=0.875):fraction of size to crop + mean_value(list, default=[0.0,0.0,0.0]):means for each channel + scale(float, default=1.0):std value + + Returns: + tuple of processed image and label + """ + + def __init__(self, height, width, central_fraction=0.875, + mean_value=[0.0,0.0,0.0], scale=1.0): + """Initialize `OnnxBilinearImagenetTransform` class.""" + self.height = height + self.width = width + self.mean_value = mean_value + self.scale = scale + self.central_fraction = central_fraction + + def __call__(self, sample): + """Convert `OnnxBilinearImagenetTransform` feature.""" + image, label = sample + if isinstance(image, np.ndarray): + image = image.astype('float32') / 255. + img_shape = image.shape + depth = img_shape[2] + img_hd = float(img_shape[0]) + bbox_h_start = int((img_hd - img_hd * self.central_fraction) / 2) + img_wd = float(img_shape[1]) + bbox_w_start = int((img_wd - img_wd * self.central_fraction) / 2) + + bbox_h_size = img_shape[0] - bbox_h_start * 2 + bbox_w_size = img_shape[1] - bbox_w_start * 2 + + image = image[bbox_h_start:bbox_h_start+bbox_h_size, bbox_w_start:bbox_w_start+bbox_w_size] + + if self.height and self.width: + image = cv2.resize(image, (self.width, self.height), interpolation=cv2.INTER_LINEAR) + + image = np.subtract(image, 0.5) + image = np.multiply(image, 2.0) + means = np.broadcast_to(self.mean_value, image.shape) + image = (image - means) * self.scale + image = image.astype(np.float32) + return (image, label) + +@transform_registry(transform_type="ResizeCropImagenet", process="preprocess", \ + framework="onnxrt_qlinearops, onnxrt_integerops") +class ONNXResizeCropImagenetTransform(BaseTransform): + """Combination of a series of transforms which is applicable to images in Imagenet. + + Args: + height: Height of the result + width:Width of the result + central_fraction(float, default=0.875):fraction of size to crop + mean_value(list, default=[0.0,0.0,0.0]):means for each channel + scale(float, default=1.0):std value + + Returns: + tuple of processed image and label + """ + + def __init__(self, height, width, random_crop=False, resize_side=256, \ + mean_value=[0.0,0.0,0.0], std_value=[0.229, 0.224, 0.225], \ + resize_method='bilinear', data_format='channels_last', subpixels='RGB'): + """Initialize `ONNXResizeCropImagenetTransform` class.""" + self.height = height + self.width = width + self.mean_value = mean_value + self.std_value = std_value + self.random_crop = random_crop + self.resize_side = resize_side + self.resize_method = resize_method + self.data_format = data_format + self.subpixels = subpixels + + # sample is (images, labels) + def __call__(self, sample): + """Convert `ONNXResizeCropImagenetTransform` feature.""" + # TODO Support optional resize_method, data_format, subpixels for ONNX + image, label = sample + height, width = image.shape[0], image.shape[1] + scale = self.resize_side / width if height > width else self.resize_side / height + new_height = int(height*scale) + new_width = int(width*scale) + image = cv2.resize(image, (new_height, new_width)) + image = image / 255. + shape = image.shape + if self.random_crop: + y0 = np.random.randint(low=0, high=(shape[0] - self.height +1)) + x0 = np.random.randint(low=0, high=(shape[1] - self.width +1)) + else: + y0 = (shape[0] - self.height) // 2 + x0 = (shape[1] - self.width) // 2 + if len(image.shape) == 2: + image = np.array([image]) + image = np.repeat(image, 3, axis=0) + image = image.transpose(1, 2, 0) + image = image[y0:y0+self.height, x0:x0+self.width, :] + image = ((image - self.mean_value)/self.std_value).astype(np.float32) + return (image.transpose(2, 0, 1), label) + +@transform_registry(transform_type="ResizeWithAspectRatio", process="preprocess", \ + framework="onnxrt_qlinearops, onnxrt_integerops") +class ResizeWithAspectRatio(BaseTransform): + """Resize the image with aspect ratio. + + Returns: + image and label + """ + + def __init__(self, height, width, scale=87.5, inter_pol=cv2.INTER_AREA): + """Initialize `ResizeWithAspectRatio` class.""" + self.height = height + self.width = width + self.scale = scale + self.inter_pol = inter_pol + + def __call__(self, sample): + """Convert `ResizeWithAspectRatio` feature.""" + (img, label) = sample + assert len(img.shape) == 3 + height, width, _ = img.shape + new_height = int(100. * self.height / self.scale) + new_width = int(100. * self.width / self.scale) + if height > width: + w = new_width + h = int(new_height * height / width) + else: + h = new_height + w = int(new_width * width / height) + img = cv2.resize(img, (w, h), interpolation=self.inter_pol) + return img, label diff --git a/neural_compressor/experimental/graph_optimization.py b/neural_compressor/experimental/graph_optimization.py index d5b403c44da..70d65084cdf 100644 --- a/neural_compressor/experimental/graph_optimization.py +++ b/neural_compressor/experimental/graph_optimization.py @@ -188,8 +188,8 @@ def __call__(self): def dataset(self, dataset_type, *args, **kwargs): """Get dataset.""" - from .data import DATASETS - return DATASETS(self.framework)[dataset_type](*args, **kwargs) + from .data import Datasets + return Datasets(self.framework)[dataset_type](*args, **kwargs) def set_config_by_model(self, model_obj): """Set model config.""" diff --git a/neural_compressor/experimental/model_conversion.py b/neural_compressor/experimental/model_conversion.py index 30db79883df..6ab1edbd104 100644 --- a/neural_compressor/experimental/model_conversion.py +++ b/neural_compressor/experimental/model_conversion.py @@ -154,8 +154,8 @@ def dataset(self, dataset_type, *args, **kwargs): Returns: class: dataset class """ - from .data import DATASETS - return DATASETS(self.framework)[dataset_type](*args, **kwargs) + from .data import Datasets + return Datasets(self.framework)[dataset_type](*args, **kwargs) @property def source(self): diff --git a/neural_compressor/experimental/quantization.py b/neural_compressor/experimental/quantization.py index c1161d11217..82a04f7ccec 100644 --- a/neural_compressor/experimental/quantization.py +++ b/neural_compressor/experimental/quantization.py @@ -28,7 +28,8 @@ from ..utils.utility import time_limit from ..utils.create_obj_from_config import create_dataloader from ..model import BaseModel -from ..model.model import TensorflowQATModel, get_model_fwk_name +from ..model.tensorflow_model import TensorflowQATModel +from ..model.model import get_model_fwk_name from ..conf.config import QuantConf from ..conf.pythonic_config import Config from deprecated import deprecated @@ -233,8 +234,8 @@ def __call__(self): def dataset(self, dataset_type, *args, **kwargs): """Get dataset according to dataset_type.""" - from ..data import DATASETS - return DATASETS(self.framework)[dataset_type](*args, **kwargs) + from ..data import Datasets + return Datasets(self.framework)[dataset_type](*args, **kwargs) @property def calib_dataloader(self): diff --git a/neural_compressor/metric/__init__.py b/neural_compressor/metric/__init__.py index 66090298362..2866945f1c3 100644 --- a/neural_compressor/metric/__init__.py +++ b/neural_compressor/metric/__init__.py @@ -15,6 +15,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -from ..experimental.metric import METRICS, BaseMetric, metric_registry -__all__ = ["METRICS", "BaseMetric", "metric_registry"] +"""Intel Neural Compressor Metric.""" + +from .metric import METRICS, Metric, BaseMetric, metric_registry +from os.path import dirname, basename, isfile, join +import glob + +modules = glob.glob(join(dirname(__file__), "*.py")) + +for f in modules: + if isfile(f) and not f.startswith('__') and not f.endswith('__init__.py'): + __import__(basename(f)[:-3], globals(), locals(), level=1) + + +__all__ = ["METRICS", "Metric", "BaseMetric", "metric_registry"] diff --git a/neural_compressor/metric/bleu.py b/neural_compressor/metric/bleu.py new file mode 100644 index 00000000000..9a5e09df572 --- /dev/null +++ b/neural_compressor/metric/bleu.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Script for BLEU metric.""" + +import re +import six +import sys +import unicodedata +from typing import List, Sequence + +from .bleu_util import compute_bleu +from .metric import metric_registry + + +class UnicodeRegex(object): + """Ad-hoc hack to recognize all punctuation and symbols. + + Attributes: + nondigit_punct_re: The compiled regular expressions to recognize + punctuation preceded with a digit. + punct_nondigit_re: The compiled regular expressions to recognize + punctuation followed by a digit. + symbol_re: The compiled regular expressions to recognize symbols. + """ + + def __init__(self) -> None: + """Initialize the regular expressions.""" + punctuation = self.property_chars("P") + self.nondigit_punct_re = re.compile(r"([^\d])([" + punctuation + r"])") + self.punct_nondigit_re = re.compile(r"([" + punctuation + r"])([^\d])") + self.symbol_re = re.compile("([" + self.property_chars("S") + "])") + + def property_chars(self, prefix: str) -> str: + """Collect all Unicode strings starting with a specific prefix. + + Args: + prefix: The specific prefix. + + Returns: + punctuation: The join result of all Unicode strings starting + with a specific prefix. + """ + punctuation = "".join(six.unichr(x) for x in range(sys.maxunicode) \ + if unicodedata.category(six.unichr(x)).startswith(prefix)) + return punctuation + + +uregex = UnicodeRegex() + + +def bleu_tokenize(string: str) -> List[str]: + """Tokenize a string following the official BLEU implementation. + + See https://github.com/moses-smt/mosesdecoder/" + "blob/master/scripts/generic/mteval-v14.pl#L954-L983 + + Args: + string: The string to be tokenized. + + Returns: + tokens: A list of tokens. + """ + string = uregex.nondigit_punct_re.sub(r"\1 \2 ", string) + string = uregex.punct_nondigit_re.sub(r" \1 \2", string) + string = uregex.symbol_re.sub(r" \1 ", string) + tokens = string.split() + return tokens + + +@metric_registry('BLEU', 'tensorflow, tensorflow_itex') +class BLEU(object): + """Computes the BLEU (Bilingual Evaluation Understudy) score. + + BLEU is an algorithm for evaluating the quality of text which has + been machine-translated from one natural language to another. + This implementent approximate the BLEU score since we do not + glue word pieces or decode the ids and tokenize the output. + By default, we use ngram order of 4 and use brevity penalty. + Also, this does not have beam search. + + Attributes: + predictions: List of translations to score. + labels: List of the reference corresponding to the prediction result. + """ + + def __init__(self) -> None: + """Initialize predictions and labels.""" + self.predictions = [] + self.labels = [] + + def reset(self) -> None: + """Clear the predictions and labels in the cache.""" + self.predictions = [] + self.labels = [] + + def update(self, prediction: Sequence[str], label: Sequence[str]) -> None: + """Add the prediction and label. + + Args: + prediction: The prediction result. + label: The reference corresponding to the prediction result. + + Raises: + ValueError: An error occurred when the length of the prediction + and label are different. + """ + if len(label) != len(prediction): + raise ValueError("Reference and prediction files have different number " + "of lines. If training only a few steps (100-200), the " + "translation may be empty.") + label = [x.lower() for x in label] + prediction = [x.lower() for x in prediction] + label = [bleu_tokenize(x) for x in label] + prediction = [bleu_tokenize(x) for x in prediction] + self.labels.extend(label) + self.predictions.extend(prediction) + + def result(self) -> float: + """Compute the BLEU score. + + Returns: + bleu_score: The approximate BLEU score. + """ + bleu_score = compute_bleu(self.labels, self.predictions) * 100 + return bleu_score diff --git a/neural_compressor/metric/bleu_util.py b/neural_compressor/metric/bleu_util.py new file mode 100644 index 00000000000..875321f4dd3 --- /dev/null +++ b/neural_compressor/metric/bleu_util.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# Copyright 2018 The TensorFlow Authors. All Rights Reserved. +# +# Licensed under the Apache License, Version 2.0 (the 'License'); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an 'AS IS' BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# ============================================================================== +"""Script to compute BLEU score. + +Source: +https://github.com/tensorflow/tensor2tensor/blob/master/tensor2tensor/utils/bleu_hook.py +""" + +from __future__ import absolute_import +from __future__ import division +from __future__ import print_function + +import collections +import math + +import numpy as np + +from neural_compressor.utils.utility import LazyImport +from six.moves import xrange # pylint: disable=redefined-builtin +from typing import List, Sequence, Union + +tf = LazyImport('tensorflow') + +def _get_ngrams_with_counter(segment: Sequence[str], + max_order: List[int]) -> collections.Counter: + """Extract all n-grams up to a given maximum order from an input segment. + + Args: + segment: The text segment from which n-grams will be extracted. + max_order: The maximum length in tokens of the n-grams returned + by this methods. + + Returns: + ngram_counts: The Counter containing all n-grams up to max_order + in segment with a count of how many times each n-gram occurred. + """ + ngram_counts = collections.Counter() + for order in xrange(1, max_order + 1): + for i in xrange(0, len(segment) - order + 1): + ngram = tuple(segment[i:i + order]) + ngram_counts[ngram] += 1 + return ngram_counts + + +def compute_bleu(reference_corpus: Union[Sequence[str], Sequence[Sequence[str]]], + translation_corpus: Sequence[str], + max_order: int = 4, + use_bp: bool = True) -> float: + """Compute the BLEU score of translated segments against its references. + + Args: + reference_corpus: List of references for each translation. + Each reference should be tokenized into a list of tokens. + translation_corpus: List of translations to score. Each translation + should be tokenized into a list of tokens. + max_order: Maximum n-gram order to use when computing BLEU score. + use_bp: The flag to decide whether to apply brevity penalty. + + Returns: + bleu_score: The approximate BLEU score. + """ + reference_length = 0 + translation_length = 0 + bp = 1.0 + geo_mean = 0 + + matches_by_order = [0] * max_order + possible_matches_by_order = [0] * max_order + precisions = [] + + for (references, translations) in zip(reference_corpus, translation_corpus): + reference_length += len(references) + translation_length += len(translations) + ref_ngram_counts = _get_ngrams_with_counter(references, max_order) + translation_ngram_counts = _get_ngrams_with_counter(translations, max_order) + + overlap = dict((ngram, + min(count, translation_ngram_counts[ngram])) + for ngram, count in ref_ngram_counts.items()) + + for ngram in overlap: + matches_by_order[len(ngram) - 1] += overlap[ngram] + for ngram in translation_ngram_counts: + possible_matches_by_order[len(ngram) - 1] += translation_ngram_counts[ + ngram] + + precisions = [0] * max_order + smooth = 1.0 + + for i in xrange(0, max_order): + if possible_matches_by_order[i] > 0: + precisions[i] = float(matches_by_order[i]) / possible_matches_by_order[i] + if matches_by_order[i] > 0: + precisions[i] = float(matches_by_order[i]) / possible_matches_by_order[ + i] + else: + smooth *= 2 + precisions[i] = 1.0 / (smooth * possible_matches_by_order[i]) + else: + precisions[i] = 0.0 + + if max(precisions) > 0: + p_log_sum = sum(math.log(p) for p in precisions if p) + geo_mean = math.exp(p_log_sum / max_order) + + if use_bp: + ratio = translation_length / reference_length + bp = math.exp(1 - 1. / ratio) if ratio < 1.0 else 1.0 + bleu_score = np.float32(geo_mean * bp) + return bleu_score diff --git a/neural_compressor/metric/coco_label_map.py b/neural_compressor/metric/coco_label_map.py new file mode 100644 index 00000000000..82327cb6ce1 --- /dev/null +++ b/neural_compressor/metric/coco_label_map.py @@ -0,0 +1,103 @@ +# +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# + +"""The dict mapping category IDs to its names of labels.""" + +category_map = { + 1: 'person', + 2: 'bicycle', + 3: 'car', + 4: 'motorcycle', + 5: 'airplane', + 6: 'bus', + 7: 'train', + 8: 'truck', + 9: 'boat', + 10: 'traffic light', + 11: 'fire hydrant', + 13: 'stop sign', + 14: 'parking meter', + 15: 'bench', + 16: 'bird', + 17: 'cat', + 18: 'dog', + 19: 'horse', + 20: 'sheep', + 21: 'cow', + 22: 'elephant', + 23: 'bear', + 24: 'zebra', + 25: 'giraffe', + 27: 'backpack', + 28: 'umbrella', + 31: 'handbag', + 32: 'tie', + 33: 'suitcase', + 34: 'frisbee', + 35: 'skis', + 36: 'snowboard', + 37: 'sports ball', + 38: 'kite', + 39: 'baseball bat', + 40: 'baseball glove', + 41: 'skateboard', + 42: 'surfboard', + 43: 'tennis racket', + 44: 'bottle', + 46: 'wine glass', + 47: 'cup', + 48: 'fork', + 49: 'knife', + 50: 'spoon', + 51: 'bowl', + 52: 'banana', + 53: 'apple', + 54: 'sandwich', + 55: 'orange', + 56: 'broccoli', + 57: 'carrot', + 58: 'hot dog', + 59: 'pizza', + 60: 'donut', + 61: 'cake', + 62: 'chair', + 63: 'couch', + 64: 'potted plant', + 65: 'bed', + 67: 'dining table', + 70: 'toilet', + 72: 'tv', + 73: 'laptop', + 74: 'mouse', + 75: 'remote', + 76: 'keyboard', + 77: 'cell phone', + 78: 'microwave', + 79: 'oven', + 80: 'toaster', + 81: 'sink', + 82: 'refrigerator', + 84: 'book', + 85: 'clock', + 86: 'vase', + 87: 'scissors', + 88: 'teddy bear', + 89: 'hair drier', + 90: 'toothbrush' +} diff --git a/neural_compressor/metric/coco_tools.py b/neural_compressor/metric/coco_tools.py new file mode 100644 index 00000000000..93d7a34b231 --- /dev/null +++ b/neural_compressor/metric/coco_tools.py @@ -0,0 +1,713 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +"""Wrappers for third party pycocotools to be used within object_detection. + +Note that nothing in this file is tensorflow related and thus cannot +be called directly as a slim metric, for example. + +TODO(jonathanhuang): wrap as a slim metric in metrics.py + + +Usage example: given a set of images with ids in the list image_ids +and corresponding lists of numpy arrays encoding groundtruth (boxes and classes) +and detections (boxes, scores and classes), where elements of each list +correspond to detections/annotations of a single image, +then evaluation (in multi-class mode) can be invoked as follows: + + groundtruth_dict = coco_tools.ExportGroundtruthToCOCO( + image_ids, groundtruth_boxes_list, groundtruth_classes_list, + max_num_classes, output_path=None) + detections_list = coco_tools.ExportDetectionsToCOCO( + image_ids, detection_boxes_list, detection_scores_list, + detection_classes_list, output_path=None) + groundtruth = coco_tools.COCOWrapper(groundtruth_dict) + detections = groundtruth.LoadAnnotations(detections_list) + evaluator = coco_tools.COCOEvalWrapper(groundtruth, detections, + agnostic_mode=False) + metrics = evaluator.ComputeMetrics() + +""" + +import copy +import time + +import numpy as np + +from collections import OrderedDict +from neural_compressor.utils import logger +from pycocotools import coco +from pycocotools import cocoeval +from pycocotools import mask +from typing import Any, Dict, List, Set, Union + + +class COCOWrapper(coco.COCO): + """Wrapper for the pycocotools COCO class. + + Attributes: + dataset: a dictionary holding bounding box annotations in the COCO format. + detection_type: type of detections being wrapped. Can be one of ['bbox', + 'segmentation'] + """ + + def __init__(self, dataset: Dict[str, Any], detection_type: str = 'bbox'): + """Construct a COCOWrapper. + + See http://mscoco.org/dataset/#format for a description of the format. + By default, the coco.COCO class constructor reads from a JSON file. + This function duplicates the same behavior but loads from a dictionary, + allowing us to perform evaluation without writing to external storage. + + Args: + dataset: a dictionary holding bounding box annotations in the COCO format. + detection_type: type of detections being wrapped. Can be one of ['bbox', + 'segmentation'] + + Raises: + ValueError: if detection_type is unsupported. + """ + supported_detection_types = ['bbox', 'segmentation'] + if detection_type not in supported_detection_types: + raise ValueError('Unsupported detection type: {}. ' + 'Supported values are: {}'.format( + detection_type, supported_detection_types)) + self._detection_type = detection_type + coco.COCO.__init__(self) + self.dataset = dataset + self.createIndex() + + def LoadAnnotations(self, annotations: list) -> coco.COCO: + """Load annotations dictionary into COCO datastructure. + + See http://mscoco.org/dataset/#format for a description of the annotations + format. As above, this function replicates the default behavior of the API + but does not require writing to external storage. + + Args: + annotations: python list holding object detection results where each + detection is encoded as a dict with required keys ['image_id', + 'category_id', 'score'] and one of ['bbox', 'segmentation'] based on + `detection_type`. + + Returns: + a coco.COCO datastructure holding object detection annotations results + + Raises: + ValueError: if (1) annotations is not a list or annotations do not + correspond to the images contained in self. + """ + results = coco.COCO() + results.dataset['images'] = [img for img in self.dataset['images']] + + logger.info("Load and prepare annotation results.") + tic = time.time() + + if not isinstance(annotations, list): + raise ValueError('annotations is not a list of objects') + annotation_img_ids = [ann['image_id'] for ann in annotations] + if (set(annotation_img_ids) != (set(annotation_img_ids) + & set(self.getImgIds()))): + raise ValueError('Results do not correspond to current coco set') + results.dataset['categories'] = copy.deepcopy( + self.dataset['categories']) + if self._detection_type == 'bbox': + for idx, ann in enumerate(annotations): + bb = ann['bbox'] + ann['area'] = bb[2] * bb[3] + ann['id'] = idx + 1 + ann['iscrowd'] = 0 + elif self._detection_type == 'segmentation': + for idx, ann in enumerate(annotations): + ann['area'] = mask.area(ann['segmentation']) + ann['bbox'] = mask.toBbox(ann['segmentation']) + ann['id'] = idx + 1 + ann['iscrowd'] = 0 + logger.info('DONE (t=%0.2fs)', (time.time() - tic)) + + results.dataset['annotations'] = annotations + results.createIndex() + return results + + +class COCOEvalWrapper(cocoeval.COCOeval): + """Wrapper for the pycocotools COCOeval class. + + To evaluate, create two objects (groundtruth_dict and detections_list) + using the conventions listed at http://mscoco.org/dataset/#format. + Then call evaluation as follows: + + groundtruth = coco_tools.COCOWrapper(groundtruth_dict) + detections = groundtruth.LoadAnnotations(detections_list) + evaluator = coco_tools.COCOEvalWrapper(groundtruth, detections, + agnostic_mode=False) + metrics = evaluator.ComputeMetrics() + """ + + def __init__(self, + groundtruth: coco.COCO = None, + detections: coco.COCO = None, + agnostic_mode = False, + iou_type: str = 'bbox', + iou_thrs: Union[str, float] = None, + map_points=None): + """Construct a COCOEvalWrapper. + + Note that for the area-based metrics to be meaningful, detection and + groundtruth boxes must be in image coordinates measured in pixels. + + Args: + groundtruth: a coco.COCO (or coco_tools.COCOWrapper) object holding + groundtruth annotations + detections: a coco.COCO (or coco_tools.COCOWrapper) object holding + detections + agnostic_mode: boolean (default: False). If True, evaluation ignores + class labels, treating all detections as proposals. + iou_thrs: Minimal value for intersection over union that allows to + make decision that prediction bounding box is true positive. + You can specify one float value between 0 to 1 or + string "05:0.05:0.95" for standard COCO thresholds. + iou_type: IOU type to use for evaluation. Supports `bbox` or `segm`. + map_points: The way to calculate mAP. 101 for 101-point interpolated AP, 11 for + 11-point interpolated AP, 0 for area under PR curve. + """ + cocoeval.COCOeval.__init__(self, + groundtruth, + detections, + iouType=iou_type) + if agnostic_mode: + self.params.useCats = 0 + if iou_thrs == '0.5:0.05:0.95': + self.params.iouThrs = np.linspace(.5, 0.95, int(np.round((0.95 - .5) / .05)) + 1, \ + endpoint=True) + elif isinstance(iou_thrs, float): + self.params.iouThrs = [iou_thrs] + + if map_points == 101: + self.params.recThrs = np.linspace(.0, 1.00, int(np.round((1.00 - .0) / .01)) + 1, \ + endpoint=True) + if map_points == 11: + self.params.recThrs = np.linspace(.0, 1.00, int(np.round((1.00 - .0) / .1)) + 1, \ + endpoint=True) + if map_points == 0: + self.params.recThrs = [-1] + + + def GetCategory(self, category_id: int) -> dict: + """Fetch dictionary holding category information given category id. + + Args: + category_id: integer id + + Returns: + dictionary holding 'id', 'name'. + """ + return self.cocoGt.cats[category_id] + + def GetAgnosticMode(self) -> bool: + """Return whether COCO Eval is configured to evaluate in agnostic mode.""" + return self.params.useCats == 0 + + def GetCategoryIdList(self) -> List[int]: + """Return the list of IDs of all valid categories.""" + return self.params.catIds + + def accumulate(self, p: cocoeval.Params = None): + """Accumulate evaluation results per image and store it to self.eval. + + Args: + p: input params for evaluation + """ + print('Accumulating evaluation results...') + tic = time.time() + if not self.evalImgs: + print('Please run evaluate() first') + # allows input customized parameters + if p is None: + p = self.params + p.catIds = p.catIds if p.useCats == 1 else [-1] + T = len(p.iouThrs) + R = len(p.recThrs) + K = len(p.catIds) if p.useCats else 1 + A = len(p.areaRng) + M = len(p.maxDets) + precision = -np.ones((T,R,K,A,M)) # -1 for the precision of absent categories + recall = -np.ones((T,K,A,M)) + scores = -np.ones((T,R,K,A,M)) + + # create dictionary for future indexing + _pe = self._paramsEval + print('-pe', _pe) + catIds = _pe.catIds if _pe.useCats else [-1] + setK = set(catIds) + setA = set(map(tuple, _pe.areaRng)) + setM = set(_pe.maxDets) + setI = set(_pe.imgIds) + # get inds to evaluate + k_list = [n for n, k in enumerate(p.catIds) if k in setK] + m_list = [m for n, m in enumerate(p.maxDets) if m in setM] + a_list = [n for n, a in enumerate(map(lambda x: tuple(x), p.areaRng)) if a in setA] + i_list = [n for n, i in enumerate(p.imgIds) if i in setI] + I0 = len(_pe.imgIds) + A0 = len(_pe.areaRng) + # retrieve E at each category, area range, and max number of detections + for k, k0 in enumerate(k_list): + Nk = k0*A0*I0 + for a, a0 in enumerate(a_list): + Na = a0*I0 + for m, maxDet in enumerate(m_list): + E = [self.evalImgs[Nk + Na + i] for i in i_list] + E = [e for e in E if not e is None] + if len(E) == 0: continue + dtScores = np.concatenate([e['dtScores'][0:maxDet] for e in E]) + + # different sorting method generates slightly different results. + # mergesort is used to be consistent as Matlab implementation. + inds = np.argsort(-dtScores, kind='mergesort') + dtScoresSorted = dtScores[inds] + + dtm = np.concatenate([e['dtMatches'][:,0:maxDet] for e in E], axis=1)[:,inds] + dtIg = np.concatenate([e['dtIgnore'][:,0:maxDet] for e in E], axis=1)[:,inds] + gtIg = np.concatenate([e['gtIgnore'] for e in E]) + npig = np.count_nonzero(gtIg==0 ) + if npig == 0: continue + tps = np.logical_and( dtm, np.logical_not(dtIg) ) + fps = np.logical_and(np.logical_not(dtm), np.logical_not(dtIg) ) + + tp_sum = np.cumsum(tps, axis=1).astype(dtype=np.float) + fp_sum = np.cumsum(fps, axis=1).astype(dtype=np.float) + for t, (tp, fp) in enumerate(zip(tp_sum, fp_sum)): + tp = np.array(tp) + fp = np.array(fp) + nd = len(tp) + rc = tp / npig + pr = tp / (fp+tp+np.spacing(1)) + + # calculate precision + if R == 1: + rc = np.concatenate(([0.], rc, [1.])) + pr = np.concatenate(([0.], pr, [0.])) + + # compute the precision envelope + for i in range(pr.size - 1, 0, -1): + pr[i - 1] = np.maximum(pr[i - 1], pr[i]) + + # to calculate area under PR curve, look for points + # where X axis (recall) changes value + change_point = np.where(rc[1:] != rc[:-1])[0] + # and sum (\Delta recall) * recall + res = np.sum((rc[change_point + 1] - rc[change_point]) \ + * pr[change_point + 1]) + precision[t,:,k,a,m] = np.array([res]) + else: + q = np.zeros((R,)) + + # numpy is slow without cython optimization for accessing elements + # use python array gets significant speed improvement + pr = pr.tolist(); q = q.tolist() + + for i in range(nd-1, 0, -1): + if pr[i] > pr[i-1]: + pr[i-1] = pr[i] + + inds = np.searchsorted(rc, p.recThrs, side='left') + try: + for ri, pi in enumerate(inds): + q[ri] = pr[pi] + except: + pass + precision[t,:,k,a,m] = np.array(q) + + # calculate recall + if nd: + recall[t,k,a,m] = rc[-1] + else: + recall[t,k,a,m] = 0 + + # calculate score + ss = np.zeros((R,)) + inds = np.searchsorted(rc, p.recThrs, side='left') + try: + for ri, pi in enumerate(inds): + ss[ri] = dtScoresSorted[pi] + except: + pass + scores[t,:,k,a,m] = np.array(ss) + # exit(0) + self.eval = { + 'params': p, + 'counts': [T, R, K, A, M], + 'precision': precision, + 'recall': recall, + 'scores': scores, + } + toc = time.time() + print('DONE (t={:0.2f}s).'.format( toc-tic)) + + + def ComputeMetrics(self, + include_metrics_per_category: bool = False, + all_metrics_per_category: bool = False): # pragma: no cover + """Compute detection metrics. + + Args: + include_metrics_per_category: Whether include metrics per category. + all_metrics_per_category: Whether include all the summery metrics for + each category in per_category_ap. Be careful with setting it to true if + you have more than handful of categories, because it will pollute + your mldash. + + Returns: + A tuple of (summary_metrics, per_category_ap), in which + (1) summary_metrics is a dictionary holding: + 'Precision/mAP': mean average precision over classes averaged over IOU + thresholds ranging from .5 to .95 with .05 increments; + 'Precision/mAP@.50IOU': mean average precision at 50% IOU; + 'Precision/mAP@.75IOU': mean average precision at 75% IOU; + 'Precision/mAP (small)': mean average precision for small objects + (area < 32^2 pixels); + 'Precision/mAP (medium)': mean average precision for medium sized + objects (32^2 pixels < area < 96^2 pixels); + 'Precision/mAP (large)': mean average precision for large objects + (96^2 pixels < area < 10000^2 pixels); + 'Recall/AR@1': average recall with 1 detection; + 'Recall/AR@10': average recall with 10 detections; + 'Recall/AR@100': average recall with 100 detections; + 'Recall/AR@100 (small)': average recall for small objects with 100 + detections; + 'Recall/AR@100 (medium)': average recall for medium objects with 100 + detections; + 'Recall/AR@100 (large)': average recall for large objects with 100 + detections; + and (2) per_category_ap is a dictionary holding category specific results with + keys of the form: 'Precision mAP ByCategory/category' + (without the supercategory part if no supercategories exist). + + For backward compatibility 'PerformanceByCategory' is included in the + output regardless of all_metrics_per_category. If evaluating class-agnostic + mode, per_category_ap is an empty dictionary. + + Raises: + ValueError: If category_stats does not exist. + """ + self.evaluate() + self.accumulate() + self.summarize() + + summary_metrics = OrderedDict([ + ('Precision/mAP', self.stats[0]), + ('Precision/mAP@.50IOU', self.stats[1]), + ('Precision/mAP@.75IOU', self.stats[2]), + ('Precision/mAP (small)', self.stats[3]), + ('Precision/mAP (medium)', self.stats[4]), + ('Precision/mAP (large)', self.stats[5]), + ('Recall/AR@1', self.stats[6]), ('Recall/AR@10', self.stats[7]), + ('Recall/AR@100', self.stats[8]), + ('Recall/AR@100 (small)', self.stats[9]), + ('Recall/AR@100 (medium)', self.stats[10]), + ('Recall/AR@100 (large)', self.stats[11]) + ]) + if not include_metrics_per_category: + return summary_metrics, {} + if not hasattr(self, 'category_stats'): + raise ValueError('Category stats do not exist') + per_category_ap = OrderedDict([]) + if self.GetAgnosticMode(): + return summary_metrics, per_category_ap + for category_index, category_id in enumerate(self.GetCategoryIdList()): + category = self.GetCategory(category_id)['name'] + # Kept for backward compatilbility + # pylint: disable=no-member + per_category_ap['PerformanceByCategory/mAP/{}'.format( + category)] = self.category_stats[0][category_index] + if all_metrics_per_category: + per_category_ap['Precision mAP ByCategory/{}'.format( + category)] = self.category_stats[0][category_index] + per_category_ap['Precision mAP@.50IOU ByCategory/{}'.format( + category)] = self.category_stats[1][category_index] + per_category_ap['Precision mAP@.75IOU ByCategory/{}'.format( + category)] = self.category_stats[2][category_index] + per_category_ap['Precision mAP (small) ByCategory/{}'.format( + category)] = self.category_stats[3][category_index] + per_category_ap['Precision mAP (medium) ByCategory/{}'.format( + category)] = self.category_stats[4][category_index] + per_category_ap['Precision mAP (large) ByCategory/{}'.format( + category)] = self.category_stats[5][category_index] + per_category_ap['Recall AR@1 ByCategory/{}'.format( + category)] = self.category_stats[6][category_index] + per_category_ap['Recall AR@10 ByCategory/{}'.format( + category)] = self.category_stats[7][category_index] + per_category_ap['Recall AR@100 ByCategory/{}'.format( + category)] = self.category_stats[8][category_index] + per_category_ap['Recall AR@100 (small) ByCategory/{}'.format( + category)] = self.category_stats[9][category_index] + per_category_ap['Recall AR@100 (medium) ByCategory/{}'.format( + category)] = self.category_stats[10][category_index] + per_category_ap['Recall AR@100 (large) ByCategory/{}'.format( + category)] = self.category_stats[11][category_index] + + return summary_metrics, per_category_ap + + +def _ConvertBoxToCOCOFormat(box): + """Convert a box in [ymin, xmin, ymax, xmax] format to COCO format. + + This is a utility function for converting from our internal + [ymin, xmin, ymax, xmax] convention to the convention used by the COCO API + i.e., [xmin, ymin, width, height]. + + Args: + box: a numpy array in format of [ymin, xmin, ymax, xmax] + + Returns: + A list of floats, in COCO format, representing [xmin, ymin, width, height] + """ + return [ + float(box[1]), + float(box[0]), + float(box[3] - box[1]), + float(box[2] - box[0]) + ] + + +def _RleCompress(masks): + """Compresses mask using Run-length encoding provided by pycocotools. + + Args: + masks: uint8 numpy array of shape [mask_height, mask_width] with values in + {0, 1}. + + Returns: + A pycocotools Run-length encoding of the mask. + """ + return mask.encode(np.asfortranarray(masks)) + + +def ExportSingleImageGroundtruthToCoco(image_id: Union[int, str], + next_annotation_id: int, + category_id_set: Set[str], + groundtruth_boxes: np.array, + groundtruth_classes: np.array, + groundtruth_masks: Union[np.array, None] = None, + groundtruth_is_crowd: Union[np.array, None] = None) -> list: + """Export groundtruth of a single image to COCO format. + + This function converts groundtruth detection annotations represented as numpy + arrays to dictionaries that can be ingested by the COCO evaluation API. Note + that the image_ids provided here must match the ones given to + ExportSingleImageDetectionsToCoco. We assume that boxes and classes are in + correspondence - that is: groundtruth_boxes[i, :], and + groundtruth_classes[i] are associated with the same groundtruth annotation. + + In the exported result, "area" fields are always set to the area of the + groundtruth bounding box. + + Args: + image_id: a unique image identifier either of type integer or string. + next_annotation_id: integer specifying the first id to use for the + groundtruth annotations. All annotations are assigned a continuous integer + id starting from this value. + category_id_set: A set of valid class ids. Groundtruth with classes not in + category_id_set are dropped. + groundtruth_boxes: numpy array (float32) with shape [num_gt_boxes, 4] + groundtruth_classes: numpy array (int) with shape [num_gt_boxes] + groundtruth_masks: optional uint8 numpy array of shape [num_detections, + image_height, image_width] containing detection_masks. + groundtruth_is_crowd: optional numpy array (int) with shape [num_gt_boxes] + indicating whether groundtruth boxes are crowd. + + Returns: + A list of groundtruth annotations for a single image in the COCO format. + + Raises: + ValueError: if (1) groundtruth_boxes and groundtruth_classes do not have the + right lengths or (2) if each of the elements inside these lists do not + have the correct shapes or (3) if image_ids are not integers + """ + if len(groundtruth_classes.shape) != 1: + raise ValueError('groundtruth_classes is ' 'expected to be of rank 1.') + if len(groundtruth_boxes.shape) != 2: + raise ValueError('groundtruth_boxes is expected to be of ' 'rank 2.') + if groundtruth_boxes.shape[1] != 4: + raise ValueError('groundtruth_boxes should have ' 'shape[1] == 4.') + num_boxes = groundtruth_classes.shape[0] + if num_boxes != groundtruth_boxes.shape[0]: + raise ValueError( + 'Corresponding entries in groundtruth_classes, ' + 'and groundtruth_boxes should have ' + 'compatible shapes (i.e., agree on the 0th dimension).' + 'Classes shape: %d. Boxes shape: %d. Image ID: %s' % + (groundtruth_classes.shape[0], groundtruth_boxes.shape[0], + image_id)) + has_is_crowd = groundtruth_is_crowd is not None + if has_is_crowd and len(groundtruth_is_crowd.shape) != 1: + raise ValueError('groundtruth_is_crowd is expected to be of rank 1.') + groundtruth_list = [] + for i in range(num_boxes): + if groundtruth_classes[i] in category_id_set: + iscrowd = groundtruth_is_crowd[i] if has_is_crowd else 0 + export_dict = { + 'id': + next_annotation_id + i, + 'image_id': + image_id, + 'category_id': + int(groundtruth_classes[i]), + 'bbox': + list(_ConvertBoxToCOCOFormat(groundtruth_boxes[i, :])), + 'area': + float((groundtruth_boxes[i, 2] - groundtruth_boxes[i, 0]) * + (groundtruth_boxes[i, 3] - groundtruth_boxes[i, 1])), + 'iscrowd': + iscrowd + } + if groundtruth_masks is not None: + export_dict['segmentation'] = _RleCompress( + groundtruth_masks[i]) + groundtruth_list.append(export_dict) + return groundtruth_list + + +def ExportSingleImageDetectionBoxesToCoco(image_id: Union[int, str], + category_id_set: Set[int], + detection_boxes: np.array, + detection_scores: np.array, + detection_classes: np.array) -> list: + """Export detections of a single image to COCO format. + + This function converts detections represented as numpy arrays to dictionaries + that can be ingested by the COCO evaluation API. Note that the image_ids + provided here must match the ones given to the + ExporSingleImageDetectionBoxesToCoco. We assume that boxes, and classes are in + correspondence - that is: boxes[i, :], and classes[i] + are associated with the same groundtruth annotation. + + Args: + image_id: unique image identifier either of type integer or string. + category_id_set: A set of valid class ids. Detections with classes not in + category_id_set are dropped. + detection_boxes: float numpy array of shape [num_detections, 4] containing + detection boxes. + detection_scores: float numpy array of shape [num_detections] containing + scored for the detection boxes. + detection_classes: integer numpy array of shape [num_detections] containing + the classes for detection boxes. + + Returns: + A list of detection annotations for a single image in the COCO format. + + Raises: + ValueError: if (1) detection_boxes, detection_scores and detection_classes + do not have the right lengths or (2) if each of the elements inside these + lists do not have the correct shapes or (3) if image_ids are not integers. + """ + if len(detection_classes.shape) != 1 or len(detection_scores.shape) != 1: + raise ValueError( + 'All entries in detection_classes and detection_scores' + 'expected to be of rank 1.') + if len(detection_boxes.shape) != 2: + raise ValueError('All entries in detection_boxes expected to be of ' + 'rank 2.') + if detection_boxes.shape[1] != 4: + raise ValueError('All entries in detection_boxes should have ' + 'shape[1] == 4.') + num_boxes = detection_classes.shape[0] + if not num_boxes == detection_boxes.shape[0] == detection_scores.shape[0]: + raise ValueError( + 'Corresponding entries in detection_classes, ' + 'detection_scores and detection_boxes should have ' + 'compatible shapes (i.e., agree on the 0th dimension). ' + 'Classes shape: %d. Boxes shape: %d. ' + 'Scores shape: %d' % + (detection_classes.shape[0], detection_boxes.shape[0], + detection_scores.shape[0])) + detections_list = [] + for i in range(num_boxes): + if detection_classes[i] in category_id_set: + detections_list.append({ + 'image_id': + image_id, + 'category_id': + int(detection_classes[i]), + 'bbox': + list(_ConvertBoxToCOCOFormat(detection_boxes[i, :])), + 'score': + float(detection_scores[i]) + }) + return detections_list + + +def ExportSingleImageDetectionMasksToCoco(image_id: Union[str, int], + category_id_set: Set[int], + detection_masks: np.array, + detection_scores: np.array, + detection_classes: np.array) -> list: + """Export detection masks of a single image to COCO format. + + This function converts detections represented as numpy arrays to dictionaries + that can be ingested by the COCO evaluation API. We assume that + detection_masks, detection_scores, and detection_classes are in correspondence + - that is: detection_masks[i, :], detection_classes[i] and detection_scores[i] + are associated with the same annotation. + + Args: + image_id: unique image identifier either of type integer or string. + category_id_set: A set of valid class ids. Detections with classes not in + category_id_set are dropped. + detection_masks: uint8 numpy array of shape [num_detections, image_height, + image_width] containing detection_masks. + detection_scores: float numpy array of shape [num_detections] containing + scores for detection masks. + detection_classes: integer numpy array of shape [num_detections] containing + the classes for detection masks. + + Returns: + A list of detection mask annotations for a single image in the COCO format. + + Raises: + ValueError: if (1) detection_masks, detection_scores and detection_classes + do not have the right lengths or (2) if each of the elements inside these + lists do not have the correct shapes or (3) if image_ids are not integers. + """ + if len(detection_classes.shape) != 1 or len(detection_scores.shape) != 1: + raise ValueError( + 'All entries in detection_classes and detection_scores' + 'expected to be of rank 1.') + num_boxes = detection_classes.shape[0] + if not num_boxes == len(detection_masks) == detection_scores.shape[0]: + raise ValueError('Corresponding entries in detection_classes, ' + 'detection_scores and detection_masks should have ' + 'compatible lengths and shapes ' + 'Classes length: %d. Masks length: %d. ' + 'Scores length: %d' % + (detection_classes.shape[0], len(detection_masks), + detection_scores.shape[0])) + detections_list = [] + for i in range(num_boxes): + if detection_classes[i] in category_id_set: + detections_list.append({ + 'image_id': + image_id, + 'category_id': + int(detection_classes[i]), + 'segmentation': + _RleCompress(detection_masks[i]), + 'score': + float(detection_scores[i]) + }) + return detections_list diff --git a/neural_compressor/metric/evaluate_squad.py b/neural_compressor/metric/evaluate_squad.py new file mode 100644 index 00000000000..20fedd74538 --- /dev/null +++ b/neural_compressor/metric/evaluate_squad.py @@ -0,0 +1,120 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Official evaluation script for v1.1 of the SQuAD dataset. + +From https://github.com/allenai/bi-att-flow/blob/master/squad/evaluate-v1.1.py +""" + +from __future__ import print_function +import sys +from collections import Counter +from .f1 import normalize_answer + + +def f1_score(prediction, ground_truth): + """Calculate the F1 score of the prediction and the ground_truth. + + Args: + prediction: The predicted result. + ground_truth: The ground truth. + + Returns: + The F1 score of prediction. Float point number. + """ + prediction_tokens = normalize_answer(prediction).split() + ground_truth_tokens = normalize_answer(ground_truth).split() + common = Counter(prediction_tokens) & Counter(ground_truth_tokens) + num_same = sum(common.values()) + if num_same == 0: + return 0 + precision = 1.0 * num_same / len(prediction_tokens) + recall = 1.0 * num_same / len(ground_truth_tokens) + f1 = (2 * precision * recall) / (precision + recall) + return f1 + + +def metric_max_over_ground_truths(metric_fn, prediction, ground_truths): + """Calculate the max metric for each ground truth. + + For each answer in ground_truths, evaluate the metric of prediction with + this answer, and return the max metric. + + Args: + metric_fn: The function to calculate the metric. + prediction: The prediction result. + ground_truths: A list of correct answers. + + Returns: + The max metric. Float point number. + """ + scores_for_ground_truths = [] + for ground_truth in ground_truths: + score = metric_fn(prediction, ground_truth) + scores_for_ground_truths.append(score) + return max(scores_for_ground_truths) + + +def exact_match_score(prediction, ground_truth): + """Compute the exact match score between prediction and ground truth. + + Args: + prediction: The result of predictions to be evaluated. + ground_truth: The ground truth. + + Returns: + The exact match score. + """ + return (normalize_answer(prediction) == normalize_answer(ground_truth)) + + +def evaluate(dataset, predictions): + """Evaluate the average F1 score and the exact match score for Question-Answering results. + + Args: + dataset: The dataset to evaluate the prediction. A list instance of articles. + An article contains a list of paragraphs, a paragraph contains a list of + question-and-answers (qas), and a question-and-answer cantains an id, a question, + and a list of correct answers. For example: + predictions: The result of predictions to be evaluated. A dict mapping the id of + a question to the predicted answer of the question. + + Returns: + The F1 score and the exact match score. + + """ + f1 = exact_match = total = 0 + for article in dataset: + for paragraph in article['paragraphs']: + for qa in paragraph['qas']: + total += 1 + if qa['id'] not in predictions: + message = 'Unanswered question ' + qa['id'] + \ + ' will receive score 0.' + print(message, file=sys.stderr) + continue + ground_truths = list(map(lambda x: x['text'], qa['answers'])) + prediction = predictions[qa['id']] + exact_match += metric_max_over_ground_truths( + exact_match_score, prediction, ground_truths) + f1 += metric_max_over_ground_truths( + f1_score, prediction, ground_truths) + + exact_match = 100.0 * exact_match / total + f1 = 100.0 * f1 / total + + return {'exact_match': exact_match, 'f1': f1} \ No newline at end of file diff --git a/neural_compressor/metric/f1.py b/neural_compressor/metric/f1.py new file mode 100644 index 00000000000..d6b0811ae3c --- /dev/null +++ b/neural_compressor/metric/f1.py @@ -0,0 +1,144 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Official evaluation script for v1.1 of the SQuAD dataset. + +From https://github.com/allenai/bi-att-flow/blob/master/squad/evaluate-v1.1.py +""" + +from collections import Counter, abc +import string +import re +from typing import Any, Callable, Dict, List, TypeVar +from neural_compressor.utils import logger + +def normalize_answer(text: str) -> str: + """Normalize the answer text. + + Lower text, remove punctuation, articles and extra whitespace, + and replace other whitespace (newline, tab, etc.) to space. + + Args: + s: The text to be normalized. + + Returns: + The normalized text. + """ + + def _remove_articles(text): + return re.sub(r'\b(a|an|the)\b', ' ', text) + + def _white_space_fix(text): + return ' '.join(text.split()) + + def _remove_punc(text): + exclude = set(string.punctuation) + return ''.join(ch for ch in text if ch not in exclude) + + def _lower(text): + return text.lower() + + return _white_space_fix(_remove_articles(_remove_punc(_lower(text)))) + + +def f1_score(prediction: abc.Sequence, ground_truth: abc.Sequence): + """Calculate the F1 score of the prediction and the ground_truth. + + Args: + prediction: the predicted answer. + ground_truth: the correct answer. + + Returns: + The F1 score of prediction. Float point number. + """ + assert isinstance(prediction, abc.Sequence) and isinstance(ground_truth, abc.Sequence),\ + 'prediction and ground_truth should be Sequence' + common = Counter(prediction) & Counter(ground_truth) + num_same = sum(common.values()) + if num_same == 0: + return 0 + precision = 1.0 * num_same / len(prediction) + recall = 1.0 * num_same / len(ground_truth) + f1 = (2 * precision * recall) / (precision + recall) + return f1 + +T = TypeVar('T') +def metric_max_over_ground_truths(metric_fn: Callable[[T, T], float], + prediction: str, ground_truths: List[str]) -> float: + """Calculate the max metric for each ground truth. + + For each answer in ground_truths, evaluate the metric of prediction with + this answer, and return the max metric. + + Args: + metric_fn: the function to calculate the metric. + prediction: the prediction result. + ground_truths: the list of correct answers. + + Returns: + The max metric. Float point number. + """ + scores_for_ground_truths = [] + for ground_truth in ground_truths: + prediction_tokens = normalize_answer(prediction).split() + ground_truth_tokens = normalize_answer(ground_truth).split() + score = metric_fn(prediction_tokens, ground_truth_tokens) + scores_for_ground_truths.append(score) + return max(scores_for_ground_truths) + +def evaluate(predictions: Dict[str, str], dataset: List[Dict[str, Any]]) -> float: + """Evaluate the average F1 score of Question-Answering results. + + The F1 score is the harmonic mean of the precision and recall. It can be computed + with the equation: F1 = 2 * (precision * recall) / (precision + recall). + For all question-and-answers in dataset, it evaluates the f1-score + + Args: + predictions: The result of predictions to be evaluated. A dict mapping the id of + a question to the predicted answer of the question. + dataset: The dataset to evaluate the prediction. A list instance of articles. + An article contains a list of paragraphs, a paragraph contains a list of + question-and-answers (qas), and a question-and-answer cantains an id, a question, + and a list of correct answers. For example: + + [{'paragraphs': + [{'qas':[{'answers': [{'answer_start': 177, 'text': 'Denver Broncos'}, ...], + 'question': 'Which NFL team represented the AFC at Super Bowl 50?', + 'id': '56be4db0acb8001400a502ec'}]}]}] + + Returns: + The F1 score of this prediction. Float point number in forms of a percentage. + """ + f1 = total = 0 + for article in dataset: + for paragraph in article['paragraphs']: + for qa in paragraph['qas']: + total += 1 + if qa['id'] not in predictions: + message = 'Unanswered question ' + qa['id'] + \ + ' will receive score 0.' + logger.warning(message) + continue + + ground_truths = list(map(lambda x: x['text'], qa['answers'])) + prediction = predictions[qa['id']] + + f1 += metric_max_over_ground_truths( + f1_score, prediction, ground_truths) + + f1 = 100.0 * f1 / total + return f1 diff --git a/neural_compressor/metric/metric.py b/neural_compressor/metric/metric.py new file mode 100644 index 00000000000..a4786e63ab8 --- /dev/null +++ b/neural_compressor/metric/metric.py @@ -0,0 +1,1613 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Neural Compressor metrics.""" + + + +import numpy as np +from abc import abstractmethod +from ctypes import Union +from neural_compressor.utils.utility import LazyImport, singleton +from neural_compressor.utils import logger +from sklearn.metrics import accuracy_score + +torch = LazyImport('torch') +tf = LazyImport('tensorflow') +mx = LazyImport('mxnet') +transformers = LazyImport('transformers') + +class Metric(object): + """A wrapper of the information needed to construct a Metric. + + The metric class should take the outputs of the model as the metric's inputs, + neural_compressor built-in metric always take (predictions, labels) as inputs, it's + recommended to design metric_cls to take (predictions, labels) as inputs. + """ + + def __init__(self, metric_cls, name='user_metric', **kwargs): + """Initialize a Metric with needed information. + + Args: + metric_cls (cls): Should be a sub_class of neural_compressor.metric.BaseMetric, + which takes (predictions, labels) as inputs + name (str, optional): Name for metric. Defaults to 'user_metric'. + """ + self.metric_cls = metric_cls + self.name = name + self.kwargs = kwargs + +@singleton +class TensorflowMetrics(object): + """Tensorflow metrics collection. + + Attributes: + metrics: A dict to maintain all metrics for Tensorflow model. + """ + + def __init__(self) -> None: + """Initialize the metrics collection.""" + self.metrics = {} + self.metrics.update(TENSORFLOW_METRICS) + + +@singleton +class PyTorchMetrics(object): + """PyTorch metrics collection. + + Attributes: + metrics: A dict to maintain all metrics for PyTorch model. + """ + + def __init__(self) -> None: + """Initialize the metrics collection.""" + self.metrics = {} + self.metrics.update(PYTORCH_METRICS) + + +@singleton +class MXNetMetrics(object): + """MXNet metrics collection. + + Attributes: + metrics: A dict to maintain all metrics for MXNet model. + """ + + def __init__(self) -> None: + """Initialize the metrics collection.""" + from neural_compressor.adaptor.mxnet_utils.util import check_mx_version + if check_mx_version('2.0.0'): + import mxnet.gluon.metric as mx_metrics + else: + import mxnet.metric as mx_metrics + self.metrics = { + "Accuracy": WrapMXNetMetric(mx_metrics.Accuracy), + "MAE": WrapMXNetMetric(mx_metrics.MAE), + "MSE": WrapMXNetMetric(mx_metrics.MSE), + "Loss": WrapMXNetMetric(mx_metrics.Loss), + } + self.metrics.update(MXNET_METRICS) + + +@singleton +class ONNXRTQLMetrics(object): + """ONNXRT QLinear metrics collection. + + Attributes: + metrics: A dict to maintain all metrics for ONNXRT QLinear model. + """ + + def __init__(self) -> None: + """Initialize the metrics collection.""" + self.metrics = {} + self.metrics.update(ONNXRT_QL_METRICS) + + +@singleton +class ONNXRTITMetrics(object): + """ONNXRT Integer metrics collection. + + Attributes: + metrics: A dict to maintain all metrics for ONNXRT Integer model. + """ + + def __init__(self) -> None: + """Initialize the metrics collection.""" + self.metrics = {} + self.metrics.update(ONNXRT_IT_METRICS) + + +framework_metrics = {"tensorflow": TensorflowMetrics, + "tensorflow_itex": TensorflowMetrics, + "mxnet": MXNetMetrics, + "pytorch": PyTorchMetrics, + "pytorch_ipex": PyTorchMetrics, + "pytorch_fx": PyTorchMetrics, + "onnxrt_qlinearops": ONNXRTQLMetrics, + "onnxrt_integerops": ONNXRTITMetrics, + "onnxrt_qdq": ONNXRTQLMetrics, + "onnxrt_qoperator": ONNXRTQLMetrics} + +# user/model specific metrics will be registered here +TENSORFLOW_METRICS = {} +TENSORFLOW_ITEX_METRICS = {} +MXNET_METRICS = {} +PYTORCH_METRICS = {} +ONNXRT_QL_METRICS = {} +ONNXRT_IT_METRICS = {} + +registry_metrics = {"tensorflow": TENSORFLOW_METRICS, + "tensorflow_itex": TENSORFLOW_ITEX_METRICS, + "mxnet": MXNET_METRICS, + "pytorch": PYTORCH_METRICS, + "pytorch_ipex": PYTORCH_METRICS, + "pytorch_fx": PYTORCH_METRICS, + "onnxrt_qlinearops": ONNXRT_QL_METRICS, + "onnxrt_qdq": ONNXRT_QL_METRICS, + "onnxrt_integerops": ONNXRT_IT_METRICS, + "onnxrt_qoperator": ONNXRT_QL_METRICS, + } + + +class METRICS(object): + """Intel Neural Compressor Metrics. + + Attributes: + metrics: The collection of registered metrics for the specified framework. + """ + + def __init__(self, framework: str): + """Initialize the metrics collection based on the framework name. + + Args: + framework: The framwork name. + """ + assert framework in ("tensorflow", "tensorflow_itex", + "pytorch", "pytorch_ipex", "pytorch_fx", "onnxrt_qdq", + "onnxrt_qlinearops", "onnxrt_integerops", "mxnet", + "onnxrt_qoperator"), \ + "framework support tensorflow pytorch mxnet onnxrt" + self.metrics = framework_metrics[framework]().metrics + + def __getitem__(self, metric_type: str): + """Get the metric based on the specified type. + + Args: + metric_type: The metric type. + + Returns: + The metric with the specified type. + """ + assert metric_type in self.metrics.keys(), "only support metrics in {}".\ + format(self.metrics.keys()) + + return self.metrics[metric_type] + + def register(self, name, metric_cls) -> None: + """Register a metric. + + Args: + name: The name of metric. + metric_cls: The metric class. + """ + assert name not in self.metrics.keys(), 'registered metric name already exists.' + self.metrics.update({name: metric_cls}) + +def metric_registry(metric_type: str, framework: str): + """Decorate for registering all Metric subclasses. + + The cross-framework metric is supported by specifying the framework param + as one of tensorflow, pytorch, mxnet, onnxrt. + + Args: + metric_type: The metric type. + framework: The framework name. + + Returns: + decorator_metric: The function to register metric class. + """ + + def decorator_metric(cls): + for single_framework in [fwk.strip() for fwk in framework.split(',')]: + assert single_framework in [ + "tensorflow", + "tensorflow_itex", + "mxnet", + "onnxrt_qlinearops", + "onnxrt_integerops", + "onnxrt_qdq", + "onnxrt_qoperator", + "pytorch", + "pytorch_ipex", + "pytorch_fx", + ], "The framework support tensorflow mxnet pytorch onnxrt" + + if metric_type in registry_metrics[single_framework].keys(): + raise ValueError('Cannot have two metrics with the same name') + registry_metrics[single_framework][metric_type] = cls + return cls + return decorator_metric + + +class BaseMetric(object): + """The base class of Metric.""" + + def __init__(self, metric, single_output = False, hvd = None): + """Initialize the basic metric. + + Args: + metric: The metric class. + single_output: Whether the output is single or not, defaults to False. + hvd: The Horovod class for distributed trainig, defaults to None. + """ + self._metric_cls = metric + self._single_output = single_output + self._hvd = hvd + + def __call__(self, *args, **kwargs): + """Evaluate the model predictions, and the reference. + + Returns: + The class itself. + """ + self._metric = self._metric_cls(*args, **kwargs) + return self + + @abstractmethod + def update(self, preds, labels=None, sample_weight=None): + """Update the state that need to be evaluated. + + Args: + preds: The prediction result. + labels: The reference. Defaults to None. + sample_weight: The sampling weight. Defaults to None. + + Raises: + NotImplementedError: The method should be implemented by subclass. + """ + raise NotImplementedError + + @abstractmethod + def reset(self): + """Clear the predictions and labels. + + Raises: + NotImplementedError: The method should be implemented by subclass. + """ + raise NotImplementedError + + @abstractmethod + def result(self): + """Evaluate the difference between predictions and labels. + + Raises: + NotImplementedError: The method should be implemented by subclass. + """ + raise NotImplementedError + + @property + def metric(self): + """Return its metric class. + + Returns: + The metric class. + """ + return self._metric + + @property + def hvd(self): + """Return its hvd class. + + Returns: + The hvd class. + """ + return self._hvd + + @hvd.setter + def hvd(self, hvd): + """Set its hvd. + + Args: + hvd: The Horovod class for distributed trainig. + """ + self._hvd = hvd + + +class WrapPyTorchMetric(BaseMetric): + """The wrapper of Metric class for PyTorch.""" + + def update(self, preds, labels=None, sample_weight=None): + """Convert the prediction to torch. + + Args: + preds: The prediction result. + labels: The reference. Defaults to None. + sample_weight: The sampling weight. Defaults to None. + """ + if self._single_output: + output = torch.as_tensor(preds) + else: + output = (torch.as_tensor(preds), torch.as_tensor(labels)) + self._metric.update(output) + + def reset(self): + """Clear the predictions and labels.""" + self._metric.reset() + + def result(self): + """Evaluate the difference between predictions and labels.""" + return self._metric.compute() + + +class WrapMXNetMetric(BaseMetric): + """The wrapper of Metric class for MXNet.""" + + def update(self, preds, labels=None, sample_weight=None): + """Convert the prediction to MXNet array. + + Args: + preds: The prediction result. + labels: The reference. Defaults to None. + sample_weight: The sampling weight. Defaults to None. + """ + preds = mx.nd.array(preds) + labels = mx.nd.array(labels) + self._metric.update(labels=labels, preds=preds) + + def reset(self): + """Clear the predictions and labels.""" + self._metric.reset() + + def result(self): + """Evaluate the difference between predictions and labels. + + Returns: + acc: The evaluated result. + """ + acc_name, acc = self._metric.get() + return acc + +class WrapONNXRTMetric(BaseMetric): + """The wrapper of Metric class for ONNXRT.""" + + def update(self, preds, labels=None, sample_weight=None): + """Convert the prediction to NumPy array. + + Args: + preds: The prediction result. + labels: The reference. Defaults to None. + sample_weight: The sampling weight. Defaults to None. + """ + preds = np.array(preds) + labels = np.array(labels) + self._metric.update(labels=labels, preds=preds) + + def reset(self): + """Clear the predictions and labels.""" + self._metric.reset() + + def result(self): + """Evaluate the difference between predictions and labels. + + Returns: + acc: The evaluated result. + """ + acc_name, acc = self._metric.get() + return acc + +def _topk_shape_validate(preds, labels): + # preds shape can be Nxclass_num or class_num(N=1 by default) + # it's more suitable for 'Accuracy' with preds shape Nx1(or 1) output from argmax + if isinstance(preds, int): + preds = [preds] + preds = np.array(preds) + elif isinstance(preds, np.ndarray): + preds = np.array(preds) + elif isinstance(preds, list): + preds = np.array(preds) + preds = preds.reshape((-1, preds.shape[-1])) + + # consider labels just int value 1x1 + if isinstance(labels, int): + labels = [labels] + labels = np.array(labels) + elif isinstance(labels, tuple): + labels = np.array([labels]) + labels = labels.reshape((labels.shape[-1], -1)) + elif isinstance(labels, list): + if isinstance(labels[0], int): + labels = np.array(labels) + labels = labels.reshape((labels.shape[0], 1)) + elif isinstance(labels[0], tuple): + labels = np.array(labels) + labels = labels.reshape((labels.shape[-1], -1)) + else: + labels = np.array(labels) + # labels most have 2 axis, 2 cases: N(or Nx1 sparse) or Nxclass_num(one-hot) + # only support 2 dimension one-shot labels + # or 1 dimension one-hot class_num will confuse with N + + if len(preds.shape) == 1: + N = 1 + class_num = preds.shape[0] + preds = preds.reshape([-1, class_num]) + elif len(preds.shape) >= 2: + N = preds.shape[0] + preds = preds.reshape([N, -1]) + class_num = preds.shape[1] + + label_N = labels.shape[0] + assert label_N == N, 'labels batch size should same with preds' + labels = labels.reshape([N, -1]) + # one-hot labels will have 2 dimension not equal 1 + if labels.shape[1] != 1: + labels = labels.argsort()[..., -1:] + return preds, labels + +def _shape_validate(preds, labels): + assert type(preds) in [int, list, np.ndarray], 'preds must be in int or list, ndarray' + assert type(labels) in [int, list, np.ndarray], 'labels must be in int or list, ndarray' + if isinstance(preds, int): + preds = [np.array([preds])] + elif isinstance(preds[0], int): + preds = [np.array(preds)] + else: + preds = [np.array(pred) for pred in preds] + if isinstance(labels, int): + labels = [np.array([labels])] + elif isinstance(labels[0], int): + labels = [np.array(labels)] + else: + labels = [np.array(label) for label in labels] + for (pred, label) in zip(preds, labels): + assert pred.shape == label.shape, \ + 'Shape mismatch, label shape {} vs pred shape {}'.format(label.shape, pred.shape) + return preds, labels + + +@metric_registry('F1', 'tensorflow, tensorflow_itex, pytorch, mxnet, onnxrt_qlinearops, onnxrt_integerops') +class F1(BaseMetric): + """F1 score of a binary classification problem. + + The F1 score is the harmonic mean of the precision and recall. + It can be computed with the equation: + F1 = 2 * (precision * recall) / (precision + recall) + """ + + def __init__(self): + """Initialize the F1 score list.""" + self._score_list = [] + + def update(self, preds, labels): + """Add the predictions and labels. + + Args: + preds: The predictions. + labels: The labels corresponding to the predictions. + """ + from .f1 import f1_score + if getattr(self, '_hvd', None) is not None: + gathered_preds_list = self._hvd.allgather_object(preds) + gathered_labels_list = self._hvd.allgather_object(labels) + temp_preds_list, temp_labels_list = [], [] + for i in range(0, self._hvd.size()): + temp_preds_list += gathered_preds_list[i] + temp_labels_list += gathered_labels_list[i] + preds = temp_preds_list + labels = temp_labels_list + result = f1_score(preds, labels) + self._score_list.append(result) + + def reset(self): + """Clear the predictions and labels.""" + self._score_list = [] + + def result(self): + """Compute the F1 score.""" + return np.array(self._score_list).mean() + +def _accuracy_shape_check(preds, labels): + """Check and conver the shape of predictions and labels. + + Args: + preds: The predictions. + labels: The labels corresponding to the predictions. + + Returns: + preds: The predictions in the format of NumPy array. + labels: The labels in the format of NumPy array. + """ + if isinstance(preds, int): + preds = [preds] + preds = np.array(preds) + if isinstance(labels, int): + labels = [labels] + labels = np.array(labels) + if len(labels.shape) != len(preds.shape) and len(labels.shape)+1 != len(preds.shape): + raise ValueError( + 'labels must have shape of (batch_size, ..) and preds must have' + 'shape of (batch_size, num_classes, ...) or (batch_size, ..),' + 'but given {} and {}.'.format(labels.shape, preds.shape)) + return preds, labels + +def _accuracy_type_check(preds, labels): + """Determine the type of prediction. + + Args: + preds: The predictions. + labels: The labels corresponding to the predictions. + + Returns: + update_type: The type of predictions. + """ + if len(preds.shape) == len(labels.shape)+1: + num_classes = preds.shape[1] + if num_classes == 1: + update_type = 'binary' + else: + update_type = 'multiclass' + elif len(preds.shape) == len(labels.shape): + if len(preds.shape) == 1 or preds.shape[1] ==1: + update_type = 'binary' + else: + update_type = 'multilabel' + return update_type + + +@metric_registry('Accuracy', 'tensorflow, tensorflow_itex, pytorch, onnxrt_qlinearops, onnxrt_integerops') +class Accuracy(BaseMetric): + """The Accuracy for the classification tasks. + + The accuracy score is the proportion of the total number of predictions + that were correct classified. + + Attributes: + pred_list: List of prediction to score. + label_list: List of labels to score. + sample: The total number of samples. + """ + + def __init__(self): + """Initialize predictions, labels and sample.""" + self.pred_list = [] + self.label_list = [] + self.sample = 0 + + def update(self, preds, labels, sample_weight=None): + """Add the predictions and labels. + + Args: + preds: The predictions. + labels: The labels corresponding to the predictions. + sample_weight: The sample weight. + """ + preds, labels = _accuracy_shape_check(preds, labels) + update_type = _accuracy_type_check(preds, labels) + if update_type == 'binary': + self.pred_list.extend(preds) + self.label_list.extend(labels) + self.sample += labels.shape[0] + elif update_type == 'multiclass': + self.pred_list.extend(np.argmax(preds, axis=1).astype('int32')) + self.label_list.extend(labels) + self.sample += labels.shape[0] + elif update_type == 'multilabel': + #(N, C, ...) -> (N*..., C) + num_label = preds.shape[1] + last_dim = len(preds.shape) + if last_dim-1 != 1: + trans_list = [0] + trans_list.extend(list(range(2, len(preds.shape)))) + trans_list.extend([1]) + preds = preds.transpose(trans_list).reshape(-1, num_label) + labels = labels.transpose(trans_list).reshape(-1, num_label) + self.sample += preds.shape[0]*preds.shape[1] + self.pred_list.append(preds) + self.label_list.append(labels) + + def reset(self): + """Clear the predictions and labels.""" + self.pred_list = [] + self.label_list = [] + self.sample = 0 + + def result(self): + """Compute the accuracy.""" + correct_num = np.sum( + np.array(self.pred_list) == np.array(self.label_list)) + if getattr(self, '_hvd', None) is not None: + allghter_correct_num = sum(self._hvd.allgather_object(correct_num)) + allgather_sample = sum(self._hvd.allgather_object(self.sample)) + return allghter_correct_num / allgather_sample + return correct_num / self.sample + + +class PyTorchLoss(): + """A dummy PyTorch Metric. + + A dummy metric that computes the average of predictions and prints it directly. + """ + + def __init__(self): + """Initialize the number of examples, sum of prediction. and device.""" + self._num_examples = 0 + self._device = torch.device('cpu') + self._sum = torch.tensor(0.0, device=self._device) + + def reset(self): + """Reset the number of samples and total cases to zero.""" + self._num_examples = 0 + self._sum = torch.tensor(0.0, device=self._device) + + def update(self, output): + """Add the predictions. + + Args: + output: The predictions. + """ + y_pred, y = output[0].detach(), output[1].detach() + loss = torch.sum(y_pred) + self._sum += loss.to(self._device) + self._num_examples += y.shape[0] + + def compute(self): + """Compute the average of predictions. + + Raises: + ValueError: There must have at least one example. + + Returns: + The dummy loss. + """ + if self._num_examples == 0: + raise ValueError("Loss must have at least one example \ + before it can be computed.") + return self._sum.item() / self._num_examples + + +@metric_registry('Loss', 'tensorflow, tensorflow_itex, pytorch, onnxrt_qlinearops, onnxrt_integerops') +class Loss(BaseMetric): + """A dummy Metric. + + A dummy metric that computes the average of predictions and prints it directly. + + Attributes: + sample: The number of samples. + sum: The sum of prediction. + """ + + def __init__(self): + """Initialize the number of samples, sum of prediction.""" + self.sample = 0 + self.sum = 0 + + def update(self, preds, labels, sample_weight=None): + """Add the predictions and labels. + + Args: + preds: The predictions. + labels: The labels corresponding to the predictions. + sample_weight: The sample weight. + """ + preds, labels = _shape_validate(preds, labels) + self.sample += labels[0].shape[0] + self.sum += sum([np.sum(pred) for pred in preds]) + + def reset(self): + """Reset the number of samples and total cases to zero.""" + self.sample = 0 + self.sum = 0 + + def result(self): + """Compute the average of predictions. + + Returns: + The dummy loss. + """ + if getattr(self, '_hvd', None) is not None: + allgather_sum = sum(self._hvd.allgather_object(self.sum)) + allgather_sample = sum(self._hvd.allgather_object(self.sample)) + return allgather_sum / allgather_sample + return self.sum / self.sample + + +@metric_registry('MAE', 'tensorflow, tensorflow_itex, pytorch, onnxrt_qlinearops, onnxrt_integerops') +class MAE(BaseMetric): + """Computes Mean Absolute Error (MAE) loss. + + Mean Absolute Error (MAE) is the mean of the magnitude of + difference between the predicted and actual numeric values. + + Attributes: + pred_list: List of prediction to score. + label_list: List of references corresponding to the prediction result. + compare_label (bool): Whether to compare label. False if there are no + labels and will use FP32 preds as labels. + """ + + def __init__(self, compare_label=True): + """Initialize the list of prediction and labels. + + Args: + compare_label: Whether to compare label. False if there are no + labels and will use FP32 preds as labels. + """ + self.label_list = [] + self.pred_list = [] + self.compare_label = compare_label + + def update(self, preds, labels, sample_weight=None): + """Add the predictions and labels. + + Args: + preds: The predictions. + labels: The labels corresponding to the predictions. + sample_weight: The sample weight. + """ + preds, labels = _shape_validate(preds, labels) + self.label_list.extend(labels) + self.pred_list.extend(preds) + + def reset(self): + """Clear the predictions and labels.""" + self.label_list = [] + self.pred_list = [] + + def result(self): + """Compute the MAE score. + + Returns: + The MAE score. + """ + aes = [abs(a-b) for (a,b) in zip(self.label_list, self.pred_list)] + aes_sum = sum([np.sum(ae) for ae in aes]) + aes_size = sum([ae.size for ae in aes]) + assert aes_size, "predictions shouldn't be none" + if getattr(self, '_hvd', None) is not None: + aes_sum = sum(self._hvd.allgather_object(aes_sum)) + aes_size = sum(self._hvd.allgather_object(aes_size)) + return aes_sum / aes_size + + +@metric_registry('RMSE', 'tensorflow, tensorflow_itex, pytorch, mxnet, onnxrt_qlinearops, onnxrt_integerops') +class RMSE(BaseMetric): + """Computes Root Mean Squared Error (RMSE) loss. + + Attributes: + mse: The instance of MSE Metric. + + """ + + def __init__(self, compare_label=True): + """Initialize the mse. + + Args: + compare_label (bool): Whether to compare label. False if there are no labels + and will use FP32 preds as labels. + """ + self.mse = MSE(compare_label) + + def update(self, preds, labels, sample_weight=None): + """Add the predictions and labels. + + Args: + preds: The predictions. + labels: The labels corresponding to the predictions. + sample_weight: The sample weight. + """ + self.mse.update(preds, labels, sample_weight) + + def reset(self): + """Clear the predictions and labels.""" + self.mse.reset() + + def result(self): + """Compute the RMSE score. + + Returns: + The RMSE score. + """ + if getattr(self, '_hvd', None) is not None: + self.mse._hvd = self._hvd + return np.sqrt(self.mse.result()) + + + +@metric_registry('MSE', 'tensorflow, tensorflow_itex, pytorch, onnxrt_qlinearops, onnxrt_integerops') +class MSE(BaseMetric): + """Computes Mean Squared Error (MSE) loss. + + Mean Squared Error(MSE) represents the average of the squares of errors. + For example, the average squared difference between the estimated values + and the actual values. + + Attributes: + pred_list: List of prediction to score. + label_list: List of references corresponding to the prediction result. + compare_label (bool): Whether to compare label. False if there are no labels + and will use FP32 preds as labels. + """ + + def __init__(self, compare_label=True): + """Initialize the list of prediction and labels. + + Args: + compare_label: Whether to compare label. False if there are no + labels and will use FP32 preds as labels. + """ + self.label_list = [] + self.pred_list = [] + self.compare_label = compare_label + + def update(self, preds, labels, sample_weight=None): + """Add the predictions and labels. + + Args: + preds: The predictions. + labels: The labels corresponding to the predictions. + sample_weight: The sample weight. + """ + preds, labels = _shape_validate(preds, labels) + self.pred_list.extend(preds) + self.label_list.extend(labels) + + def reset(self): + """Clear the predictions and labels.""" + self.label_list = [] + self.pred_list = [] + + def result(self): + """Compute the MSE score. + + Returns: + The MSE score. + """ + squares = [(a-b)**2.0 for (a,b) in zip(self.label_list, self.pred_list)] + squares_sum = sum([np.sum(square) for square in squares]) + squares_size = sum([square.size for square in squares]) + assert squares_size, "predictions should't be None" + if getattr(self, '_hvd', None) is not None: + squares_sum = sum(self._hvd.allgather_object(squares_sum)) + squares_size = sum(self._hvd.allgather_object(squares_size)) + return squares_sum / squares_size + + +@metric_registry('topk', 'tensorflow, tensorflow_itex') +class TensorflowTopK(BaseMetric): + """Compute Top-k Accuracy classification score for Tensorflow model. + + This metric computes the number of times where the correct label is among + the top k labels predicted. + + Attributes: + k (int): The number of most likely outcomes considered to find the correct label. + num_correct: The number of predictions that were correct classified. + num_sample: The total number of predictions. + """ + + def __init__(self, k=1): + """Initialize the k, number of samples and correct predictions. + + Args: + k: The number of most likely outcomes considered to find the correct label. + """ + self.k = k + self.num_correct = 0 + self.num_sample = 0 + + def update(self, preds, labels, sample_weight=None): + """Add the predictions and labels. + + Args: + preds: The predictions. + labels: The labels corresponding to the predictions. + sample_weight: The sample weight. + """ + preds, labels = _topk_shape_validate(preds, labels) + + labels = labels.reshape([len(labels)]) + with tf.Graph().as_default() as acc_graph: + topk = tf.nn.in_top_k(predictions=tf.constant(preds, dtype=tf.float32), + targets=tf.constant(labels, dtype=tf.int32), k=self.k) + fp32_topk = tf.cast(topk, tf.float32) + correct_tensor = tf.reduce_sum(input_tensor=fp32_topk) + + with tf.compat.v1.Session() as acc_sess: + correct = acc_sess.run(correct_tensor) + + self.num_sample += len(labels) + self.num_correct += correct + + def reset(self): + """Reset the number of samples and correct predictions.""" + self.num_correct = 0 + self.num_sample = 0 + + def result(self): + """Compute the top-k score. + + Returns: + The top-k score. + """ + if self.num_sample == 0: + logger.warning("Sample num during evaluation is 0.") + return 0 + elif getattr(self, '_hvd', None) is not None: + allgather_num_correct = sum(self._hvd.allgather_object(self.num_correct)) + allgather_num_sample = sum(self._hvd.allgather_object(self.num_sample)) + return allgather_num_correct / allgather_num_sample + return self.num_correct / self.num_sample + + +@metric_registry('topk', 'pytorch, mxnet, onnxrt_qlinearops, onnxrt_integerops') +class GeneralTopK(BaseMetric): + """Compute Top-k Accuracy classification score. + + This metric computes the number of times where the correct label is among + the top k labels predicted. + + Attributes: + k (int): The number of most likely outcomes considered to find the correct label. + num_correct: The number of predictions that were correct classified. + num_sample: The total number of predictions. + """ + + def __init__(self, k=1): + """Initialize the k, number of samples and correct predictions. + + Args: + k: The number of most likely outcomes considered to find the correct label. + """ + self.k = k + self.num_correct = 0 + self.num_sample = 0 + + def update(self, preds, labels, sample_weight=None): + """Add the predictions and labels. + + Args: + preds: The predictions. + labels: The labels corresponding to the predictions. + sample_weight: The sample weight. + """ + preds, labels = _topk_shape_validate(preds, labels) + preds = preds.argsort()[..., -self.k:] + if self.k == 1: + correct = accuracy_score(preds, labels, normalize=False) + self.num_correct += correct + + else: + for p, l in zip(preds, labels): + # get top-k labels with np.argpartition + # p = np.argpartition(p, -self.k)[-self.k:] + l = l.astype('int32') + if l in p: + self.num_correct += 1 + + self.num_sample += len(labels) + + def reset(self): + """Reset the number of samples and correct predictions.""" + self.num_correct = 0 + self.num_sample = 0 + + def result(self): + """Compute the top-k score. + + Returns: + The top-k score. + """ + if self.num_sample == 0: + logger.warning("Sample num during evaluation is 0.") + return 0 + elif getattr(self, '_hvd', None) is not None: + allgather_num_correct = sum(self._hvd.allgather_object(self.num_correct)) + allgather_num_sample = sum(self._hvd.allgather_object(self.num_sample)) + return allgather_num_correct / allgather_num_sample + return self.num_correct / self.num_sample + + +@metric_registry('COCOmAPv2', 'tensorflow, tensorflow_itex, onnxrt_qlinearops, onnxrt_integerops') +class COCOmAPv2(BaseMetric): + """Compute mean average precision of the detection task.""" + + def __init__(self, + anno_path=None, + iou_thrs='0.5:0.05:0.95', + map_points=101, + map_key='DetectionBoxes_Precision/mAP', + output_index_mapping={'num_detections':-1, 'boxes':0, 'scores':1, 'classes':2}): + """Initialize the metric. + + Args: + anno_path: The path of annotation file. + iou_thrs: Minimal value for intersection over union that allows to make decision + that prediction bounding box is true positive. You can specify one float value + between 0 to 1 or string "05:0.05:0.95" for standard COCO thresholds. + map_points: The way to calculate mAP. 101 for 101-point interpolated AP, 11 for + 11-point interpolated AP, 0 for area under PR curve. + map_key: The key that mapping to pycocotools COCOeval. + Defaults to 'DetectionBoxes_Precision/mAP'. + output_index_mapping: The output index mapping. + Defaults to {'num_detections':-1, 'boxes':0, 'scores':1, 'classes':2}. + """ + self.output_index_mapping = output_index_mapping + from .coco_label_map import category_map + if anno_path: + import os + import yaml + assert os.path.exists(anno_path), 'Annotation path does not exists!' + with open(anno_path, 'r') as f: + label_map = yaml.safe_load(f.read()) + self.category_map_reverse = {k: v for k,v in label_map.items()} + else: + # label: index + self.category_map_reverse = {v: k for k, v in category_map.items()} + self.image_ids = [] + self.ground_truth_list = [] + self.detection_list = [] + self.annotation_id = 1 + self.category_map = category_map + self.category_id_set = set( + [cat for cat in self.category_map]) #index + self.iou_thrs = iou_thrs + self.map_points = map_points + self.map_key = map_key + + def update(self, predicts, labels, sample_weight=None): + """Add the predictions and labels. + + Args: + predicts: The predictions. + labels: The labels corresponding to the predictions. + sample_weight: The sample weight. Defaults to None. + """ + from .coco_tools import ExportSingleImageGroundtruthToCoco,\ + ExportSingleImageDetectionBoxesToCoco + detections = [] + if 'num_detections' in self.output_index_mapping and \ + self.output_index_mapping['num_detections'] > -1: + for item in zip(*predicts): + detection = {} + num = int(item[self.output_index_mapping['num_detections']]) + detection['boxes'] = np.asarray( + item[self.output_index_mapping['boxes']])[0:num] + detection['scores'] = np.asarray( + item[self.output_index_mapping['scores']])[0:num] + detection['classes'] = np.asarray( + item[self.output_index_mapping['classes']])[0:num] + detections.append(detection) + else: + for item in zip(*predicts): + detection = {} + detection['boxes'] = np.asarray(item[self.output_index_mapping['boxes']]) + detection['scores'] = np.asarray(item[self.output_index_mapping['scores']]) + detection['classes'] = np.asarray(item[self.output_index_mapping['classes']]) + detections.append(detection) + + bboxes, str_labels,int_labels, image_ids = labels + labels = [] + if len(int_labels[0]) == 0: + for str_label in str_labels: + str_label = [ + x if type(x) == 'str' else x.decode('utf-8') + for x in str_label + ] + labels.append([self.category_map_reverse[x] for x in str_label]) + elif len(str_labels[0]) == 0: + for int_label in int_labels: + labels.append([x for x in int_label]) + + for idx, image_id in enumerate(image_ids): + image_id = image_id if type( + image_id) == 'str' else image_id.decode('utf-8') + if image_id in self.image_ids: + continue + self.image_ids.append(image_id) + + ground_truth = {} + ground_truth['boxes'] = np.asarray(bboxes[idx]) + ground_truth['classes'] = np.asarray(labels[idx]) + + self.ground_truth_list.extend( + ExportSingleImageGroundtruthToCoco( + image_id=image_id, + next_annotation_id=self.annotation_id, + category_id_set=self.category_id_set, + groundtruth_boxes=ground_truth['boxes'], + groundtruth_classes=ground_truth['classes'])) + self.annotation_id += ground_truth['boxes'].shape[0] + + self.detection_list.extend( + ExportSingleImageDetectionBoxesToCoco( + image_id=image_id, + category_id_set=self.category_id_set, + detection_boxes=detections[idx]['boxes'], + detection_scores=detections[idx]['scores'], + detection_classes=detections[idx]['classes'])) + + def reset(self): + """Reset the prediction and labels.""" + self.image_ids = [] + self.ground_truth_list = [] + self.detection_list = [] + self.annotation_id = 1 + + def result(self): + """Compute mean average precision. + + Returns: + The mean average precision score. + """ + from .coco_tools import COCOWrapper, COCOEvalWrapper + if len(self.ground_truth_list) == 0: + logger.warning("Sample num during evaluation is 0.") + return 0 + else: + groundtruth_dict = { + 'annotations': + self.ground_truth_list, + 'images': [{ + 'id': image_id + } for image_id in self.image_ids], + 'categories': [{ + 'id': k, + 'name': v + } for k, v in self.category_map.items()] + } + coco_wrapped_groundtruth = COCOWrapper(groundtruth_dict) + coco_wrapped_detections = coco_wrapped_groundtruth.LoadAnnotations( + self.detection_list) + box_evaluator = COCOEvalWrapper(coco_wrapped_groundtruth, + coco_wrapped_detections, + agnostic_mode=False, + iou_thrs = self.iou_thrs, + map_points = self.map_points) + box_metrics, box_per_category_ap = box_evaluator.ComputeMetrics( + include_metrics_per_category=False, all_metrics_per_category=False) + box_metrics.update(box_per_category_ap) + box_metrics = { + 'DetectionBoxes_' + key: value + for key, value in iter(box_metrics.items()) + } + + return box_metrics[self.map_key] + +@metric_registry('mAP', 'tensorflow, tensorflow_itex, onnxrt_qlinearops, onnxrt_integerops') +class TensorflowMAP(BaseMetric): + """Computes mean average precision.""" + + def __init__(self, + anno_path=None, + iou_thrs=0.5, + map_points=0, + map_key='DetectionBoxes_Precision/mAP'): + """Initialize the metric. + + Args: + anno_path: The path of annotation file. + iou_thrs: Minimal value for intersection over union that allows to make decision + that prediction bounding box is true positive. You can specify one float value + between 0 to 1 or string "05:0.05:0.95" for standard COCO thresholds. + map_points: The way to calculate mAP. 101 for 101-point interpolated AP, 11 for + 11-point interpolated AP, 0 for area under PR curve. + map_key: The key that mapping to pycocotools COCOeval. + Defaults to 'DetectionBoxes_Precision/mAP'. + """ + from .coco_label_map import category_map + if anno_path: + import os + import yaml + assert os.path.exists(anno_path), 'Annotation path does not exists!' + with open(anno_path, 'r') as f: + label_map = yaml.safe_load(f.read()) + self.category_map_reverse = {k: v for k,v in label_map.items()} + else: + # label: index + self.category_map_reverse = {v: k for k, v in category_map.items()} + self.image_ids = [] + self.ground_truth_list = [] + self.detection_list = [] + self.annotation_id = 1 + self.category_map = category_map + self.category_id_set = set( + [cat for cat in self.category_map]) #index + self.iou_thrs = iou_thrs + self.map_points = map_points + self.map_key = map_key + + + def update(self, predicts, labels, sample_weight=None): + """Add the predictions and labels. + + Args: + predicts: The predictions. + labels: The labels corresponding to the predictions. + sample_weight: The sample weight. + """ + if getattr(self, '_hvd', None) is not None: + raise NotImplementedError("Metric TensorflowMAP currently do not support distribued inference.") + + from .coco_tools import ExportSingleImageGroundtruthToCoco,\ + ExportSingleImageDetectionBoxesToCoco + detections = [] + if len(predicts) == 3: + for bbox, score, cls in zip(*predicts): + detection = {} + detection['boxes'] = np.asarray(bbox) + detection['scores'] = np.asarray(score) + detection['classes'] = np.asarray(cls) + detections.append(detection) + elif len(predicts) == 4: + for num, bbox, score, cls in zip(*predicts): + detection = {} + num = int(num) + detection['boxes'] = np.asarray(bbox)[0:num] + detection['scores'] = np.asarray(score)[0:num] + detection['classes'] = np.asarray(cls)[0:num] + detections.append(detection) + else: + raise ValueError("Unsupported prediction format!") + + bboxes, str_labels,int_labels, image_ids = labels + labels = [] + if len(int_labels[0]) == 0: + for str_label in str_labels: + str_label = [ + x if type(x) == 'str' else x.decode('utf-8') + for x in str_label + ] + labels.append([self.category_map_reverse[x] for x in str_label]) + elif len(str_labels[0]) == 0: + for int_label in int_labels: + labels.append([x for x in int_label]) + + for idx, image_id in enumerate(image_ids): + image_id = image_id if type( + image_id) == 'str' else image_id.decode('utf-8') + if image_id in self.image_ids: + continue + self.image_ids.append(image_id) + + ground_truth = {} + ground_truth['boxes'] = np.asarray(bboxes[idx]) + ground_truth['classes'] = np.asarray(labels[idx]) + + self.ground_truth_list.extend( + ExportSingleImageGroundtruthToCoco( + image_id=image_id, + next_annotation_id=self.annotation_id, + category_id_set=self.category_id_set, + groundtruth_boxes=ground_truth['boxes'], + groundtruth_classes=ground_truth['classes'])) + self.annotation_id += ground_truth['boxes'].shape[0] + + self.detection_list.extend( + ExportSingleImageDetectionBoxesToCoco( + image_id=image_id, + category_id_set=self.category_id_set, + detection_boxes=detections[idx]['boxes'], + detection_scores=detections[idx]['scores'], + detection_classes=detections[idx]['classes'])) + + def reset(self): + """Reset the prediction and labels.""" + self.image_ids = [] + self.ground_truth_list = [] + self.detection_list = [] + self.annotation_id = 1 + + def result(self): + """Compute mean average precision. + + Returns: + The mean average precision score. + """ + from .coco_tools import COCOWrapper, COCOEvalWrapper + if len(self.ground_truth_list) == 0: + logger.warning("Sample num during evaluation is 0.") + return 0 + else: + groundtruth_dict = { + 'annotations': + self.ground_truth_list, + 'images': [{ + 'id': image_id + } for image_id in self.image_ids], + 'categories': [{ + 'id': k, + 'name': v + } for k, v in self.category_map.items()] + } + coco_wrapped_groundtruth = COCOWrapper(groundtruth_dict) + coco_wrapped_detections = coco_wrapped_groundtruth.LoadAnnotations( + self.detection_list) + box_evaluator = COCOEvalWrapper(coco_wrapped_groundtruth, + coco_wrapped_detections, + agnostic_mode=False, + iou_thrs = self.iou_thrs, + map_points = self.map_points) + box_metrics, box_per_category_ap = box_evaluator.ComputeMetrics( + include_metrics_per_category=False, all_metrics_per_category=False) + box_metrics.update(box_per_category_ap) + box_metrics = { + 'DetectionBoxes_' + key: value + for key, value in iter(box_metrics.items()) + } + + return box_metrics[self.map_key] + +@metric_registry('COCOmAP', 'tensorflow, tensorflow_itex, onnxrt_qlinearops, onnxrt_integerops') +class TensorflowCOCOMAP(TensorflowMAP): + """Computes mean average precision using algorithm in COCO.""" + + def __init__(self, + anno_path=None, + iou_thrs=None, + map_points=None, + map_key='DetectionBoxes_Precision/mAP'): + """Initialize the iou threshold and max points. + + Args: + anno_path: The path of annotation file. + iou_thrs: Minimal value for intersection over union that allows to make decision + that prediction bounding box is true positive. You can specify one float value + between 0 to 1 or string "05:0.05:0.95" for standard COCO thresholds. + map_points: The way to calculate mAP. 101 for 101-point interpolated AP, 11 for + 11-point interpolated AP, 0 for area under PR curve. + map_key: The key that mapping to pycocotools COCOeval. + Defaults to 'DetectionBoxes_Precision/mAP'. + """ + super(TensorflowCOCOMAP, self).__init__(anno_path, iou_thrs, map_points, map_key) + self.iou_thrs = '0.5:0.05:0.95' + self.map_points = 101 + +@metric_registry('VOCmAP', 'tensorflow, tensorflow_itex, onnxrt_qlinearops, onnxrt_integerops') +class TensorflowVOCMAP(TensorflowMAP): + """Computes mean average precision using algorithm in VOC.""" + + def __init__(self, + anno_path=None, + iou_thrs=None, + map_points=None, + map_key='DetectionBoxes_Precision/mAP'): + """Initialize the iou threshold and max points. + + Args: + anno_path: The path of annotation file. + iou_thrs: Minimal value for intersection over union that allows to make decision + that prediction bounding box is true positive. You can specify one float value + between 0 to 1 or string "05:0.05:0.95" for standard COCO thresholds. + map_points: The way to calculate mAP. 101 for 101-point interpolated AP, 11 for + 11-point interpolated AP, 0 for area under PR curve. + map_key: The key that mapping to pycocotools COCOeval. + Defaults to 'DetectionBoxes_Precision/mAP'. + """ + super(TensorflowVOCMAP, self).__init__(anno_path, iou_thrs, map_points, map_key) + self.iou_thrs = 0.5 + self.map_points = 0 + + +@metric_registry('SquadF1', 'tensorflow, tensorflow_itex') +class SquadF1(BaseMetric): + """Evaluate for v1.1 of the SQuAD dataset.""" + + def __init__(self): + """Initialize the score list.""" + self._score_list = [] # squad metric only work when all data preds collected + + def update(self, preds, labels, sample_weight=None): + """Add the predictions and labels. + + Args: + preds: The predictions. + labels: The labels corresponding to the predictions. + sample_weight: The sample weight. + """ + if preds: + from .evaluate_squad import evaluate + if getattr(self, '_hvd', None) is not None: + gathered_preds_list = self._hvd.allgather_object(preds) + gathered_labels_list = self._hvd.allgather_object(labels) + temp_preds_list, temp_labels_list = [], [] + for i in range(0, self._hvd.size()): + temp_preds_list += gathered_preds_list[i] + temp_labels_list += gathered_labels_list[i] + preds = temp_preds_list + labels = temp_labels_list + result = evaluate(labels, preds) + self._score_list.append(result["f1"]) + + def reset(self): + """Reset the score list.""" + self._score_list = [] + + def result(self): + """Compute F1 score.""" + if len(self._score_list) == 0: + return 0. + return np.array(self._score_list).mean() + +@metric_registry('mIOU', 'tensorflow, tensorflow_itex') +class mIOU(BaseMetric): + """Compute the mean IOU(Intersection over Union) score.""" + + def __init__(self, num_classes=21): + """Initialize the number of classes. + + Args: + num_classes: The number of classes. + """ + self.num_classes = num_classes + self.hist = np.zeros((num_classes, num_classes)) + + def update(self, preds, labels): + """Add the predictions and labels. + + Args: + preds: The predictions. + labels: The labels corresponding to the predictions. + """ + preds = preds.flatten() + labels = labels.flatten() + p_dtype = preds.dtype + l_dtype = labels.dtype + if getattr(self, '_hvd', None) is not None: + preds = self._hvd.allgather_object(preds) + labels = self._hvd.allgather_object(labels) + preds_list, labels_list = np.array([], dtype = p_dtype), np.array([], dtype = l_dtype) + for i in range(self._hvd.size()): + preds_list = np.append(preds_list, preds[i]) + labels_list = np.append(labels_list, labels[i]) + preds, labels = preds_list, labels_list + mask = (labels >= 0) & (labels < self.num_classes) + self.hist += np.bincount( + self.num_classes * labels[mask].astype(int) + + preds[mask], minlength=self.num_classes ** 2).reshape(self.num_classes, + self.num_classes) + + def reset(self): + """Reset the hist.""" + self.hist = np.zeros((self.num_classes, self.num_classes)) + + def result(self): + """Compute mean IOU. + + Returns: + The mean IOU score. + """ + iu = np.diag(self.hist) / (self.hist.sum(axis=1) + self.hist.sum(axis=0) - + np.diag(self.hist)) + mean_iu = np.nanmean(iu) + return mean_iu + +@metric_registry('GLUE', 'onnxrt_qlinearops, onnxrt_integerops') +class ONNXRTGLUE(BaseMetric): + """Compute the GLUE score.""" + + def __init__(self, task='mrpc'): + """Initialize the metric. + + Args: + task:The name of the task (Choices: mrpc, qqp, qnli, rte, + sts-b, cola, mnli, wnli.). + """ + assert task in ['mrpc', 'qqp', 'qnli', 'rte', 'sts-b', 'cola', \ + 'mnli', 'wnli', 'sst-2'], 'Unsupported task type' + self.pred_list = None + self.label_list = None + self.task = task + self.return_key = { + "cola": "mcc", + "mrpc": "acc", + "sts-b": "corr", + "qqp": "acc", + "mnli": "mnli/acc", + "qnli": "acc", + "rte": "acc", + "wnli": "acc", + "sst-2": "acc" + } + + def update(self, preds, labels): + """Add the predictions and labels. + + Args: + preds: The predictions. + labels: The labels corresponding to the predictions. + """ + if getattr(self, '_hvd', None) is not None: + raise NotImplementedError("Metric ONNXRTGLUE currently do not support distribued inference.") + if isinstance(preds, list) and len(preds) == 1: + preds = preds[0] + if isinstance(labels, list) and len(labels) == 1: + labels = labels[0] + if self.pred_list is None: + self.pred_list = preds + self.label_list = labels + else: + self.pred_list = np.append(self.pred_list, preds, axis=0) + self.label_list = np.append(self.label_list, labels, axis=0) + + def reset(self): + """Reset the prediction and labels.""" + self.pred_list = None + self.label_list = None + + def result(self): + """Compute the GLUE score.""" + output_mode = transformers.glue_output_modes[self.task] + + if output_mode == "classification": + processed_preds = np.argmax(self.pred_list, axis=1) + elif output_mode == "regression": + processed_preds = np.squeeze(self.pred_list) + result = transformers.glue_compute_metrics(\ + self.task, processed_preds, self.label_list) + return result[self.return_key[self.task]] + +@metric_registry('ROC', 'pytorch') +class ROC(BaseMetric): + """Computes ROC score.""" + + def __init__(self, task='dlrm'): + """Initialize the metric. + + Args: + task:The name of the task (Choices: dlrm, dien, wide_deep.). + """ + assert task in ['dlrm', 'dien', 'wide_deep'], 'Unsupported task type' + self.pred_list = None + self.label_list = None + self.task = task + self.return_key = { + "dlrm": "acc", + "dien": "acc", + "wide_deep": "acc", + } + + def update(self, preds, labels): + """Add the predictions and labels. + + Args: + preds: The predictions. + labels: The labels corresponding to the predictions. + """ + if isinstance(preds, list) and len(preds) == 1: + preds = preds[0] + if isinstance(labels, list) and len(labels) == 1: + labels = labels[0] + if self.pred_list is None: + self.pred_list = preds + self.label_list = labels + else: + self.pred_list = np.append(self.pred_list, preds, axis=0) + self.label_list = np.append(self.label_list, labels, axis=0) + + def reset(self): + """Reset the prediction and labels.""" + self.pred_list = None + self.label_list = None + + def result(self): + """Compute the ROC score.""" + import sklearn.metrics + scores = np.squeeze(self.pred_list) + targets = np.squeeze(self.label_list) + roc_auc = sklearn.metrics.roc_auc_score(targets, scores) + acc = sklearn.metrics.accuracy_score(targets, np.round(scores)) + return acc diff --git a/neural_compressor/model/__init__.py b/neural_compressor/model/__init__.py index e0793e43f96..5f84dec2e4c 100644 --- a/neural_compressor/model/__init__.py +++ b/neural_compressor/model/__init__.py @@ -15,7 +15,9 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .model import MODELS, BaseModel +from .model import MODELS, Model +from .base_model import BaseModel + +__all__ = ["MODELS", "Model", "BaseModel"] -__all__ = ["MODELS", "BaseModel"] diff --git a/neural_compressor/model/model.py b/neural_compressor/model/model.py index 0d2acf6ccea..b723fa57a7c 100644 --- a/neural_compressor/model/model.py +++ b/neural_compressor/model/model.py @@ -17,18 +17,20 @@ import copy import os -import shutil import importlib -from abc import abstractmethod -import tempfile import sys -from neural_compressor.utils.utility import LazyImport, compute_sparsity -from neural_compressor.utils.utility import version1_lt_version2, version1_gt_version2, version1_gte_version2 +from neural_compressor.utils.utility import LazyImport from neural_compressor.utils import logger from neural_compressor.conf import config as cfg from neural_compressor.model.base_model import BaseModel from neural_compressor.model.onnx_model import ONNXModel +from neural_compressor.model.mxnet_model import MXNetModel from neural_compressor.model.keras_model import KerasModel +from neural_compressor.model.tensorflow_model import ( + TensorflowBaseModel, + TensorflowModel, + get_model_type + ) TORCH = False if importlib.util.find_spec('torch'): @@ -44,62 +46,21 @@ json = LazyImport('json') np = LazyImport('numpy') -tensor_to_node = lambda s: list(set([x.split(':')[0] for x in s])) - -def get_model_type(model): - """Get mode type - - Args: - model (string or model object): model path or model object - - Returns: - type (string): model type - """ - - from neural_compressor.adaptor.tf_utils.util import is_saved_model_format, is_ckpt_format - if isinstance(model, str): - model = os.path.abspath(os.path.expanduser(model)) - if (model.endswith('.h5') and os.path.isfile(model)) or \ - is_saved_model_format(os.path.dirname(model)) or \ - (os.path.isdir(model) and is_saved_model_format(model)): - if version1_lt_version2(tf.version.VERSION, '2.10.0'): - logger.warn("keras model running on tensorflow 2.10.0 and" - " lower not support intel ITEX.") - try: - model = tf.keras.models.load_model(model) - except: - pass - if isinstance(model, tf.keras.Model) and hasattr(model, 'to_json'): - return 'keras' - if isinstance(model, tf.Graph): - return 'graph' - elif isinstance(model, tf.compat.v1.GraphDef): - return 'graph_def' - elif isinstance(model, tf.compat.v1.estimator.Estimator): - return 'estimator' - elif isinstance(model, str): - model = os.path.abspath(os.path.expanduser(model)) - if (model.endswith('.pb') and os.path.isfile(model)): - if is_saved_model_format(os.path.dirname(model)): - return 'saved_model' - else: - return 'frozen_pb' - elif model.endswith('.ckpt') and os.path.isfile(model): - return 'slim' - elif os.path.isdir(model): - if is_ckpt_format(model): - return 'checkpoint' - elif is_saved_model_format(model): - return 'saved_model' - elif os.path.isfile(model + '.pb'): - return 'frozen_pb' - - raise ValueError('model {} has not recognized model type....'.format(model)) - +MODELS = {'tensorflow': TensorflowModel, + 'tensorflow_itex': TensorflowModel, + 'keras': KerasModel, + 'mxnet': MXNetModel, + 'pytorch': PyTorchModel if TORCH else None, + 'pytorch_ipex': IPEXModel if TORCH else None, + 'pytorch_fx': PyTorchFXModel if TORCH else None, + 'onnxruntime': ONNXModel, + 'onnxrt_qlinearops': ONNXModel, + 'onnxrt_qdq': ONNXModel, + 'onnxrt_integerops': ONNXModel + } def get_model_fwk_name(model): """Detect the input model belongs to which framework - Args: model (string): framework name that supported by Neural Compressor, if there's no available fwk info, then return 'NA'. @@ -181,996 +142,35 @@ def _is_mxnet(model): return fwk_name -def validate_graph_node(graph_def, node_names): - """Validate nodes exist in the graph_def - - Args: - graph_def (tf.compat.v1.GraphDef): tf.compat.v1.GraphDef object - node_names (list of string): node names to be validated - """ - - if len(node_names) == 0: - return False - all_node_name = [node.name for node in graph_def.node] - for user_name in node_names: - if user_name not in all_node_name: - logger.warn( - str("Node name {} specified in yaml doesn't exist in the model."). - format(user_name)) - return False - return True - -def validate_and_inference_input_output(graph_def, \ - input_tensor_names, output_tensor_names): - """validate and inference the input and output tensor names of graph_def - - Args: - graph_def (tf.compat.v1.GraphDef): tf.compat.v1.GraphDef object - input_tensor_names (list of string): input_tensor_names of graph_def - output_tensor_names (list of string): output_tensor_names of graph_def - - Returns: - input_tensor_names (list of string): validated input_tensor_names - output_tensor_names (list of string): validated output_tensor_names - """ - from neural_compressor.adaptor.tf_utils.util import get_input_output_node_names - temp_output_tensor_names = [] - if validate_graph_node(graph_def, tensor_to_node(input_tensor_names)): - input_tensor_names = input_tensor_names - else: - input_tensor_names, temp_output_tensor_names = get_input_output_node_names(graph_def) - - if validate_graph_node(graph_def, tensor_to_node(output_tensor_names)): - output_tensor_names = output_tensor_names - elif temp_output_tensor_names: - output_tensor_names = temp_output_tensor_names - else: - _, output_tensor_names = get_input_output_node_names(graph_def) - - return input_tensor_names, output_tensor_names - -def graph_session(model, input_tensor_names, output_tensor_names, **kwargs): - """Build session with tf.compat.v1.Graph - - Args: - model (tf.compat.v1.Graph): tf.compat.v1.Graph object - input_tensor_names (list of string): input_tensor_names of model - output_tensor_names (list of string): output_tensor_names of model - - Returns: - sess (tf.compat.v1.Session): tf.compat.v1.Session object - input_tensor_names (list of string): validated input_tensor_names - output_tensor_names (list of string): validated output_tensor_names - """ - - config = tf.compat.v1.ConfigProto() - config.use_per_session_threads = 1 - config.inter_op_parallelism_threads = 1 - sess = tf.compat.v1.Session(graph=model, config=config) - - input_tensor_names, output_tensor_names = validate_and_inference_input_output(\ - model.as_graph_def(), input_tensor_names, output_tensor_names) - - return sess, input_tensor_names, output_tensor_names - -def graph_def_session(model, input_tensor_names, output_tensor_names, **kwargs): - """Build session with tf.compat.v1.GraphDef - - Args: - model (tf.compat.v1.GraphDef): tf.compat.v1.GraphDef object - input_tensor_names (list of string): input_tensor_names of model - output_tensor_names (list of string): output_tensor_names of model - - Returns: - sess (tf.compat.v1.Session): tf.compat.v1.Session object - input_tensor_names (list of string): validated input_tensor_names - output_tensor_names (list of string): validated output_tensor_names - """ - - graph = tf.Graph() - try: - with graph.as_default(): - tf.import_graph_def(model, name='') - except: - input_tensor_names, output_tensor_names = validate_and_inference_input_output(\ - model, input_tensor_names, output_tensor_names) - from neural_compressor.adaptor.tf_utils.util import fix_ref_type_of_graph_def - from neural_compressor.adaptor.tf_utils.util import strip_unused_nodes - model = fix_ref_type_of_graph_def(model) - input_node_names = tensor_to_node(input_tensor_names) - output_node_names = tensor_to_node(output_tensor_names) - model = strip_unused_nodes(model, input_node_names, output_node_names) - with graph.as_default(): - tf.import_graph_def(model, name='') - - return graph_session(graph, input_tensor_names, output_tensor_names, **kwargs) - -def frozen_pb_session(model, input_tensor_names, output_tensor_names, **kwargs): - """Build session with frozen pb - - Args: - model (string): model path - input_tensor_names (list of string): input_tensor_names of model - output_tensor_names (list of string): output_tensor_names of model - - Returns: - sess (tf.compat.v1.Session): tf.compat.v1.Session object - input_tensor_names (list of string): validated input_tensor_names - output_tensor_names (list of string): validated output_tensor_names - """ - - graph_def = tf.compat.v1.GraphDef() - model = model if model.endswith('.pb') else model + '.pb' - with open(model, 'rb') as f: - graph_def.ParseFromString(f.read()) - return graph_def_session(graph_def, input_tensor_names, \ - output_tensor_names, **kwargs) - -def _contains_function_with_implements_attr(saved_model_proto): - meta_graph = saved_model_proto.meta_graphs[0] - for function in meta_graph.graph_def.library.function: - if function.attr.get("_implements", None) or function.attr.get( - "api_implements", None): - return True - return False - -def load_saved_model(model, saved_model_tags, input_tensor_names, output_tensor_names): - """Load graph_def from saved model with the default serving signature key. - - Args: - saved_model_dir: Directory of the SavedModel. - saved_model_tags: Set of tags identifying the MetaGraphDef within the - SavedModel to analyze. - - Returns: - graph_def: The loaded GraphDef. - input_tensors: List of input tensors. - output_tensors: List of output tensors. - """ - config = tf.compat.v1.ConfigProto() - config.use_per_session_threads = 1 - config.inter_op_parallelism_threads = 1 - if not os.listdir(os.path.join(model,'variables')): - sess = tf.compat.v1.Session(graph=tf.Graph(), config=config) - loader = tf.compat.v1.saved_model.loader.load(sess, ["serve"], model) - if len(input_tensor_names) == 0: - input_tensor_names = [i.name for _, i in \ - loader.signature_def['serving_default'].inputs.items()] - else: - assert validate_graph_node(\ - sess.graph.as_graph_def(), tensor_to_node(input_tensor_names)), \ - 'tensor names {} not in the graph'.format(input_tensor_names) - - if len(output_tensor_names) == 0: - output_tensor_names = [i.name for _, i in \ - loader.signature_def['serving_default'].outputs.items()] - else: - assert validate_graph_node(\ - sess.graph.as_graph_def(), tensor_to_node(output_tensor_names)), \ - 'tensor names {} not in the graph'.format(output_tensor_names) - - return sess.graph.as_graph_def(), input_tensor_names, output_tensor_names - else: - from tensorflow.python.eager import context - from tensorflow.python.saved_model import load - from tensorflow.python.saved_model import tag_constants - from tensorflow.python.saved_model import signature_constants - from tensorflow.python.framework.convert_to_constants import \ - convert_variables_to_constants_v2 - from tensorflow.python.training import saver - from tensorflow.core.protobuf import config_pb2 - from tensorflow.python.grappler import tf_optimizer - from tensorflow.core.protobuf import meta_graph_pb2 - _saved_model = load.load(model, [tag_constants.SERVING]) - func = _saved_model.signatures[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] - frozen_func = convert_variables_to_constants_v2(func) - grappler_meta_graph_def = saver.export_meta_graph( - graph_def=frozen_func.graph.as_graph_def(), graph=frozen_func.graph) - if len(input_tensor_names) == 0: - input_tensor_names = [i.name.split(':')[0] for i in frozen_func.inputs] - if len(output_tensor_names) == 0: - output_tensor_names = [i.name.split(':')[0] for i in frozen_func.outputs] - # Add a collection 'train_op' so that Grappler knows the outputs. - fetch_collection = meta_graph_pb2.CollectionDef() - for array in frozen_func.inputs + frozen_func.outputs: - fetch_collection.node_list.value.append(array.name) - grappler_meta_graph_def.collection_def["train_op"].CopyFrom( - fetch_collection) - from tensorflow.python.eager import context - grappler_session_config = config_pb2.ConfigProto() - rewrite_options = grappler_session_config.graph_options.rewrite_options - rewrite_options.min_graph_nodes = -1 - opt = tf_optimizer.OptimizeGraph(grappler_session_config, - grappler_meta_graph_def, graph_id=b"tf_graph") - return opt, input_tensor_names, output_tensor_names - -def get_graph_from_saved_model_v2(saved_model_dir, - input_tensor_names, output_tensor_names): - from tensorflow.python.saved_model import tag_constants - from tensorflow.python.saved_model import signature_constants - saved_model_exported_names = [ - signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY - ] - saved_model_tags = set([tag_constants.SERVING]) - return load_saved_model(saved_model_dir, saved_model_tags, - input_tensor_names, output_tensor_names) - -def get_graph_from_original_keras_v2(model, output_dir): - from tensorflow.python.eager import def_function - from tensorflow.lite.python.util import trace_model_call - from tensorflow.lite.python.util import model_input_signature - from tensorflow.python.framework import convert_to_constants - from tensorflow.python.framework import dtypes - from tensorflow.lite.python.util import run_graph_optimizations - from tensorflow.lite.python.convert import OpsSet - from tensorflow.lite.python.util import get_grappler_config - input_signature = None - # If the model's call is not a `tf.function`, then we need to first get its - # input signature from `model_input_signature` method. - if not isinstance(model.call, def_function.Function): - input_signature = model_input_signature(model, keep_original_batch_size=False) - - func = trace_model_call(model, input_signature) - concrete_func = func.get_concrete_function() - funcs = [concrete_func] - - frozen_func, graph_def = ( - convert_to_constants.convert_variables_to_constants_v2_as_graph( - funcs[0], lower_control_flow=False)) - - input_tensors = [ - tensor for tensor in frozen_func.inputs - if tensor.dtype != dtypes.resource - ] - output_tensors = frozen_func.outputs - # Grappler will also try to lower while loop into switch merge - # representation which is undesired for Ophints, so we simply remove - # those attributes to prevent Grappler from doing so. - graph = convert_to_constants.disable_lower_using_switch_merge(graph_def) - # Run function inlining optimization to ensure any models generated - # through the from_frozen_graph path have been inlined. - # grappler_config = get_grappler_config(['function']) - # graph_def = run_graph_optimizations( - # graph, - # input_tensors, - # output_tensors, - # config=grappler_config) - input_names = [tensor.name.split(':')[0] for tensor in input_tensors] - output_names = [tensor.name.split(':')[0] for tensor in output_tensors] - return graph_def, input_names, output_names - -def check_keras_format(model, saved_model_dir): - from tensorflow.python import saved_model - from tensorflow.python.saved_model.load import load - from tensorflow.python.saved_model import save_options - from tensorflow.python.saved_model.loader_impl import parse_saved_model_with_debug_info - version = 'saved_model_v2' - try: - saved_model.save( - model, - saved_model_dir, - options=save_options.SaveOptions(save_debug_info=True)) - except: - return 'trackable_object' - saved_model_proto, _ = parse_saved_model_with_debug_info(saved_model_dir) - saved_model_version = saved_model_proto.saved_model_schema_version - if saved_model_version == 0: - return 'saved_model_v1' - if saved_model_version not in [1, 2]: - raise ValueError("SavedModel file format({0}) is not supported".format( - saved_model_version)) - return version - -def get_graph_from_saved_model_v1(model): - from tensorflow.python.framework import ops - from tensorflow.python.saved_model import constants - from tensorflow.python.client import session - from tensorflow.python.saved_model import tag_constants - from tensorflow.python.saved_model import signature_constants - from tensorflow.lite.python.convert_saved_model import get_meta_graph_def - from tensorflow.lite.python.convert_saved_model import get_signature_def - from tensorflow.lite.python.convert_saved_model import get_inputs_outputs - saved_model_tags = set([tag_constants.SERVING]) - signature_key = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY - - meta_graph = get_meta_graph_def(model, saved_model_tags) - signature_def = get_signature_def(meta_graph, signature_key) - inputs, outputs = get_inputs_outputs(signature_def) - # Check SavedModel for assets directory. - collection_def = meta_graph.collection_def - if constants.ASSETS_KEY in collection_def: - raise ValueError("SavedModels with assets/ directory are not supported.") - - from tensorflow.python.saved_model import loader - from tensorflow.python.framework import graph_util as tf_graph_util - graph = ops.Graph() - import tensorflow as tf - with session.Session(graph=graph) as sess: - loader.load(sess, meta_graph.meta_info_def.tags, model) - sess.run(tf.compat.v1.global_variables_initializer()) - sess.run(tf.compat.v1.tables_initializer()) - output_nodes = list(set([output.split(':')[0] for output in outputs])) - node_ops = [node.op for node in graph.as_graph_def().node] - if 'MakeIterator' in node_ops: - output_nodes.append('MakeIterator') - table_ops = tf.compat.v1.get_collection( - tf.compat.v1.GraphKeys.TABLE_INITIALIZERS) - # For table initialization - for table_op in table_ops: - output_nodes.append(table_op.name) - if len(table_ops) > 0: - output_nodes.append('init_all_tables') - graph_def = tf_graph_util.convert_variables_to_constants( - sess, graph.as_graph_def(), output_nodes) - return graph_def, inputs, outputs - -def keras_session(model, input_tensor_names, output_tensor_names, **kwargs): - """Build session with keras model - Args: - model (string or tf.keras.Model): model path or tf.keras.Model object - input_tensor_names (list of string): input_tensor_names of model - output_tensor_names (list of string): output_tensor_names of model - Returns: - sess (tf.compat.v1.Session): tf.compat.v1.Session object - input_tensor_names (list of string): validated input_tensor_names - output_tensor_names (list of string): validated output_tensor_names - """ - temp_dir = tempfile.mkdtemp() - if tf.version.VERSION > '2.1.0': - if not isinstance(model, tf.keras.Model): - model = tf.keras.models.load_model(model) - keras_format = check_keras_format(model, temp_dir) - if keras_format == 'saved_model_v2': - try: - graph_def, input_names, output_names = get_graph_from_saved_model_v2( - temp_dir, input_tensor_names, output_tensor_names) - if '_FusedBatchNormEx' in [node.op for node in graph_def.node]: - keras_format = 'trackable_object' - except: - keras_format = 'trackable_object' - if keras_format == 'trackable_object': - try: - graph_def, input_names, output_names = get_graph_from_original_keras_v2( - model, temp_dir) - except: - keras_format = 'saved_model_v1' - if keras_format == 'saved_model_v1': - try: - tf.keras.backend.set_learning_phase(0) - graph_def, input_names, output_names = get_graph_from_saved_model_v1(model) - except: - raise ValueError('Not supported keras model type...') - - # tensorflow 1.x use v1 convert method - else: - tf.keras.backend.set_learning_phase(0) - graph_def, input_names, output_names = get_graph_from_saved_model_v1(model) - shutil.rmtree(temp_dir, True) - return graph_def_session(graph_def, input_names, output_names, **kwargs) - - -def slim_session(model, input_tensor_names, output_tensor_names, **kwargs): - """Build session with slim model - - Args: - model (string): model path - input_tensor_names (list of string): input_tensor_names of model - output_tensor_names (list of string): output_tensor_names of model - - Returns: - sess (tf.compat.v1.Session): tf.compat.v1.Session object - input_tensor_names (list of string): validated input_tensor_names - output_tensor_names (list of string): validated output_tensor_names - """ - - assert version1_lt_version2(tf.version.VERSION, '2.0.0'), 'slim model only used in tensorflow 1.x' - from .nets_factory import TFSlimNetsFactory - factory = TFSlimNetsFactory() - assert 'name' in kwargs, 'model name should be set in slim checkpoint....' - assert kwargs['name'] in factory.default_slim_models, \ - 'only support topology {}'.format(factory.default_slim_models) - net = copy.deepcopy(factory.networks_map[kwargs['name']]) - model_func = net.pop('model') - arg_scope = net.pop('arg_scope')() - inputs_shape = net.pop('input_shape') - kwargs = net - import tf_slim as slim - with tf.Graph().as_default(): - images = tf.compat.v1.placeholder(name='input', dtype=tf.float32, \ - shape=inputs_shape) - with tf.compat.v1.Session() as sess: - with slim.arg_scope(arg_scope) as scope: # pylint: disable=not-context-manager - model_func(images, is_training=False, **kwargs) - graph_def = sess.graph.as_graph_def() - output_tensor_names = output_tensor_names if len(output_tensor_names) > 0 \ - else [graph_def.node[-1].name] - - from tensorflow.python.tools.freeze_graph import freeze_graph_with_def_protos - graph_def = freeze_graph_with_def_protos( - input_graph_def=graph_def, - input_saver_def=None, - input_checkpoint=model, - output_node_names=','.join(output_tensor_names), - restore_op_name='save/restore_all', - filename_tensor_name='save/Const:0', - output_graph='', - clear_devices=True, - initializer_nodes='') - - return graph_def_session(graph_def, ['input'], output_tensor_names) - -def checkpoint_session(model, input_tensor_names, output_tensor_names, **kwargs): - """Build session with ckpt model - - Args: - model (string): model path - input_tensor_names (list of string): input_tensor_names of model - output_tensor_names (list of string): validated output_tensor_names of model - - Returns: - sess (tf.compat.v1.Session): tf.compat.v1.Session object - input_tensor_names (list of string): validated input_tensor_names - output_tensor_names (list of string): validated output_tensor_names - """ - - assert output_tensor_names is not None and len(output_tensor_names) > 0, \ - 'outputs should not be None of checkpoint....' - - ckpt_prefix = [os.path.splitext(i)[0] for i in os.listdir(model) \ - if i.endswith(".meta")][0] - - config = tf.compat.v1.ConfigProto() - config.use_per_session_threads = 1 - config.inter_op_parallelism_threads = 1 - graph = tf.Graph() - sess = tf.compat.v1.Session(graph=graph, config=config) - with graph.as_default(): - saver = tf.compat.v1.train.import_meta_graph(\ - os.path.join(model, ckpt_prefix + '.meta'), clear_devices=True) - - sess.run(tf.compat.v1.global_variables_initializer()) - saver.restore(sess, os.path.join(model, ckpt_prefix)) - - from neural_compressor.adaptor.tf_utils.util import get_input_output_node_names - if validate_graph_node(sess.graph.as_graph_def(), tensor_to_node(input_tensor_names)): - input_tensor_names = input_tensor_names - else: - input_tensor_names, _ = get_input_output_node_names(sess.graph.as_graph_def()) - return sess, input_tensor_names, output_tensor_names - -def estimator_session(model, input_tensor_names, output_tensor_names, **kwargs): - """Build session with estimator model - - Args: - model (tf.estimator.Estimator): tf.estimator.Estimator object - input_tensor_names (list of string): input_tensor_names of model - output_tensor_names (list of string): output_tensor_names of model - kwargs (dict): other required parameters, like input_fn - - Returns: - sess (tf.compat.v1.Session): tf.compat.v1.Session object - input_tensor_names (list of string): validated input_tensor_names - output_tensor_names (list of string): validated output_tensor_names - """ - - assert 'input_fn' in kwargs, 'input func should be supplied for estimator session....' - with tf.Graph().as_default() as g: - features, input_hooks = model._get_features_from_input_fn( - kwargs['input_fn'], tf.estimator.ModeKeys.PREDICT) - estimator_spec = model._call_model_fn(features, None, - tf.estimator.ModeKeys.PREDICT, model.config) - - if len(output_tensor_names) == 0: - outputs = [tensor.name for tensor in estimator_spec.predictions.values()] if\ - isinstance(estimator_spec.predictions, dict) else \ - [estimator_spec.predictions.name] - else: - outputs = output_tensor_names - - logger.info("Estimator output tensor names are {}.".format(outputs)) - with tf.compat.v1.Session(graph=g) as sess: - sess.run(tf.compat.v1.global_variables_initializer()) - # Freezing a graph requires output_node_names, which can be found in - # estimator_spec.predictions that contains prediction tensors as a - # dictionary - # When a model uses Iterator, we need to have 'MakeIterator' (default - # name used by TF) in the output_node_names as well. - output_nodes = list(set([output.split(':')[0] for output in outputs])) - if 'MakeIterator' in [node.op for node in g.as_graph_def().node]: - output_nodes.append('MakeIterator') - - graph_def = tf.compat.v1.graph_util.convert_variables_to_constants(sess, - g.as_graph_def(), output_nodes) - - return graph_def_session(graph_def, input_tensor_names, outputs) - -def saved_model_session(model, input_tensor_names, output_tensor_names, **kwargs): - """Build session with saved model - - Args: - model (string): model path - input_tensor_names (list of string): input_tensor_names of model - output_tensor_names (list of string): output_tensor_names of model - - Returns: - sess (tf.compat.v1.Session): tf.compat.v1.Session object - input_tensor_names (list of string): validated input_tensor_names - output_tensor_names (list of string): validated output_tensor_names - """ - try: - graph_def, input_names, output_names = get_graph_from_saved_model_v2( - model, input_tensor_names, output_tensor_names) - except: - graph_def, input_names, output_names = get_graph_from_saved_model_v1(model) - assert graph_def is not None, 'Can not parse the saved model...' - return graph_def_session(graph_def, input_names, output_names, **kwargs) - -# it's necessary that a session with input output tensors to run the model -SESSIONS = {'frozen_pb': frozen_pb_session, - 'graph_def': graph_def_session, - 'graph': graph_session, - 'saved_model': saved_model_session, - 'keras': keras_session, - 'checkpoint': checkpoint_session, - 'estimator': estimator_session, - 'slim': slim_session,} - - -class TensorflowBaseModel(BaseModel): - """Build TensorflowBaseModel object - - Args: - model (string or tensorflow model object): model path or model object - kwargs (dict): other required parameters, like input_fn - - """ - - def __init__(self, model, **kwargs): - - self._model = model - self._name = '' - self._weights = None - self.kwargs = kwargs - self._graph_info = {} - self._input_tensor_names = [] - self._output_tensor_names = [] - self._model_type = '' - self._sess = None - self._iter_op = None - self._workspace_path = '' - self._q_config = None - - def framework(self): - return 'tensorflow' - - @property - def name(self): - return self._name - - @name.setter - def name(self, name): - self.kwargs.update({'name': name}) - self._name = name - - @property - def weights(self): - """ Getter to weights """ - return self._weights - - @weights.setter - def weights(self, new_weights): - """ Setter to weights """ - self._weights = new_weights - - @property - def q_config(self): - return self._q_config - - @q_config.setter - def q_config(self, q_config): - self._q_config = q_config - - @property - def workspace_path(self): - return self._workspace_path - - @workspace_path.setter - def workspace_path(self, path): - self._workspace_path = path - - @property - def model_type(self): - return self._model_type - - @model_type.setter - def model_type(self, model_type): - assert model_type in SESSIONS, 'model type not supported....' - self._model_type = model_type - - @property - def model(self): - return self.graph - - @property - def graph_def(self): - return self.graph.as_graph_def() - - @property - def graph_info(self): - self._graph_info = {} - for node in self.graph_def.node: - self._graph_info[node.name] = node.op - return self._graph_info - - @property - def sess(self): - if self._sess is None: - self._load_sess(self._model, **self.kwargs) - return self._sess - - @property - def graph(self): - return self.sess.graph - - @graph_def.setter - def graph_def(self, graph_def): - if self._sess is not None: - self._sess.close() - output_sess = SESSIONS['graph_def'](graph_def,\ - self._input_tensor_names, \ - self._output_tensor_names) - - self._sess = output_sess[0] - self._input_tensor_names = output_sess[1] - self._output_tensor_names = output_sess[2] - self.model_type = 'graph_def' - - def _load_sess(self, model, **kwargs): - if self.name: - kwargs.update({'name': self.name}) - # assert self.model_type, 'model type not set....' - output_sess = SESSIONS[self.model_type](model, - self._input_tensor_names, \ - self._output_tensor_names, - **kwargs) - self._sess = output_sess[0] - self._input_tensor_names = output_sess[1] - self._output_tensor_names = output_sess[2] - - tf.compat.v1.get_variable_scope().reuse_variables() - return self._sess - - @property - def iter_op(self): - self._iter_op = [] - if self._sess is None: - self._load_sess(self._model, **self.kwargs) - op_list = [node.op for node in self._sess.graph.as_graph_def().node] - if 'MakeIterator' in op_list: - self._iter_op.append(self._sess.graph.get_operation_by_name('MakeIterator')) - return self._iter_op - - @property - def input_tensor_names(self): - if self._sess is None: - self._load_sess(self._model, **self.kwargs) - return copy.deepcopy(self._input_tensor_names) - - @input_tensor_names.setter - def input_tensor_names(self, tensor_names): - if len(tensor_names) == 0: - logger.warn("Input tensor names is empty.") - return - if self._sess is not None: - assert validate_graph_node(\ - self.graph_def, tensor_to_node(tensor_names)), \ - 'tensor names {} not in graph'.format(tensor_names) - self._input_tensor_names = tensor_names - - @property - def output_tensor_names(self): - if len(self._output_tensor_names) == 0: - self._load_sess(self._model, **self.kwargs) - return copy.deepcopy(self._output_tensor_names) - - @output_tensor_names.setter - def output_tensor_names(self, tensor_names): - if len(tensor_names) == 0: - logger.warn("Output tensor names should not be empty.") - return - if self._sess is not None: - assert validate_graph_node(\ - self.graph_def, tensor_to_node(tensor_names)), \ - 'tensor names {} not in graph'.format(tensor_names) - self._output_tensor_names = tensor_names - - # input/output node names and input/output tensor - # come from input/output tensor names, so do not support assign these values - @property - def input_node_names(self): - return copy.deepcopy(tensor_to_node(self.input_tensor_names)) - - @property - def output_node_names(self): - output_node_names = tensor_to_node(self.output_tensor_names) - iter_op_list = self.iter_op - if iter_op_list != []: - output_node_names += [iter_op.name for iter_op in iter_op_list] - return copy.deepcopy(output_node_names) - - @property - def input_tensor(self): - from neural_compressor.adaptor.tf_utils.util import get_tensor_by_name - return [get_tensor_by_name(\ - self.graph, x) for x in self.input_tensor_names] - - @property - def output_tensor(self): - from neural_compressor.adaptor.tf_utils.util import get_tensor_by_name - return [get_tensor_by_name(\ - self.graph, x) for x in self.output_tensor_names] - - def save(self, root=None): - if not root: - root = cfg.default_workspace + '/save.pb' - root = os.path.abspath(os.path.expanduser(root)) - # if not have suffix, default append .pb - os.makedirs(os.path.dirname(root), exist_ok=True) - pb_file = root if os.path.split(root)[-1].endswith('.pb') else root + '.pb' - f = tf.io.gfile.GFile(pb_file, 'wb') - f.write(self.graph_def.SerializeToString()) - logger.info("Save quantized model to {}.".format(pb_file)) - - -class TensorflowSavedModelModel(TensorflowBaseModel): - def get_all_weight_names(self): - import tensorflow as tf - names = [] - for index, layer in enumerate(tf.keras.models.load_model(self._model).layers): - if len(layer.weights): - names.append(index) - return names - - def update_weights(self, tensor_name, new_tensor): - pass - - def get_weight(self, tensor_name): - return self.weights[tensor_name] - - @property - def model(self): - import time - import shutil - root = os.path.abspath(os.path.expanduser(cfg.default_workspace)) - root += str(time.time()) - if os.path.exists(root): - shutil.rmtree(root) - os.makedirs(root, exist_ok=True) - if not self._sess: - self._load_sess(self._model, **self.kwargs) - _, builder = self.build_saved_model(root) - builder.save() - model = tf.saved_model.load(root) - shutil.rmtree(root) - return model - - def report_sparsity(self): - """ Get sparsity of the model - +class Model(object): + """A wrapper of the information needed to construct a Model.""" + + def __new__(cls, root, **kwargs): + """Create a new instance object of Model. Args: - + root (object): raw model format. For Tensorflow model, could be path to frozen pb file, + path to ckpt or savedmodel folder, loaded estimator/graph_def/graph/keras model object. + For PyTorch model, it's torch.nn.model instance. For MXNet model, it's mxnet.symbol.Symbol + or gluon.HybirdBlock instance. For ONNX model, it's path to onnx model or loaded ModelProto + model object. + Returns: - df (DataFrame): DataFrame of sparsity of each weight - total_sparsity (float): total sparsity of model - + BaseModel: neural_compressor built-in model """ - import pandas as pd - import tensorflow as tf - import numpy as np - df = pd.DataFrame(columns=['Name', 'Shape', 'NNZ (dense)', 'NNZ (sparse)', "Sparsity(%)", - 'Std', 'Mean', 'Abs-Mean']) - pd.set_option('display.precision', 2) - param_dims = [2, 4] - params_size = 0 - sparse_params_size = 0 - for index, layer in enumerate(tf.keras.models.load_model(self._model).layers): - if not len(layer.weights): - continue - # Extract just the actual parameter's name, which in this context we treat - # as its "type" - weights = layer.get_weights()[0] - if weights.ndim in param_dims: - param_size, sparse_param_size, dense_param_size = compute_sparsity( - weights) - density = dense_param_size / param_size - params_size += param_size - sparse_params_size += sparse_param_size - df.loc[len(df.index)] = ([ - index, - list(weights.shape), - dense_param_size, - sparse_param_size, - (1 - density) * 100, - np.std(weights), - np.mean(weights), - np.mean(np.abs(weights)) - ]) - - total_sparsity = sparse_params_size / params_size * 100 - - df.loc[len(df.index)] = ([ - 'Total sparsity:', - params_size, - "-", - int(sparse_params_size), - total_sparsity, - 0, 0, 0]) - - return df, total_sparsity - - def build_saved_model(self, root=None): - if not root: - root = cfg.default_workspace - root = os.path.abspath(os.path.expanduser(root)) - if os.path.exists(root): - import shutil - shutil.rmtree(root) - - os.makedirs(root, exist_ok=True) - - from tensorflow.python.saved_model import signature_constants - from tensorflow.python.saved_model import tag_constants - from neural_compressor.adaptor.tf_utils.util import get_tensor_by_name - builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(root) - sigs = {} - with tf.compat.v1.Session(graph=tf.Graph()) as sess: - #(TODO) not directly use self._sess.graph, use self.graph - tf.import_graph_def(self.graph.as_graph_def(), name="") - g = tf.compat.v1.get_default_graph() - inp = [get_tensor_by_name(g, x) for x in self._input_tensor_names] - out = [get_tensor_by_name(g, x) for x in self._output_tensor_names] - sigs[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = \ - tf.compat.v1.saved_model.signature_def_utils.predict_signature_def( - {k: v for k, v in zip(self._input_tensor_names, inp)}, - {k: v for k, v in zip(self._output_tensor_names, out)}) - builder.add_meta_graph_and_variables(sess, - [tag_constants.SERVING], - signature_def_map=sigs) - return root, builder - - def save(self, root=None): - root, builder = self.build_saved_model(root) - builder.save() - logger.info("Save quantized model to {}.".format(root)) - -class TensorflowQATModel(TensorflowSavedModelModel): - def __init__(self, model='', **kwargs): - super(TensorflowQATModel, self).__init__(model) - self.keras_model = None - self.model_type = 'keras' - - @property - def model(self): - if self.keras_model == None: - self.keras_model = tf.keras.models.load_model(self._model) - return self.keras_model - - @model.setter - def model(self, q_model): - self.keras_model = q_model - - def save(self, root=None): - if not root: - root = cfg.default_workspace + '/saved_model' - root = os.path.abspath(os.path.expanduser(root)) - # if not have suffix, default append .pb - os.makedirs(os.path.dirname(root), exist_ok=True) - q_aware_model = self.keras_model - q_aware_model.save(root) - saved_format = 'saved_model' - if root.endswith('.h5'): - saved_format = 'h5 file' - logger.info("Save quantized model to {}.".format(saved_format)) - return root - -class TensorflowCheckpointModel(TensorflowBaseModel): - - @property - def graph_def(self): - if self.model_type == 'graph_def': - return self.sess.graph.as_graph_def() - from neural_compressor.adaptor.tf_utils.util import _parse_ckpt_bn_input - from tensorflow.python.framework import graph_util - graph_def = self.sess.graph.as_graph_def() - graph_def = _parse_ckpt_bn_input(graph_def) - return graph_util.convert_variables_to_constants( - sess=self._sess, - input_graph_def=graph_def, - output_node_names=self.output_node_names) - - @graph_def.setter - def graph_def(self, graph_def): - if self._sess is not None: - self._sess.close() - output_sess = SESSIONS['graph_def'](graph_def, - self._input_tensor_names, \ - self._output_tensor_names) - self._sess = output_sess[0] - self._input_tensor_names = output_sess[1] - self._output_tensor_names = output_sess[2] - self.model_type = 'graph_def' - - -TENSORFLOW_MODELS = {'frozen_pb': TensorflowBaseModel, - 'graph_def': TensorflowBaseModel, - 'graph': TensorflowBaseModel, - 'checkpoint': TensorflowCheckpointModel, - 'estimator': TensorflowBaseModel, - 'slim': TensorflowBaseModel, - 'saved_model': TensorflowSavedModelModel, - 'keras': TensorflowSavedModelModel - } + framework = kwargs.get("framework", "NA") + if framework == "NA": + framework = get_model_fwk_name(root) -class TensorflowModel(object): - def __new__(cls, model_type, root, **kwargs): - os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" - os.environ["CUDA_VISIBLE_DEVICES"] = "-1" - model = TENSORFLOW_MODELS[model_type](root, **kwargs) - model.model_type = model_type - return model - - -class MXNetModel(BaseModel): - """Build MXNetModel object - - Args: - model (mxnet model): model path - """ - - def __init__(self, model, **kwargs): - #(TODO) MXNet does not support recover model from tuning history currently - self.q_config = None - self._model = model - self.calib_cache = {} - - def framework(self): - return 'mxnet' - - @property - def model(self): - return self._model - - @model.setter - def model(self, model): - self._model = model - - def save(self, root=None): - if root is None: - root = cfg.default_workspace - root = os.path.abspath(os.path.expanduser(root)) - os.makedirs(os.path.dirname(root), exist_ok=True) - - if isinstance(self._model, mx.gluon.HybridBlock): - self._model.export(root, remove_amp_cast=False) - logger.info("Save quantized hybrid block model to {}.".format(root)) + if 'tensorflow' in framework: + if 'modelType' in kwargs: + model_type = kwargs['modelType'] + else: + model_type = get_model_type(root) + model = MODELS['tensorflow'](model_type, root, **kwargs) + elif framework == 'keras': + model = MODELS['keras'](root, **kwargs) + elif framework == 'pytorch': + model = MODELS[framework](root, **kwargs) else: - symnet, args, auxs = self._model - symnet = symnet.as_nd_ndarray() - args = {k:v.as_nd_ndarray() for k, v in args.items()} - auxs = {k:v.as_nd_ndarray() for k, v in auxs.items()} - mx.model.save_checkpoint(root, 0, symnet, args, auxs, remove_amp_cast=False) - logger.info("Save quantized symbol model to {}.".format(root)) - - -MODELS = {'tensorflow': TensorflowModel, - 'tensorflow_itex': TensorflowModel, - 'keras': KerasModel, - 'mxnet': MXNetModel, - 'pytorch': PyTorchModel if TORCH else None, - 'pytorch_ipex': PyTorchIpexModel if TORCH else None, - 'pytorch_fx': PyTorchFXModel if TORCH else None, - 'onnxruntime': ONNXModel, - 'onnxrt_qlinearops': ONNXModel, - 'onnxrt_qdq': ONNXModel, - 'onnxrt_integerops': ONNXModel - } + model = MODELS[framework](root, **kwargs) + return model \ No newline at end of file diff --git a/neural_compressor/model/mxnet_model.py b/neural_compressor/model/mxnet_model.py new file mode 100644 index 00000000000..84188b8efc9 --- /dev/null +++ b/neural_compressor/model/mxnet_model.py @@ -0,0 +1,64 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +from neural_compressor.conf import config as cfg +from neural_compressor.utils.utility import LazyImport +from neural_compressor.utils import logger +from .base_model import BaseModel +mx = LazyImport('mxnet') + +class MXNetModel(BaseModel): + """Build MXNetModel object + + Args: + model (mxnet model): model path + """ + + def __init__(self, model, **kwargs): + #(TODO) MXNet does not support recover model from tuning history currently + self.q_config = None + self._model = model + self.calib_cache = {} + + def framework(self): + return 'mxnet' + + @property + def model(self): + return self._model + + @model.setter + def model(self, model): + self._model = model + + def save(self, root=None): + if root is None: + root = cfg.default_workspace + root = os.path.abspath(os.path.expanduser(root)) + os.makedirs(os.path.dirname(root), exist_ok=True) + + if isinstance(self._model, mx.gluon.HybridBlock): + self._model.export(root, remove_amp_cast=False) + logger.info("Save quantized hybrid block model to {}.".format(root)) + else: + symnet, args, auxs = self._model + symnet = symnet.as_nd_ndarray() + args = {k:v.as_nd_ndarray() for k, v in args.items()} + auxs = {k:v.as_nd_ndarray() for k, v in auxs.items()} + mx.model.save_checkpoint(root, 0, symnet, args, auxs, remove_amp_cast=False) + logger.info("Save quantized symbol model to {}.".format(root)) diff --git a/neural_compressor/model/tensorflow_model.py b/neural_compressor/model/tensorflow_model.py new file mode 100644 index 00000000000..8070935a337 --- /dev/null +++ b/neural_compressor/model/tensorflow_model.py @@ -0,0 +1,998 @@ +#!/usr/bin/env python +# -*- coding: utf-8 -*- +# +# Copyright (c) 2021 Intel Corporation +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import os +import shutil +import importlib +from abc import abstractmethod +import tempfile +import sys +from neural_compressor.utils.utility import LazyImport, compute_sparsity +from neural_compressor.utils.utility import version1_lt_version2, version1_gt_version2, version1_gte_version2 +from neural_compressor.utils import logger +from neural_compressor.conf import config as cfg +from neural_compressor.model.base_model import BaseModel + +tf = LazyImport('tensorflow') +np = LazyImport('numpy') + +tensor_to_node = lambda s: list(set([x.split(':')[0] for x in s])) + +def get_model_type(model): + """Get mode type + Args: + model (string or model object): model path or model object + Returns: + type (string): model type + """ + + from neural_compressor.adaptor.tf_utils.util import is_saved_model_format, is_ckpt_format + if isinstance(model, str): + model = os.path.abspath(os.path.expanduser(model)) + if (model.endswith('.h5') and os.path.isfile(model)) or \ + is_saved_model_format(os.path.dirname(model)) or \ + (os.path.isdir(model) and is_saved_model_format(model)): + if version1_lt_version2(tf.version.VERSION, '2.10.0'): + logger.warn("keras model running on tensorflow 2.10.0 and" + " lower not support intel ITEX.") + try: + model = tf.keras.models.load_model(model) + except: + pass + if isinstance(model, tf.keras.Model) and hasattr(model, 'to_json'): + return 'keras' + if isinstance(model, tf.Graph): + return 'graph' + elif isinstance(model, tf.compat.v1.GraphDef): + return 'graph_def' + elif isinstance(model, tf.compat.v1.estimator.Estimator): + return 'estimator' + elif isinstance(model, str): + model = os.path.abspath(os.path.expanduser(model)) + if (model.endswith('.pb') and os.path.isfile(model)): + if is_saved_model_format(os.path.dirname(model)): + return 'saved_model' + else: + return 'frozen_pb' + elif model.endswith('.ckpt') and os.path.isfile(model): + return 'slim' + elif os.path.isdir(model): + if is_ckpt_format(model): + return 'checkpoint' + elif is_saved_model_format(model): + return 'saved_model' + elif os.path.isfile(model + '.pb'): + return 'frozen_pb' + + raise ValueError('model {} has not recognized model type....'.format(model)) + + + +def validate_graph_node(graph_def, node_names): + """Validate nodes exist in the graph_def + Args: + graph_def (tf.compat.v1.GraphDef): tf.compat.v1.GraphDef object + node_names (list of string): node names to be validated + """ + + if len(node_names) == 0: + return False + all_node_name = [node.name for node in graph_def.node] + for user_name in node_names: + if user_name not in all_node_name: + logger.warn( + str("Node name {} specified in yaml doesn't exist in the model."). + format(user_name)) + return False + return True + +def validate_and_inference_input_output(graph_def, \ + input_tensor_names, output_tensor_names): + """validate and inference the input and output tensor names of graph_def + Args: + graph_def (tf.compat.v1.GraphDef): tf.compat.v1.GraphDef object + input_tensor_names (list of string): input_tensor_names of graph_def + output_tensor_names (list of string): output_tensor_names of graph_def + Returns: + input_tensor_names (list of string): validated input_tensor_names + output_tensor_names (list of string): validated output_tensor_names + """ + from neural_compressor.adaptor.tf_utils.util import get_input_output_node_names + temp_output_tensor_names = [] + if validate_graph_node(graph_def, tensor_to_node(input_tensor_names)): + input_tensor_names = input_tensor_names + else: + input_tensor_names, temp_output_tensor_names = get_input_output_node_names(graph_def) + + if validate_graph_node(graph_def, tensor_to_node(output_tensor_names)): + output_tensor_names = output_tensor_names + elif temp_output_tensor_names: + output_tensor_names = temp_output_tensor_names + else: + _, output_tensor_names = get_input_output_node_names(graph_def) + + return input_tensor_names, output_tensor_names + +def graph_session(model, input_tensor_names, output_tensor_names, **kwargs): + """Build session with tf.compat.v1.Graph + Args: + model (tf.compat.v1.Graph): tf.compat.v1.Graph object + input_tensor_names (list of string): input_tensor_names of model + output_tensor_names (list of string): output_tensor_names of model + Returns: + sess (tf.compat.v1.Session): tf.compat.v1.Session object + input_tensor_names (list of string): validated input_tensor_names + output_tensor_names (list of string): validated output_tensor_names + """ + + config = tf.compat.v1.ConfigProto() + config.use_per_session_threads = 1 + config.inter_op_parallelism_threads = 1 + sess = tf.compat.v1.Session(graph=model, config=config) + + input_tensor_names, output_tensor_names = validate_and_inference_input_output(\ + model.as_graph_def(), input_tensor_names, output_tensor_names) + + return sess, input_tensor_names, output_tensor_names + +def graph_def_session(model, input_tensor_names, output_tensor_names, **kwargs): + """Build session with tf.compat.v1.GraphDef + Args: + model (tf.compat.v1.GraphDef): tf.compat.v1.GraphDef object + input_tensor_names (list of string): input_tensor_names of model + output_tensor_names (list of string): output_tensor_names of model + Returns: + sess (tf.compat.v1.Session): tf.compat.v1.Session object + input_tensor_names (list of string): validated input_tensor_names + output_tensor_names (list of string): validated output_tensor_names + """ + + graph = tf.Graph() + try: + with graph.as_default(): + tf.import_graph_def(model, name='') + except: + input_tensor_names, output_tensor_names = validate_and_inference_input_output(\ + model, input_tensor_names, output_tensor_names) + from neural_compressor.adaptor.tf_utils.util import fix_ref_type_of_graph_def + from neural_compressor.adaptor.tf_utils.util import strip_unused_nodes + model = fix_ref_type_of_graph_def(model) + input_node_names = tensor_to_node(input_tensor_names) + output_node_names = tensor_to_node(output_tensor_names) + model = strip_unused_nodes(model, input_node_names, output_node_names) + with graph.as_default(): + tf.import_graph_def(model, name='') + + return graph_session(graph, input_tensor_names, output_tensor_names, **kwargs) + +def frozen_pb_session(model, input_tensor_names, output_tensor_names, **kwargs): + """Build session with frozen pb + Args: + model (string): model path + input_tensor_names (list of string): input_tensor_names of model + output_tensor_names (list of string): output_tensor_names of model + Returns: + sess (tf.compat.v1.Session): tf.compat.v1.Session object + input_tensor_names (list of string): validated input_tensor_names + output_tensor_names (list of string): validated output_tensor_names + """ + + graph_def = tf.compat.v1.GraphDef() + model = model if model.endswith('.pb') else model + '.pb' + with open(model, 'rb') as f: + graph_def.ParseFromString(f.read()) + return graph_def_session(graph_def, input_tensor_names, \ + output_tensor_names, **kwargs) + +def _contains_function_with_implements_attr(saved_model_proto): + meta_graph = saved_model_proto.meta_graphs[0] + for function in meta_graph.graph_def.library.function: + if function.attr.get("_implements", None) or function.attr.get( + "api_implements", None): + return True + return False + +def load_saved_model(model, saved_model_tags, input_tensor_names, output_tensor_names): + """Load graph_def from saved model with the default serving signature key. + Args: + saved_model_dir: Directory of the SavedModel. + saved_model_tags: Set of tags identifying the MetaGraphDef within the + SavedModel to analyze. + Returns: + graph_def: The loaded GraphDef. + input_tensors: List of input tensors. + output_tensors: List of output tensors. + """ + config = tf.compat.v1.ConfigProto() + config.use_per_session_threads = 1 + config.inter_op_parallelism_threads = 1 + if not os.listdir(os.path.join(model,'variables')): + sess = tf.compat.v1.Session(graph=tf.Graph(), config=config) + loader = tf.compat.v1.saved_model.loader.load(sess, ["serve"], model) + if len(input_tensor_names) == 0: + input_tensor_names = [i.name for _, i in \ + loader.signature_def['serving_default'].inputs.items()] + else: + assert validate_graph_node(\ + sess.graph.as_graph_def(), tensor_to_node(input_tensor_names)), \ + 'tensor names {} not in the graph'.format(input_tensor_names) + + if len(output_tensor_names) == 0: + output_tensor_names = [i.name for _, i in \ + loader.signature_def['serving_default'].outputs.items()] + else: + assert validate_graph_node(\ + sess.graph.as_graph_def(), tensor_to_node(output_tensor_names)), \ + 'tensor names {} not in the graph'.format(output_tensor_names) + + return sess.graph.as_graph_def(), input_tensor_names, output_tensor_names + else: + from tensorflow.python.eager import context + from tensorflow.python.saved_model import load + from tensorflow.python.saved_model import tag_constants + from tensorflow.python.saved_model import signature_constants + from tensorflow.python.framework.convert_to_constants import \ + convert_variables_to_constants_v2 + from tensorflow.python.training import saver + from tensorflow.core.protobuf import config_pb2 + from tensorflow.python.grappler import tf_optimizer + from tensorflow.core.protobuf import meta_graph_pb2 + _saved_model = load.load(model, [tag_constants.SERVING]) + func = _saved_model.signatures[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] + frozen_func = convert_variables_to_constants_v2(func) + grappler_meta_graph_def = saver.export_meta_graph( + graph_def=frozen_func.graph.as_graph_def(), graph=frozen_func.graph) + if len(input_tensor_names) == 0: + input_tensor_names = [i.name.split(':')[0] for i in frozen_func.inputs] + if len(output_tensor_names) == 0: + output_tensor_names = [i.name.split(':')[0] for i in frozen_func.outputs] + # Add a collection 'train_op' so that Grappler knows the outputs. + fetch_collection = meta_graph_pb2.CollectionDef() + for array in frozen_func.inputs + frozen_func.outputs: + fetch_collection.node_list.value.append(array.name) + grappler_meta_graph_def.collection_def["train_op"].CopyFrom( + fetch_collection) + from tensorflow.python.eager import context + grappler_session_config = config_pb2.ConfigProto() + rewrite_options = grappler_session_config.graph_options.rewrite_options + rewrite_options.min_graph_nodes = -1 + opt = tf_optimizer.OptimizeGraph(grappler_session_config, + grappler_meta_graph_def, graph_id=b"tf_graph") + return opt, input_tensor_names, output_tensor_names + +def get_graph_from_saved_model_v2(saved_model_dir, + input_tensor_names, output_tensor_names): + from tensorflow.python.saved_model import tag_constants + from tensorflow.python.saved_model import signature_constants + saved_model_exported_names = [ + signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY + ] + saved_model_tags = set([tag_constants.SERVING]) + return load_saved_model(saved_model_dir, saved_model_tags, + input_tensor_names, output_tensor_names) + +def get_graph_from_original_keras_v2(model, output_dir): + from tensorflow.python.eager import def_function + from tensorflow.lite.python.util import trace_model_call + from tensorflow.lite.python.util import model_input_signature + from tensorflow.python.framework import convert_to_constants + from tensorflow.python.framework import dtypes + from tensorflow.lite.python.util import run_graph_optimizations + from tensorflow.lite.python.convert import OpsSet + from tensorflow.lite.python.util import get_grappler_config + input_signature = None + # If the model's call is not a `tf.function`, then we need to first get its + # input signature from `model_input_signature` method. + if not isinstance(model.call, def_function.Function): + input_signature = model_input_signature(model, keep_original_batch_size=False) + + func = trace_model_call(model, input_signature) + concrete_func = func.get_concrete_function() + funcs = [concrete_func] + + frozen_func, graph_def = ( + convert_to_constants.convert_variables_to_constants_v2_as_graph( + funcs[0], lower_control_flow=False)) + + input_tensors = [ + tensor for tensor in frozen_func.inputs + if tensor.dtype != dtypes.resource + ] + output_tensors = frozen_func.outputs + # Grappler will also try to lower while loop into switch merge + # representation which is undesired for Ophints, so we simply remove + # those attributes to prevent Grappler from doing so. + graph = convert_to_constants.disable_lower_using_switch_merge(graph_def) + # Run function inlining optimization to ensure any models generated + # through the from_frozen_graph path have been inlined. + # grappler_config = get_grappler_config(['function']) + # graph_def = run_graph_optimizations( + # graph, + # input_tensors, + # output_tensors, + # config=grappler_config) + input_names = [tensor.name.split(':')[0] for tensor in input_tensors] + output_names = [tensor.name.split(':')[0] for tensor in output_tensors] + return graph_def, input_names, output_names + +def check_keras_format(model, saved_model_dir): + from tensorflow.python import saved_model + from tensorflow.python.saved_model.load import load + from tensorflow.python.saved_model import save_options + from tensorflow.python.saved_model.loader_impl import parse_saved_model_with_debug_info + version = 'saved_model_v2' + try: + saved_model.save( + model, + saved_model_dir, + options=save_options.SaveOptions(save_debug_info=True)) + except: + return 'trackable_object' + saved_model_proto, _ = parse_saved_model_with_debug_info(saved_model_dir) + saved_model_version = saved_model_proto.saved_model_schema_version + if saved_model_version == 0: + return 'saved_model_v1' + if saved_model_version not in [1, 2]: + raise ValueError("SavedModel file format({0}) is not supported".format( + saved_model_version)) + return version + +def get_graph_from_saved_model_v1(model): + from tensorflow.python.framework import ops + from tensorflow.python.saved_model import constants + from tensorflow.python.client import session + from tensorflow.python.saved_model import tag_constants + from tensorflow.python.saved_model import signature_constants + from tensorflow.lite.python.convert_saved_model import get_meta_graph_def + from tensorflow.lite.python.convert_saved_model import get_signature_def + from tensorflow.lite.python.convert_saved_model import get_inputs_outputs + saved_model_tags = set([tag_constants.SERVING]) + signature_key = signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY + + meta_graph = get_meta_graph_def(model, saved_model_tags) + signature_def = get_signature_def(meta_graph, signature_key) + inputs, outputs = get_inputs_outputs(signature_def) + # Check SavedModel for assets directory. + collection_def = meta_graph.collection_def + if constants.ASSETS_KEY in collection_def: + raise ValueError("SavedModels with assets/ directory are not supported.") + + from tensorflow.python.saved_model import loader + from tensorflow.python.framework import graph_util as tf_graph_util + graph = ops.Graph() + import tensorflow as tf + with session.Session(graph=graph) as sess: + loader.load(sess, meta_graph.meta_info_def.tags, model) + sess.run(tf.compat.v1.global_variables_initializer()) + sess.run(tf.compat.v1.tables_initializer()) + output_nodes = list(set([output.split(':')[0] for output in outputs])) + node_ops = [node.op for node in graph.as_graph_def().node] + if 'MakeIterator' in node_ops: + output_nodes.append('MakeIterator') + table_ops = tf.compat.v1.get_collection( + tf.compat.v1.GraphKeys.TABLE_INITIALIZERS) + # For table initialization + for table_op in table_ops: + output_nodes.append(table_op.name) + if len(table_ops) > 0: + output_nodes.append('init_all_tables') + graph_def = tf_graph_util.convert_variables_to_constants( + sess, graph.as_graph_def(), output_nodes) + return graph_def, inputs, outputs + +def keras_session(model, input_tensor_names, output_tensor_names, **kwargs): + """Build session with keras model + Args: + model (string or tf.keras.Model): model path or tf.keras.Model object + input_tensor_names (list of string): input_tensor_names of model + output_tensor_names (list of string): output_tensor_names of model + Returns: + sess (tf.compat.v1.Session): tf.compat.v1.Session object + input_tensor_names (list of string): validated input_tensor_names + output_tensor_names (list of string): validated output_tensor_names + """ + temp_dir = tempfile.mkdtemp() + if tf.version.VERSION > '2.1.0': + if not isinstance(model, tf.keras.Model): + model = tf.keras.models.load_model(model) + keras_format = check_keras_format(model, temp_dir) + if keras_format == 'saved_model_v2': + try: + graph_def, input_names, output_names = get_graph_from_saved_model_v2( + temp_dir, input_tensor_names, output_tensor_names) + if '_FusedBatchNormEx' in [node.op for node in graph_def.node]: + keras_format = 'trackable_object' + except: + keras_format = 'trackable_object' + if keras_format == 'trackable_object': + try: + graph_def, input_names, output_names = get_graph_from_original_keras_v2( + model, temp_dir) + except: + keras_format = 'saved_model_v1' + if keras_format == 'saved_model_v1': + try: + tf.keras.backend.set_learning_phase(0) + graph_def, input_names, output_names = get_graph_from_saved_model_v1(model) + except: + raise ValueError('Not supported keras model type...') + + # tensorflow 1.x use v1 convert method + else: + tf.keras.backend.set_learning_phase(0) + graph_def, input_names, output_names = get_graph_from_saved_model_v1(model) + shutil.rmtree(temp_dir, True) + return graph_def_session(graph_def, input_names, output_names, **kwargs) + + +def slim_session(model, input_tensor_names, output_tensor_names, **kwargs): + """Build session with slim model + Args: + model (string): model path + input_tensor_names (list of string): input_tensor_names of model + output_tensor_names (list of string): output_tensor_names of model + Returns: + sess (tf.compat.v1.Session): tf.compat.v1.Session object + input_tensor_names (list of string): validated input_tensor_names + output_tensor_names (list of string): validated output_tensor_names + """ + + assert version1_lt_version2(tf.version.VERSION, '2.0.0'), 'slim model only used in tensorflow 1.x' + from .nets_factory import TFSlimNetsFactory + factory = TFSlimNetsFactory() + assert 'name' in kwargs, 'model name should be set in slim checkpoint....' + assert kwargs['name'] in factory.default_slim_models, \ + 'only support topology {}'.format(factory.default_slim_models) + net = copy.deepcopy(factory.networks_map[kwargs['name']]) + model_func = net.pop('model') + arg_scope = net.pop('arg_scope')() + inputs_shape = net.pop('input_shape') + kwargs = net + import tf_slim as slim + with tf.Graph().as_default(): + images = tf.compat.v1.placeholder(name='input', dtype=tf.float32, \ + shape=inputs_shape) + with tf.compat.v1.Session() as sess: + with slim.arg_scope(arg_scope) as scope: # pylint: disable=not-context-manager + model_func(images, is_training=False, **kwargs) + graph_def = sess.graph.as_graph_def() + output_tensor_names = output_tensor_names if len(output_tensor_names) > 0 \ + else [graph_def.node[-1].name] + + from tensorflow.python.tools.freeze_graph import freeze_graph_with_def_protos + graph_def = freeze_graph_with_def_protos( + input_graph_def=graph_def, + input_saver_def=None, + input_checkpoint=model, + output_node_names=','.join(output_tensor_names), + restore_op_name='save/restore_all', + filename_tensor_name='save/Const:0', + output_graph='', + clear_devices=True, + initializer_nodes='') + + return graph_def_session(graph_def, ['input'], output_tensor_names) + +def checkpoint_session(model, input_tensor_names, output_tensor_names, **kwargs): + """Build session with ckpt model + Args: + model (string): model path + input_tensor_names (list of string): input_tensor_names of model + output_tensor_names (list of string): validated output_tensor_names of model + Returns: + sess (tf.compat.v1.Session): tf.compat.v1.Session object + input_tensor_names (list of string): validated input_tensor_names + output_tensor_names (list of string): validated output_tensor_names + """ + + assert output_tensor_names is not None and len(output_tensor_names) > 0, \ + 'outputs should not be None of checkpoint....' + + ckpt_prefix = [os.path.splitext(i)[0] for i in os.listdir(model) \ + if i.endswith(".meta")][0] + + config = tf.compat.v1.ConfigProto() + config.use_per_session_threads = 1 + config.inter_op_parallelism_threads = 1 + graph = tf.Graph() + sess = tf.compat.v1.Session(graph=graph, config=config) + with graph.as_default(): + saver = tf.compat.v1.train.import_meta_graph(\ + os.path.join(model, ckpt_prefix + '.meta'), clear_devices=True) + + sess.run(tf.compat.v1.global_variables_initializer()) + saver.restore(sess, os.path.join(model, ckpt_prefix)) + + from neural_compressor.adaptor.tf_utils.util import get_input_output_node_names + if validate_graph_node(sess.graph.as_graph_def(), tensor_to_node(input_tensor_names)): + input_tensor_names = input_tensor_names + else: + input_tensor_names, _ = get_input_output_node_names(sess.graph.as_graph_def()) + return sess, input_tensor_names, output_tensor_names + +def estimator_session(model, input_tensor_names, output_tensor_names, **kwargs): + """Build session with estimator model + Args: + model (tf.estimator.Estimator): tf.estimator.Estimator object + input_tensor_names (list of string): input_tensor_names of model + output_tensor_names (list of string): output_tensor_names of model + kwargs (dict): other required parameters, like input_fn + Returns: + sess (tf.compat.v1.Session): tf.compat.v1.Session object + input_tensor_names (list of string): validated input_tensor_names + output_tensor_names (list of string): validated output_tensor_names + """ + + assert 'input_fn' in kwargs, 'input func should be supplied for estimator session....' + with tf.Graph().as_default() as g: + features, input_hooks = model._get_features_from_input_fn( + kwargs['input_fn'], tf.estimator.ModeKeys.PREDICT) + estimator_spec = model._call_model_fn(features, None, + tf.estimator.ModeKeys.PREDICT, model.config) + + if len(output_tensor_names) == 0: + outputs = [tensor.name for tensor in estimator_spec.predictions.values()] if\ + isinstance(estimator_spec.predictions, dict) else \ + [estimator_spec.predictions.name] + else: + outputs = output_tensor_names + + logger.info("Estimator output tensor names are {}.".format(outputs)) + with tf.compat.v1.Session(graph=g) as sess: + sess.run(tf.compat.v1.global_variables_initializer()) + # Freezing a graph requires output_node_names, which can be found in + # estimator_spec.predictions that contains prediction tensors as a + # dictionary + # When a model uses Iterator, we need to have 'MakeIterator' (default + # name used by TF) in the output_node_names as well. + output_nodes = list(set([output.split(':')[0] for output in outputs])) + if 'MakeIterator' in [node.op for node in g.as_graph_def().node]: + output_nodes.append('MakeIterator') + + graph_def = tf.compat.v1.graph_util.convert_variables_to_constants(sess, + g.as_graph_def(), output_nodes) + + return graph_def_session(graph_def, input_tensor_names, outputs) + +def saved_model_session(model, input_tensor_names, output_tensor_names, **kwargs): + """Build session with saved model + Args: + model (string): model path + input_tensor_names (list of string): input_tensor_names of model + output_tensor_names (list of string): output_tensor_names of model + Returns: + sess (tf.compat.v1.Session): tf.compat.v1.Session object + input_tensor_names (list of string): validated input_tensor_names + output_tensor_names (list of string): validated output_tensor_names + """ + try: + graph_def, input_names, output_names = get_graph_from_saved_model_v2( + model, input_tensor_names, output_tensor_names) + except: + graph_def, input_names, output_names = get_graph_from_saved_model_v1(model) + assert graph_def is not None, 'Can not parse the saved model...' + return graph_def_session(graph_def, input_names, output_names, **kwargs) + +# it's necessary that a session with input output tensors to run the model +SESSIONS = {'frozen_pb': frozen_pb_session, + 'graph_def': graph_def_session, + 'graph': graph_session, + 'saved_model': saved_model_session, + 'keras': keras_session, + 'checkpoint': checkpoint_session, + 'estimator': estimator_session, + 'slim': slim_session,} + + +class TensorflowBaseModel(BaseModel): + """Build TensorflowBaseModel object + Args: + model (string or tensorflow model object): model path or model object + kwargs (dict): other required parameters, like input_fn + """ + + def __init__(self, model, **kwargs): + + self._model = model + self._name = '' + self._weights = None + self.kwargs = kwargs + self._graph_info = {} + self._input_tensor_names = [] + self._output_tensor_names = [] + self._model_type = '' + self._sess = None + self._iter_op = None + self._workspace_path = '' + self._q_config = None + + def framework(self): + return 'tensorflow' + + @property + def name(self): + return self._name + + @name.setter + def name(self, name): + self.kwargs.update({'name': name}) + self._name = name + + @property + def weights(self): + """ Getter to weights """ + return self._weights + + @weights.setter + def weights(self, new_weights): + """ Setter to weights """ + self._weights = new_weights + + @property + def q_config(self): + return self._q_config + + @q_config.setter + def q_config(self, q_config): + self._q_config = q_config + + @property + def workspace_path(self): + return self._workspace_path + + @workspace_path.setter + def workspace_path(self, path): + self._workspace_path = path + + @property + def model_type(self): + return self._model_type + + @model_type.setter + def model_type(self, model_type): + assert model_type in SESSIONS, 'model type not supported....' + self._model_type = model_type + + @property + def model(self): + return self.graph + + @property + def graph_def(self): + return self.graph.as_graph_def() + + @property + def graph_info(self): + self._graph_info = {} + for node in self.graph_def.node: + self._graph_info[node.name] = node.op + return self._graph_info + + @property + def sess(self): + if self._sess is None: + self._load_sess(self._model, **self.kwargs) + return self._sess + + @property + def graph(self): + return self.sess.graph + + @graph_def.setter + def graph_def(self, graph_def): + if self._sess is not None: + self._sess.close() + output_sess = SESSIONS['graph_def'](graph_def,\ + self._input_tensor_names, \ + self._output_tensor_names) + + self._sess = output_sess[0] + self._input_tensor_names = output_sess[1] + self._output_tensor_names = output_sess[2] + self.model_type = 'graph_def' + + def _load_sess(self, model, **kwargs): + if self.name: + kwargs.update({'name': self.name}) + # assert self.model_type, 'model type not set....' + output_sess = SESSIONS[self.model_type](model, + self._input_tensor_names, \ + self._output_tensor_names, + **kwargs) + self._sess = output_sess[0] + self._input_tensor_names = output_sess[1] + self._output_tensor_names = output_sess[2] + + tf.compat.v1.get_variable_scope().reuse_variables() + return self._sess + + @property + def iter_op(self): + self._iter_op = [] + if self._sess is None: + self._load_sess(self._model, **self.kwargs) + op_list = [node.op for node in self._sess.graph.as_graph_def().node] + if 'MakeIterator' in op_list: + self._iter_op.append(self._sess.graph.get_operation_by_name('MakeIterator')) + return self._iter_op + + @property + def input_tensor_names(self): + if self._sess is None: + self._load_sess(self._model, **self.kwargs) + return copy.deepcopy(self._input_tensor_names) + + @input_tensor_names.setter + def input_tensor_names(self, tensor_names): + if len(tensor_names) == 0: + logger.warn("Input tensor names is empty.") + return + if self._sess is not None: + assert validate_graph_node(\ + self.graph_def, tensor_to_node(tensor_names)), \ + 'tensor names {} not in graph'.format(tensor_names) + self._input_tensor_names = tensor_names + + @property + def output_tensor_names(self): + if len(self._output_tensor_names) == 0: + self._load_sess(self._model, **self.kwargs) + return copy.deepcopy(self._output_tensor_names) + + @output_tensor_names.setter + def output_tensor_names(self, tensor_names): + if len(tensor_names) == 0: + logger.warn("Output tensor names should not be empty.") + return + if self._sess is not None: + assert validate_graph_node(\ + self.graph_def, tensor_to_node(tensor_names)), \ + 'tensor names {} not in graph'.format(tensor_names) + self._output_tensor_names = tensor_names + + # input/output node names and input/output tensor + # come from input/output tensor names, so do not support assign these values + @property + def input_node_names(self): + return copy.deepcopy(tensor_to_node(self.input_tensor_names)) + + @property + def output_node_names(self): + output_node_names = tensor_to_node(self.output_tensor_names) + iter_op_list = self.iter_op + if iter_op_list != []: + output_node_names += [iter_op.name for iter_op in iter_op_list] + return copy.deepcopy(output_node_names) + + @property + def input_tensor(self): + from neural_compressor.adaptor.tf_utils.util import get_tensor_by_name + return [get_tensor_by_name(\ + self.graph, x) for x in self.input_tensor_names] + + @property + def output_tensor(self): + from neural_compressor.adaptor.tf_utils.util import get_tensor_by_name + return [get_tensor_by_name(\ + self.graph, x) for x in self.output_tensor_names] + + def save(self, root=None): + if not root: + root = cfg.default_workspace + '/save.pb' + root = os.path.abspath(os.path.expanduser(root)) + # if not have suffix, default append .pb + os.makedirs(os.path.dirname(root), exist_ok=True) + pb_file = root if os.path.split(root)[-1].endswith('.pb') else root + '.pb' + f = tf.io.gfile.GFile(pb_file, 'wb') + f.write(self.graph_def.SerializeToString()) + logger.info("Save quantized model to {}.".format(pb_file)) + + +class TensorflowSavedModelModel(TensorflowBaseModel): + def get_all_weight_names(self): + import tensorflow as tf + names = [] + for index, layer in enumerate(tf.keras.models.load_model(self._model).layers): + if len(layer.weights): + names.append(index) + return names + + def update_weights(self, tensor_name, new_tensor): + pass + + def get_weight(self, tensor_name): + return self.weights[tensor_name] + + @property + def model(self): + import time + import shutil + root = os.path.abspath(os.path.expanduser(cfg.default_workspace)) + root += str(time.time()) + if os.path.exists(root): + shutil.rmtree(root) + os.makedirs(root, exist_ok=True) + if not self._sess: + self._load_sess(self._model, **self.kwargs) + _, builder = self.build_saved_model(root) + builder.save() + model = tf.saved_model.load(root) + shutil.rmtree(root) + return model + + def report_sparsity(self): + """ Get sparsity of the model + Args: + Returns: + df (DataFrame): DataFrame of sparsity of each weight + total_sparsity (float): total sparsity of model + """ + import pandas as pd + import tensorflow as tf + import numpy as np + df = pd.DataFrame(columns=['Name', 'Shape', 'NNZ (dense)', 'NNZ (sparse)', "Sparsity(%)", + 'Std', 'Mean', 'Abs-Mean']) + pd.set_option('display.precision', 2) + param_dims = [2, 4] + params_size = 0 + sparse_params_size = 0 + for index, layer in enumerate(tf.keras.models.load_model(self._model).layers): + if not len(layer.weights): + continue + # Extract just the actual parameter's name, which in this context we treat + # as its "type" + weights = layer.get_weights()[0] + if weights.ndim in param_dims: + param_size, sparse_param_size, dense_param_size = compute_sparsity( + weights) + density = dense_param_size / param_size + params_size += param_size + sparse_params_size += sparse_param_size + df.loc[len(df.index)] = ([ + index, + list(weights.shape), + dense_param_size, + sparse_param_size, + (1 - density) * 100, + np.std(weights), + np.mean(weights), + np.mean(np.abs(weights)) + ]) + + total_sparsity = sparse_params_size / params_size * 100 + + df.loc[len(df.index)] = ([ + 'Total sparsity:', + params_size, + "-", + int(sparse_params_size), + total_sparsity, + 0, 0, 0]) + + return df, total_sparsity + + def build_saved_model(self, root=None): + if not root: + root = cfg.default_workspace + root = os.path.abspath(os.path.expanduser(root)) + if os.path.exists(root): + import shutil + shutil.rmtree(root) + + os.makedirs(root, exist_ok=True) + + from tensorflow.python.saved_model import signature_constants + from tensorflow.python.saved_model import tag_constants + from neural_compressor.adaptor.tf_utils.util import get_tensor_by_name + builder = tf.compat.v1.saved_model.builder.SavedModelBuilder(root) + sigs = {} + with tf.compat.v1.Session(graph=tf.Graph()) as sess: + #(TODO) not directly use self._sess.graph, use self.graph + tf.import_graph_def(self.graph.as_graph_def(), name="") + g = tf.compat.v1.get_default_graph() + inp = [get_tensor_by_name(g, x) for x in self._input_tensor_names] + out = [get_tensor_by_name(g, x) for x in self._output_tensor_names] + sigs[signature_constants.DEFAULT_SERVING_SIGNATURE_DEF_KEY] = \ + tf.compat.v1.saved_model.signature_def_utils.predict_signature_def( + {k: v for k, v in zip(self._input_tensor_names, inp)}, + {k: v for k, v in zip(self._output_tensor_names, out)}) + builder.add_meta_graph_and_variables(sess, + [tag_constants.SERVING], + signature_def_map=sigs) + return root, builder + + def save(self, root=None): + root, builder = self.build_saved_model(root) + builder.save() + logger.info("Save quantized model to {}.".format(root)) + +class TensorflowQATModel(TensorflowSavedModelModel): + def __init__(self, model='', **kwargs): + super(TensorflowQATModel, self).__init__(model) + self.keras_model = None + self.model_type = 'keras' + + @property + def model(self): + if self.keras_model == None: + self.keras_model = tf.keras.models.load_model(self._model) + return self.keras_model + + @model.setter + def model(self, q_model): + self.keras_model = q_model + + def save(self, root=None): + if not root: + root = cfg.default_workspace + '/saved_model' + root = os.path.abspath(os.path.expanduser(root)) + # if not have suffix, default append .pb + os.makedirs(os.path.dirname(root), exist_ok=True) + q_aware_model = self.keras_model + q_aware_model.save(root) + saved_format = 'saved_model' + if root.endswith('.h5'): + saved_format = 'h5 file' + logger.info("Save quantized model to {}.".format(saved_format)) + return root + +class TensorflowCheckpointModel(TensorflowBaseModel): + + @property + def graph_def(self): + if self.model_type == 'graph_def': + return self.sess.graph.as_graph_def() + from neural_compressor.adaptor.tf_utils.util import _parse_ckpt_bn_input + from tensorflow.python.framework import graph_util + graph_def = self.sess.graph.as_graph_def() + graph_def = _parse_ckpt_bn_input(graph_def) + return graph_util.convert_variables_to_constants( + sess=self._sess, + input_graph_def=graph_def, + output_node_names=self.output_node_names) + + @graph_def.setter + def graph_def(self, graph_def): + if self._sess is not None: + self._sess.close() + output_sess = SESSIONS['graph_def'](graph_def, + self._input_tensor_names, \ + self._output_tensor_names) + self._sess = output_sess[0] + self._input_tensor_names = output_sess[1] + self._output_tensor_names = output_sess[2] + self.model_type = 'graph_def' + + +TENSORFLOW_MODELS = {'frozen_pb': TensorflowBaseModel, + 'graph_def': TensorflowBaseModel, + 'graph': TensorflowBaseModel, + 'checkpoint': TensorflowCheckpointModel, + 'estimator': TensorflowBaseModel, + 'slim': TensorflowBaseModel, + 'saved_model': TensorflowSavedModelModel, + 'keras': TensorflowSavedModelModel + } + +class TensorflowModel(object): + def __new__(cls, model_type, root, **kwargs): + os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" + os.environ["CUDA_VISIBLE_DEVICES"] = "-1" + model = TENSORFLOW_MODELS[model_type](root, **kwargs) + model.model_type = model_type + return model diff --git a/neural_compressor/model/torch_model.py b/neural_compressor/model/torch_model.py index 67f8ae159bb..5560158358e 100644 --- a/neural_compressor/model/torch_model.py +++ b/neural_compressor/model/torch_model.py @@ -695,15 +695,15 @@ def __init__(self, model, **kwargs): super(PyTorchFXModel, self).__init__(model, **kwargs) -class PyTorchIpexModel(PyTorchBaseModel): # pragma: no cover - """Build PyTorchIpexModel object +class IPEXModel(PyTorchBaseModel): # pragma: no cover + """Build IPEXModel object Args: model (onnx model): model path """ def __init__(self, model, **kwargs): - super(PyTorchIpexModel, self).__init__(model, **kwargs) + super(IPEXModel, self).__init__(model, **kwargs) self.ipex_config_path = None @property diff --git a/neural_compressor/quantization.py b/neural_compressor/quantization.py index 36504b5c12e..6d50c622756 100644 --- a/neural_compressor/quantization.py +++ b/neural_compressor/quantization.py @@ -15,11 +15,11 @@ # See the License for the specific language governing permissions and # limitations under the License. + from .experimental import Quantization as ExpQuantization from neural_compressor.conf.pythonic_config import Config from neural_compressor.config import PostTrainingQuantConfig - def fit(model, conf, calib_dataloader=None, @@ -29,7 +29,6 @@ def fit(model, eval_metric=None, **kwargs): """Quantize the model with a given configure. - Args: model (torch.nn.Module): For Tensorflow model, it could be a path to frozen pb,loaded graph_def object or @@ -75,15 +74,12 @@ def fit(model, encapsulated in this function implementation and outputs a higher-is-better accuracy scalar value. - The pseudo code should be something like: - def eval_func(model): input, label = dataloader() output = model(input) accuracy = metric(output, label) return accuracy - """ if isinstance(conf, PostTrainingQuantConfig): diff --git a/neural_compressor/utils/create_obj_from_config.py b/neural_compressor/utils/create_obj_from_config.py index 8d22945c522..61dfcd4ac01 100644 --- a/neural_compressor/utils/create_obj_from_config.py +++ b/neural_compressor/utils/create_obj_from_config.py @@ -16,7 +16,7 @@ # limitations under the License. from neural_compressor.experimental.metric import METRICS -from neural_compressor.experimental.data import DATASETS, TRANSFORMS, FILTERS, DATALOADERS +from neural_compressor.experimental.data import Datasets, TRANSFORMS, FILTERS, DATALOADERS from neural_compressor.experimental.common import Optimizers, Criterions from collections import OrderedDict import copy @@ -65,7 +65,7 @@ def create_dataset(framework, data_source, cfg_preprocess, cfg_filter): preprocesses = TRANSFORMS(framework, 'preprocess') preprocess = get_preprocess(preprocesses, cfg_preprocess) # even we can unify transform, how can we handle the IO, or we do the transform here - datasets = DATASETS(framework) + datasets = Datasets(framework) dataset_type = list(data_source.keys())[0] # generate framework and dataset specific filters filter = None diff --git a/neural_compressor/ux/components/model/model_type_getter.py b/neural_compressor/ux/components/model/model_type_getter.py index 1d0913181a0..0fd151692b3 100644 --- a/neural_compressor/ux/components/model/model_type_getter.py +++ b/neural_compressor/ux/components/model/model_type_getter.py @@ -14,7 +14,7 @@ # limitations under the License. """Model type getter.""" -from neural_compressor.model.model import get_model_type as nc_get_model_type +from neural_compressor.model.tensorflow_model import get_model_type as nc_get_model_type from neural_compressor.ux.utils.expiring_dict import ExpiringDict model_type_cache = ExpiringDict(ttl=600) diff --git a/neural_compressor/ux/components/model/tensorflow/model.py b/neural_compressor/ux/components/model/tensorflow/model.py index 4f33b161c85..e3218634088 100644 --- a/neural_compressor/ux/components/model/tensorflow/model.py +++ b/neural_compressor/ux/components/model/tensorflow/model.py @@ -17,7 +17,7 @@ from typing import Any, List, Optional from neural_compressor.experimental.common.model import Model as NCModel -from neural_compressor.model.model import TensorflowBaseModel +from neural_compressor.model.tensorflow_model import TensorflowBaseModel from neural_compressor.utils.logger import Logger from neural_compressor.ux.components.graph.graph import Graph from neural_compressor.ux.components.graph.reader.tensorflow_reader import TensorflowReader diff --git a/test/adaptor/onnxrt_adaptor/test_adaptor_onnxrt.py b/test/adaptor/onnxrt_adaptor/test_adaptor_onnxrt.py index 8e4113deb58..0457546e50a 100644 --- a/test/adaptor/onnxrt_adaptor/test_adaptor_onnxrt.py +++ b/test/adaptor/onnxrt_adaptor/test_adaptor_onnxrt.py @@ -10,7 +10,7 @@ from onnx import onnx_pb as onnx_proto from onnx import helper, TensorProto, numpy_helper from neural_compressor.adaptor import FRAMEWORKS -from neural_compressor.data import DATASETS, DATALOADERS +from neural_compressor.data import Datasets, DATALOADERS from neural_compressor.experimental import Quantization, common from neural_compressor.experimental import Benchmark, common from neural_compressor import options @@ -529,7 +529,7 @@ def build_gemm_model(): def build_benchmark(): seq = ''' from neural_compressor.experimental import Benchmark -from neural_compressor.data import DATASETS, DATALOADERS +from neural_compressor.data import Datasets, DATALOADERS from neural_compressor import conf from onnx import onnx_pb as onnx_proto from onnx import helper, TensorProto, numpy_helper @@ -555,7 +555,7 @@ def reverse_matrix(x): graph = helper.make_graph(nodes, 'test0', [input0], [output0]) model = helper.make_model(graph, **{'opset_imports': [helper.make_opsetid('', 13)]}) -datasets = DATASETS('onnxrt_qlinearops') +datasets = Datasets('onnxrt_qlinearops') ext_dataset = datasets['dummy'](shape=(10, 2), low=0., high=1., label=True) ext_dataloader = DATALOADERS['onnxrt_qlinearops'](ext_dataset) @@ -590,26 +590,26 @@ class TestAdaptorONNXRT(unittest.TestCase): rn50_export_path = "rn50.onnx" rn50_model = torchvision.models.resnet50() - datasets = DATASETS('onnxrt_qlinearops') + datasets = Datasets('onnxrt_qlinearops') cv_dataset = datasets['dummy'](shape=(10, 3, 224, 224), low=0., high=1., label=True) cv_dataloader = DATALOADERS['onnxrt_qlinearops'](cv_dataset) ir3_dataset = datasets['dummy'](shape=(10, 2048), low=0., high=1., label=True) ir3_dataloader = DATALOADERS['onnxrt_qlinearops'](ir3_dataset) - gather_dataset = DATASETS('onnxrt_qlinearops')['dummy'](shape=(5, 100, 4), label=True) + gather_dataset = Datasets('onnxrt_qlinearops')['dummy'](shape=(5, 100, 4), label=True) gather_dataloader = DATALOADERS['onnxrt_qlinearops'](gather_dataset) ext_dataset = datasets['dummy'](shape=(10, 2), low=0., high=1., label=True) ext_dataloader = DATALOADERS['onnxrt_qlinearops'](ext_dataset) - rename_dataset = DATASETS('onnxrt_qlinearops')['dummy'](shape=(5, 1, 200), label=True) + rename_dataset = Datasets('onnxrt_qlinearops')['dummy'](shape=(5, 1, 200), label=True) rename_dataloader = DATALOADERS['onnxrt_qlinearops'](rename_dataset) matmul_dataset = MatmulDataset() matmul_dataloader = DATALOADERS['onnxrt_qlinearops'](matmul_dataset) - conv_dataset = DATASETS('onnxrt_qlinearops')['dummy'](shape=(10, 3, 1, 3), label=True) + conv_dataset = Datasets('onnxrt_qlinearops')['dummy'](shape=(10, 3, 1, 3), label=True) conv_dataloader = DATALOADERS['onnxrt_qlinearops'](conv_dataset) @classmethod diff --git a/test/adaptor/onnxrt_adaptor/test_onnxrt_augment.py b/test/adaptor/onnxrt_adaptor/test_onnxrt_augment.py index 8051b086642..05b83bd17d8 100644 --- a/test/adaptor/onnxrt_adaptor/test_onnxrt_augment.py +++ b/test/adaptor/onnxrt_adaptor/test_onnxrt_augment.py @@ -11,7 +11,7 @@ from neural_compressor.experimental.data.datasets.dataset import Dataset from neural_compressor.adaptor.ox_utils.calibration import ONNXRTAugment from neural_compressor.model.onnx_model import ONNXModel -from neural_compressor.data import DATASETS, DATALOADERS +from neural_compressor.data import Datasets, DATALOADERS def generate_input_initializer(tensor_shape, tensor_dtype, input_name): ''' @@ -55,7 +55,7 @@ def create_nlp_session(): node = onnx.helper.make_node('Gather', ['D', 'B'], ['C'], name='gather') graph = helper.make_graph([squeeze, node], 'test_graph_1', [A], [C], [B_init]) model = helper.make_model(graph, **{'opset_imports': [helper.make_opsetid('', 13)]}) - datasets = DATASETS('onnxrt_qlinearops') + datasets = Datasets('onnxrt_qlinearops') dataset = datasets['dummy_v2'](input_shape=(100, 4), label_shape=(1,)) dataloader = DATALOADERS['onnxrt_qlinearops'](dataset) diff --git a/test/adaptor/pytorch_adaptor/test_adaptor_pytorch_2.x.py b/test/adaptor/pytorch_adaptor/test_adaptor_pytorch_2.x.py index e95985787a9..9bb5a122fa0 100644 --- a/test/adaptor/pytorch_adaptor/test_adaptor_pytorch_2.x.py +++ b/test/adaptor/pytorch_adaptor/test_adaptor_pytorch_2.x.py @@ -7,8 +7,8 @@ import unittest import os from neural_compressor import PostTrainingQuantConfig, QuantizationAwareTrainingConfig, set_workspace -from neural_compressor.data import DATASETS, DATALOADERS -from neural_compressor.experimental.data.datasets.dataset import DATASETS +from neural_compressor.data import Datasets, DATALOADERS, DataLoader +from neural_compressor.experimental.data.datasets.dataset import Datasets from neural_compressor import quantization from neural_compressor.training import prepare_compression from neural_compressor.utils.pytorch import load @@ -307,7 +307,7 @@ def tearDownClass(self): def test_fx_quant(self): for approach in ["qat", "static"]: model_origin = resnet18() - dataset = DATASETS("pytorch")["dummy"]((10, 3, 224, 224), label=True) + dataset = Datasets("pytorch")["dummy"]((10, 3, 224, 224), label=True) dataloader = DATALOADERS["pytorch"](dataset) if approach == "qat": model = copy.deepcopy(model_origin) @@ -343,7 +343,7 @@ def test_fx_quant(self): for approach in ["qat", "static"]: model_origin = M() # run fx_quant in neural_compressor and save the quantized GraphModule - dataset = DATASETS("pytorch")["dummy"]((100, 3, 224, 224), label=True) + dataset = Datasets("pytorch")["dummy"]((100, 3, 224, 224), label=True) dataloader = DATALOADERS["pytorch"](dataset) if approach == "qat": model = copy.deepcopy(model_origin) @@ -419,7 +419,7 @@ def eval_func(model): nhid = 256, nlayers = 2, ) - dataset = DATASETS("pytorch")["dummy"]((3, 10)) + dataset = Datasets("pytorch")["dummy"]((3, 10)) dataloader = DATALOADERS["pytorch"](dataset) # run fx_quant in neural_compressor and save the quantized GraphModule if approach == "qat": @@ -445,7 +445,7 @@ def eval_func(model): def test_fx_sub_module_quant(self): for approach in ["qat", "static"]: model_origin = DynamicControlModel() - dataset = DATASETS("pytorch")["dummy"]((1, 3, 224, 224)) + dataset = Datasets("pytorch")["dummy"]((1, 3, 224, 224)) dataloader = DATALOADERS["pytorch"](dataset) # run fx_quant in neural_compressor and save the quantized GraphModule if approach == "qat": @@ -484,8 +484,8 @@ def test_fx_sub_module_quant(self): def test_mix_precision(self): model_origin = DynamicControlModel() # run fx_quant in neural_compressor and save the quantized GraphModule - dataset = DATASETS("pytorch")["dummy"]((100, 3, 224, 224)) - dataloader = DATALOADERS["pytorch"](dataset) + dataset = Datasets("pytorch")["dummy"]((100, 3, 224, 224)) + dataloader = DataLoader("pytorch", dataset) set_workspace=("./saved") conf = PostTrainingQuantConfig(op_name_list=ptq_fx_op_name_list) q_model = quantization.fit(model_origin, diff --git a/test/adaptor/pytorch_adaptor/test_torch2onnx.py b/test/adaptor/pytorch_adaptor/test_torch2onnx.py index a9efc2c1eb6..67b71df76d6 100644 --- a/test/adaptor/pytorch_adaptor/test_torch2onnx.py +++ b/test/adaptor/pytorch_adaptor/test_torch2onnx.py @@ -9,7 +9,7 @@ import neural_compressor.adaptor.pytorch as nc_torch from neural_compressor import quantization from neural_compressor.config import PostTrainingQuantConfig -from neural_compressor.experimental.data.datasets.dataset import DATASETS +from neural_compressor.experimental.data.datasets.dataset import Datasets from packaging.version import Version from torch.quantization import QuantStub, DeQuantStub @@ -209,7 +209,7 @@ def test_fx_quant(self): model = DynamicControlModel() # run fx_quant in neural_compressor and save the quantized GraphModule conf = PostTrainingQuantConfig(approach=approach) - dataset = DATASETS("pytorch")['dummy']((100, 3, 224, 224)) + dataset = Datasets("pytorch")['dummy']((100, 3, 224, 224)) dataloader = torch.utils.data.DataLoader(dataset) q_model = quantization.fit(model, conf, diff --git a/test/benchmark/test_benchmark.py b/test/benchmark/test_benchmark.py index 37aef1ca500..1a0450b4425 100644 --- a/test/benchmark/test_benchmark.py +++ b/test/benchmark/test_benchmark.py @@ -45,8 +45,8 @@ def build_benchmark(): arg_parser = ArgumentParser(description='Parse args') arg_parser.add_argument('--input_model', dest='input_model', default='input_model', help='input odel') args = arg_parser.parse_args() -from neural_compressor.data import DATASETS -dataset = DATASETS('tensorflow')['dummy']((100, 32, 32, 1), label=True) +from neural_compressor.data import Datasets +dataset = Datasets('tensorflow')['dummy']((100, 32, 32, 1), label=True) from neural_compressor.experimental import Benchmark, common from neural_compressor.conf.config import BenchmarkConf benchmarker = Benchmark('fake_yaml.yaml') @@ -60,8 +60,8 @@ def build_benchmark(): arg_parser = ArgumentParser(description='Parse args') arg_parser.add_argument('--input_model', dest='input_model', default='input_model', help='input odel') args = arg_parser.parse_args() -from neural_compressor.data import DATASETS -dataset = DATASETS('tensorflow')['dummy']((100, 32, 32, 1), label=True) +from neural_compressor.data import Datasets +dataset = Datasets('tensorflow')['dummy']((100, 32, 32, 1), label=True) from neural_compressor.experimental import Benchmark, common from neural_compressor.conf.config import BenchmarkConf conf = BenchmarkConf('fake_yaml.yaml') @@ -94,8 +94,8 @@ def build_benchmark2(): "arg_parser.add_argument('--input_model', dest='input_model', default='input_model', help='input model')\n", "args = arg_parser.parse_args()\n", - "from neural_compressor.data import DATASETS\n", - "dataset = DATASETS('tensorflow')['dummy']((5, 32, 32, 1), label=True)\n", + "from neural_compressor.data import Datasets\n", + "dataset = Datasets('tensorflow')['dummy']((5, 32, 32, 1), label=True)\n", "from neural_compressor.experimental import Benchmark, common\n", "benchmarker = Benchmark()\n", diff --git a/test/benchmark/test_benchmark_2.x.py b/test/benchmark/test_benchmark_2.x.py index 0829811081a..720210ddf23 100644 --- a/test/benchmark/test_benchmark_2.x.py +++ b/test/benchmark/test_benchmark_2.x.py @@ -17,9 +17,9 @@ def build_benchmark(): args = arg_parser.parse_args() from neural_compressor.benchmark import fit from neural_compressor.config import BenchmarkConfig -from neural_compressor.data import DATASETS +from neural_compressor.data import Datasets from neural_compressor.experimental import common -dataset = DATASETS('tensorflow')['dummy']((100, 32, 32, 1), label=True) +dataset = Datasets('tensorflow')['dummy']((100, 32, 32, 1), label=True) b_dataloader = common.DataLoader(dataset, batch_size=10) conf = BenchmarkConfig(warmup=5, iteration=10, cores_per_instance=4, num_of_instance=2) fit(args.input_model, conf, b_dataloader=b_dataloader) @@ -32,8 +32,8 @@ def build_benchmark(): args = arg_parser.parse_args() from neural_compressor.benchmark import fit from neural_compressor.config import BenchmarkConfig -from neural_compressor.data import DATASETS -dataset = DATASETS('tensorflow')['dummy']((100, 32, 32, 1), label=True) +from neural_compressor.data import Datasets +dataset = Datasets('tensorflow')['dummy']((100, 32, 32, 1), label=True) from neural_compressor.experimental import common conf = BenchmarkConfig(warmup=5, iteration=10, cores_per_instance=4, num_of_instance=2) b_dataloader = common.DataLoader(dataset, batch_size=10) @@ -63,8 +63,8 @@ def build_benchmark2(): "arg_parser.add_argument('--input_model', dest='input_model', default='input_model', help='input model')\n", "args = arg_parser.parse_args()\n", "from neural_compressor.benchmark import fit\n" - "from neural_compressor.data import DATASETS\n", - "dataset = DATASETS('tensorflow')['dummy']((5, 32, 32, 1), label=True)\n", + "from neural_compressor.data import Datasets\n", + "dataset = Datasets('tensorflow')['dummy']((5, 32, 32, 1), label=True)\n", "from neural_compressor.experimental import common\n", "b_dataloader = common.DataLoader(dataset)\n", diff --git a/test/config/test_pythonic_config.py b/test/config/test_pythonic_config.py index 8283324f7d5..f72686982b5 100644 --- a/test/config/test_pythonic_config.py +++ b/test/config/test_pythonic_config.py @@ -19,7 +19,7 @@ from torch import nn from neural_compressor.conf.pythonic_config import OpQuantConf, ActivationConf, WeightConf -from neural_compressor.data import DATASETS +from neural_compressor.data import Datasets from neural_compressor.experimental import Quantization, Distillation, Pruning, NAS, common from neural_compressor.experimental.data.dataloaders.pytorch_dataloader import PyTorchDataLoader from neural_compressor.adaptor import FRAMEWORKS @@ -199,7 +199,7 @@ def test_distillation(self): distiller.teacher_model = ConvNet(16, 32) # Customized train, evaluation - datasets = DATASETS('pytorch') + datasets = Datasets('pytorch') dummy_dataset = datasets['dummy'](shape=(32, 3, 64, 64), low=0., high=1., label=True) dummy_dataloader = PyTorchDataLoader(dummy_dataset) def train_func(model): @@ -241,7 +241,7 @@ def test_pruning(self): prune.model = model # Customized train, evaluation - datasets = DATASETS('pytorch') + datasets = Datasets('pytorch') dummy_dataset = datasets['dummy'](shape=(32, 3, 64, 64), low=0., high=1., label=True) dummy_dataloader = PyTorchDataLoader(dummy_dataset) def train_func(model): diff --git a/test/data/test_dataloader.py b/test/data/test_dataloader.py index de9017b4e8c..d62bb8c5488 100644 --- a/test/data/test_dataloader.py +++ b/test/data/test_dataloader.py @@ -6,7 +6,7 @@ import shutil from neural_compressor.utils.create_obj_from_config import create_dataset, create_dataloader from neural_compressor.data.dataloaders.dataloader import DataLoader -from neural_compressor.data import DATASETS, DATALOADERS, TRANSFORMS +from neural_compressor.data import Datasets, DATALOADERS, TRANSFORMS from PIL import Image class TestBuiltinDataloader(unittest.TestCase): @@ -1069,7 +1069,7 @@ def test_pytorch_bert_dataset(self): self.assertEqual(5, len(ds[0][0])) def test_tensorflow_dummy(self): - datasets = DATASETS('tensorflow') + datasets = Datasets('tensorflow') dataset = datasets['dummy'](shape=(4, 256, 256, 3)) data_loader = DATALOADERS['tensorflow'](dataset) @@ -1092,7 +1092,7 @@ def test_tensorflow_dummy(self): dataset = datasets['dummy'](shape=(4, 256, 256, 3), dtype=['float32', 'int8']) def test_tensorflow_dummy_v2(self): - datasets = DATASETS('tensorflow') + datasets = Datasets('tensorflow') # test with label dataset = datasets['dummy_v2'](\ input_shape=(256, 256, 3), label_shape=(1,)) @@ -1131,7 +1131,7 @@ def test_tensorflow_dummy_v2(self): input_shape=(256, 256, 3), dtype=['float32', 'int8']) def test_tensorflow_sparse_dummy_v2(self): - datasets = DATASETS('tensorflow') + datasets = Datasets('tensorflow') # test with label dataset = datasets['sparse_dummy_v2'](\ dense_shape=[[10, 20], [5, 3]], label_shape=[[1]], sparse_ratio=[0.98, 0.8]) @@ -1184,7 +1184,7 @@ def test_style_transfer_dataset(self): im = Image.fromarray(random_array) im.save('test.jpg') - datasets = DATASETS('tensorflow') + datasets = Datasets('tensorflow') dataset = datasets['style_transfer'](content_folder='./', style_folder='./') length = len(dataset) image, label = dataset[0] @@ -1223,7 +1223,7 @@ def test_tensorflow_list_dict(self): # self.assertEqual(data[0][1], 2) def test_pytorch_dummy(self): - datasets = DATASETS('pytorch') + datasets = Datasets('pytorch') transform = TRANSFORMS('pytorch', 'preprocess')['Resize'](**{'size':100}) dataset = datasets['dummy'](shape=[(4, 256, 256, 3), (4, 1)], \ high=[10., 10.], low=[0., 0.], transform=transform) @@ -1240,7 +1240,7 @@ def test_pytorch_dummy(self): @unittest.skipIf(platform.system().lower() == "windows", "not support mxnet on windows yet") def test_mxnet_dummy(self): - datasets = DATASETS('mxnet') + datasets = Datasets('mxnet') transform = TRANSFORMS('mxnet', 'preprocess')['Resize'](**{'size':100}) dataset = datasets['dummy'](shape=(4, 256, 256, 3), transform=transform) @@ -1258,7 +1258,7 @@ def test_mxnet_dummy(self): self.assertEqual(dataset[0][1], 0) def test_onnxrt_qlinear_dummy(self): - datasets = DATASETS('onnxrt_qlinearops') + datasets = Datasets('onnxrt_qlinearops') transform = TRANSFORMS('onnxrt_qlinearops', 'preprocess')['Resize'](**{'size':100}) dataset = datasets['dummy'](shape=(4, 256, 256, 3), transform=transform) @@ -1283,7 +1283,7 @@ def test_onnxrt_qlinear_dummy(self): shape=[(4, 256, 256, 3), (4, 256, 256, 3)], dtype=['float32', 'int8', 'int8']) def test_onnx_integer_dummy(self): - datasets = DATASETS('onnxrt_integerops') + datasets = Datasets('onnxrt_integerops') dataset = datasets['dummy'](shape=(4, 256, 256, 3)) data_loader = DATALOADERS['onnxrt_integerops'](dataset) @@ -1321,7 +1321,7 @@ def test_onnx_bert(self): tsv_w.writerow(['Quality', '#1 ID', '#2 ID', '#1 String', '#2 String']) tsv_w.writerow(['1', '702876', '702977', """Amrozi accused his brother , whom he called " the witness " , of deliberately distorting his evidence .""", """Referring to him as only " the witness " , Amrozi accused his brother of deliberately distorting his evidence ."""]) - datasets = DATASETS('onnxrt_integerops') + datasets = Datasets('onnxrt_integerops') args = {'GLUE': {'data_dir': './MRPC', 'model_name_or_path': 'bert-base-uncased', diff --git a/test/data/test_filter.py b/test/data/test_filter.py index 5b86781de68..830f8e7cf1c 100644 --- a/test/data/test_filter.py +++ b/test/data/test_filter.py @@ -4,7 +4,7 @@ import json import shutil from PIL import Image -from neural_compressor.data import FILTERS, TRANSFORMS, DATASETS, DATALOADERS +from neural_compressor.data import FILTERS, TRANSFORMS, Datasets, DATALOADERS from neural_compressor.utils.create_obj_from_config import create_dataset, get_preprocess, create_dataloader import tensorflow as tf @@ -60,7 +60,7 @@ def testLabelBalanceCOCORecord(self): preprocesses = TRANSFORMS('tensorflow', 'preprocess') filters = FILTERS('tensorflow') filter = filters['LabelBalanceCOCORecord'](2) - datasets = DATASETS('tensorflow') + datasets = Datasets('tensorflow') dataset = datasets['COCORecord']('test.record', \ transform=None, filter=filter) dataloader = DATALOADERS['tensorflow'](dataset=dataset, batch_size=1) @@ -146,7 +146,7 @@ def testLabelBalanceCOCORaw(self): filters = FILTERS('onnxrt_qlinearops') filter = filters['LabelBalanceCOCORaw'](1) - datasets = DATASETS('onnxrt_qlinearops') + datasets = Datasets('onnxrt_qlinearops') dataset = datasets['COCORaw']('./', transform=None, filter=filter) dataloader = DATALOADERS['onnxrt_qlinearops'](dataset=dataset, batch_size=1) for (inputs, labels) in dataloader: diff --git a/test/distillation/test_distillation.py b/test/distillation/test_distillation.py index a43259ed7bb..672adf86d5c 100644 --- a/test/distillation/test_distillation.py +++ b/test/distillation/test_distillation.py @@ -6,7 +6,7 @@ import torchvision import torch.nn as nn import tensorflow as tf -from neural_compressor.data import DATASETS +from neural_compressor.data import Datasets from neural_compressor.config import DistillationConfig, KnowledgeDistillationLossConfig from neural_compressor.experimental.data.dataloaders.pytorch_dataloader import PyTorchDataLoader from neural_compressor.adaptor.tf_utils.util import version1_lt_version2 @@ -216,7 +216,7 @@ def test_distillation_external(self): def test_distillation_external_new_API(self): from neural_compressor.training import prepare_compression - datasets = DATASETS('pytorch') + datasets = Datasets('pytorch') dummy_dataset = datasets['dummy'](shape=(100, 3, 224, 224), low=0., high=1., label=True) dummy_dataloader = PyTorchDataLoader(dummy_dataset) diff --git a/test/distillation/test_self_distillation.py b/test/distillation/test_self_distillation.py index e05a40ae56e..20a695ac211 100644 --- a/test/distillation/test_self_distillation.py +++ b/test/distillation/test_self_distillation.py @@ -5,7 +5,7 @@ import torch import torch.nn as nn import torchvision -from neural_compressor.data import DATASETS +from neural_compressor.data import Datasets from neural_compressor.experimental.data.dataloaders.pytorch_dataloader import \ PyTorchDataLoader @@ -84,7 +84,7 @@ def test_self_distillation(self): from neural_compressor.config import DistillationConfig, \ SelfKnowledgeDistillationLossConfig - datasets = DATASETS("pytorch") + datasets = Datasets("pytorch") dummy_dataset = datasets["dummy"]( shape=(100, 3, 224, 224), low=0.0, high=1.0, label=True ) diff --git a/test/distributed/test_distributed_pt_train.py b/test/distributed/test_distributed_pt_train.py index 672bdbb5ce8..c46d8fa03ec 100644 --- a/test/distributed/test_distributed_pt_train.py +++ b/test/distributed/test_distributed_pt_train.py @@ -8,7 +8,7 @@ import torch.nn as nn import horovod.torch as hvd -from neural_compressor.data import DATASETS +from neural_compressor.data import Datasets from neural_compressor.experimental.data.dataloaders.pytorch_dataloader import PyTorchDataLoader def build_fake_py(): @@ -22,7 +22,7 @@ def build_fake_py(): import torch.nn as nn import horovod.torch as hvd -from neural_compressor.data import DATASETS +from neural_compressor.data import Datasets from neural_compressor.experimental.data.dataloaders.pytorch_dataloader import PyTorchDataLoader diff --git a/test/distributed/test_distributed_tf_dataloader.py b/test/distributed/test_distributed_tf_dataloader.py index 9f4b1d69ba6..88f249ba522 100644 --- a/test/distributed/test_distributed_tf_dataloader.py +++ b/test/distributed/test_distributed_tf_dataloader.py @@ -11,7 +11,7 @@ from neural_compressor import data from neural_compressor.utils.create_obj_from_config import create_dataset, create_dataloader from neural_compressor.data.dataloaders.dataloader import DataLoader -from neural_compressor.data import DATASETS, DATALOADERS, TRANSFORMS +from neural_compressor.data import Datasets, DATALOADERS, TRANSFORMS from neural_compressor.utils import logger from neural_compressor.adaptor.tf_utils.util import version1_lt_version2 diff --git a/test/export/test_torch2onnx.py b/test/export/test_torch2onnx.py index 4cab83cc2b4..21183a7697e 100644 --- a/test/export/test_torch2onnx.py +++ b/test/export/test_torch2onnx.py @@ -7,10 +7,10 @@ from neural_compressor import quantization from neural_compressor.experimental.common import Model from neural_compressor.config import Torch2ONNXConfig -from neural_compressor.experimental.data.datasets.dataset import DATASETS +from neural_compressor.experimental.data.datasets.dataset import Datasets from neural_compressor import PostTrainingQuantConfig, QuantizationAwareTrainingConfig from neural_compressor.training import prepare_compression -from neural_compressor.data import DATASETS, DATALOADERS +from neural_compressor.data import Datasets, DATALOADERS from transformers import AutoModelForSequenceClassification, AutoTokenizer import torch.utils.data as data @@ -79,7 +79,7 @@ class TestPytorch2ONNX(unittest.TestCase): def setUpClass(self): from torchvision.models.quantization import resnet18 self.cv_model = resnet18() - self.cv_dataset = DATASETS("pytorch")["dummy"]((10, 3, 224, 224)) + self.cv_dataset = Datasets("pytorch")["dummy"]((10, 3, 224, 224)) self.cv_dataloader = DATALOADERS["pytorch"](self.cv_dataset) self.nlp_model = AutoModelForSequenceClassification.from_pretrained( "distilbert-base-uncased-finetuned-sst-2-english" diff --git a/test/itex/test_tensorflow_itex_new_api.py b/test/itex/test_tensorflow_itex_new_api.py index 9cb67b25fff..7046a4ff2a6 100644 --- a/test/itex/test_tensorflow_itex_new_api.py +++ b/test/itex/test_tensorflow_itex_new_api.py @@ -58,8 +58,8 @@ def test_itex_new_api(self): quant_format="QDQ", calibration_sampling_size=[200]) - from neural_compressor.data import DATASETS - dataset = DATASETS('tensorflow')['dummy'](shape=(100, 56, 56, 16), label=True) + from neural_compressor.data import Datasets + dataset = Datasets('tensorflow')['dummy'](shape=(100, 56, 56, 16), label=True) output_graph = fit( model=output_graph_def, conf=config, diff --git a/test/metric/test_metrics_2.x.py b/test/metric/test_metrics_2.x.py new file mode 100644 index 00000000000..b638ff47a59 --- /dev/null +++ b/test/metric/test_metrics_2.x.py @@ -0,0 +1,1061 @@ +"""Tests for the metrics module.""" +import numpy as np +import unittest +import platform +from neural_compressor.metric import METRICS +from neural_compressor.metric.f1 import evaluate +from neural_compressor.metric.evaluate_squad import evaluate as evaluate_squad + +class InCorrectMetric: + def __init__(self): + self.item = None + +class CorrectMetric: + def __init__(self): + self.item = [] + + def update(self, samples): + self.item.append(samples) + + def result(self): + return 0 + + def reset(self): + self.item = [] + +class TestMetrics(unittest.TestCase): + def testUserMetric(self): + from neural_compressor.experimental import common, Quantization, Benchmark, \ + Graph_Optimization + for i in [Quantization(), Benchmark(), Graph_Optimization()]: + item = i + with self.assertRaises(AssertionError): + item.metric = InCorrectMetric() + item.framework = 'tensorflow' + item.metric = common.Metric(CorrectMetric, str(i)) + + def testmIOU(self): + metrics = METRICS('tensorflow') + miou = metrics['mIOU']() + preds = np.array([0, 0, 1, 1]) + labels = np.array([0, 1, 0, 1]) + miou.update(preds, labels) + self.assertAlmostEqual(miou.result(), 0.33333334) + + miou.reset() + preds = np.array([0, 0, 1, 1]) + labels = np.array([0, 1, 1, 1]) + miou.update(preds, labels) + self.assertAlmostEqual(miou.result(), 0.58333333) + + def testBLEU(self): + metrics = METRICS('tensorflow') + bleu = metrics['BLEU']() + preds = ['Gutach: Mehr Sicherheit für Fußgänger'] + labels = ('Gutach: Noch mehr Sicherheit für Fußgänger',) + bleu.update(preds, labels) + self.assertAlmostEqual(bleu.result(), 51.1507809) + bleu.reset() + + preds = ['Dies wurde auch von Peter Arnold vom Offenburg District Office bestätigt.'] + labels = ('Dies bestätigt auch Peter Arnold vom Landratsamt Offenburg.',) + bleu.update(preds, labels) + self.assertAlmostEqual(bleu.result(), 16.108992695) + with self.assertRaises(ValueError): + bleu.update(['a','b'], ('c',)) + + def test_onnxrt_GLUE(self): + metrics = METRICS('onnxrt_qlinearops') + glue = metrics['GLUE']('mrpc') + preds = [np.array( + [[-3.2443411, 3.0909934], + [2.0500996, -2.3100944], + [1.870293 , -2.0741048], + [-2.8377204, 2.617834], + [2.008347 , -2.0215416], + [-2.9693947, 2.7782154], + [-2.9949608, 2.7887983], + [-3.0623112, 2.8748074]]) + ] + labels = [np.array([1, 0, 0, 1, 0, 1, 0, 1])] + glue.update(preds, labels) + self.assertEqual(glue.result(), 0.875) + preds_2 = [np.array( + [[-3.1296735, 2.8356276], + [-3.172515 , 2.9173899], + [-3.220131 , 3.0916846], + [2.1452675, -1.9398905], + [1.5475761, -1.9101546], + [-2.9797182, 2.721741], + [-3.2052834, 2.9934788], + [-2.7451005, 2.622343]]) + ] + labels_2 = [np.array([1, 1, 1, 0, 0, 1, 1, 1])] + glue.update(preds_2, labels_2) + self.assertEqual(glue.result(), 0.9375) + + glue.reset() + glue.update(preds, labels) + self.assertEqual(glue.result(), 0.875) + + def test_tensorflow_F1(self): + metrics = METRICS('tensorflow') + F1 = metrics['F1']() + preds = [1, 1, 1, 1] + labels = [0, 1, 1, 0] + + F1.update(preds, labels) + self.assertEqual(F1.result(), 0.5) + + def test_squad_evaluate(self): + label = [{'paragraphs':\ + [{'qas':[{'answers': [{'answer_start': 177, 'text': 'Denver Broncos'}, \ + {'answer_start': 177, 'text': 'Denver Broncos'}, \ + {'answer_start': 177, 'text': 'Denver Broncos'}], \ + 'question': 'Which NFL team represented the AFC at Super Bowl 50?', \ + 'id': '56be4db0acb8001400a502ec'}]}]}] + preds = {'56be4db0acb8001400a502ec': 'Denver Broncos'} + f1 = evaluate(preds, label) + self.assertEqual(f1, 100.) + dataset = [{'paragraphs':\ + [{'qas':[{'answers': [{'answer_start': 177, 'text': 'Denver Broncos'}, \ + {'answer_start': 177, 'text': 'Denver Broncos'}, \ + {'answer_start': 177, 'text': 'Denver Broncos'}], \ + 'question': 'Which NFL team represented the AFC at Super Bowl 50?', \ + 'id': '56be4db0acb8001400a502ec'}]}]}] + predictions = {'56be4db0acb8001400a502ec': 'Denver Broncos'} + f1_squad = evaluate_squad(dataset,predictions) + self.assertEqual(f1_squad['f1'], 100.) + self.assertEqual(f1_squad['exact_match'], 100.) + + + def test_pytorch_F1(self): + metrics = METRICS('pytorch') + F1 = metrics['F1']() + F1.reset() + preds = [1, 1] + labels = [2, 1, 1] + + F1.update(preds, labels) + self.assertEqual(F1.result(), 0.8) + + @unittest.skipIf(platform.system().lower() == "windows", "not support mxnet on windows yet") + def test_mxnet_F1(self): + metrics = METRICS('mxnet') + F1 = metrics['F1']() + preds = [0, 1, 1, 1, 1, 0] + labels = [0, 1, 1, 1] + + F1.update(preds, labels) + self.assertEqual(F1.result(), 0.8) + + def test_onnx_topk(self): + metrics = METRICS('onnxrt_qlinearops') + top1 = metrics['topk']() + top1.reset() + self.assertEqual(top1.result(), 0) + self.assertEqual(top1.result(), 0) + top2 = metrics['topk'](k=2) + top3 = metrics['topk'](k=3) + + predicts = [[0, 0.2, 0.9, 0.3], [0, 0.9, 0.8, 0]] + single_predict = [0, 0.2, 0.9, 0.3] + + labels = [[0, 1, 0, 0], [0, 0, 1, 0]] + sparse_labels = [2, 2] + single_label = 2 + + # test functionality of one-hot label + top1.update(predicts, labels) + top2.update(predicts, labels) + top3.update(predicts, labels) + self.assertEqual(top1.result(), 0.0) + self.assertEqual(top2.result(), 0.5) + self.assertEqual(top3.result(), 1) + + # test functionality of sparse label + top1.update(predicts, sparse_labels) + top2.update(predicts, sparse_labels) + top3.update(predicts, sparse_labels) + self.assertEqual(top1.result(), 0.25) + self.assertEqual(top2.result(), 0.75) + self.assertEqual(top3.result(), 1) + + # test functionality of single label + top1.update(single_predict, single_label) + top2.update(single_predict, single_label) + top3.update(single_predict, single_label) + self.assertEqual(top1.result(), 0.4) + self.assertEqual(top2.result(), 0.8) + self.assertEqual(top3.result(), 1) + + @unittest.skipIf(platform.system().lower() == "windows", "not support mxnet on windows yet") + def test_mxnet_topk(self): + metrics = METRICS('mxnet') + top1 = metrics['topk']() + top1.reset() + self.assertEqual(top1.result(), 0) + top2 = metrics['topk'](k=2) + top3 = metrics['topk'](k=3) + + predicts = [[0, 0.2, 0.9, 0.3], [0, 0.9, 0.8, 0]] + single_predict = [0, 0.2, 0.9, 0.3] + + labels = [[0, 1, 0, 0], [0, 0, 1, 0]] + sparse_labels = [2, 2] + single_label = 2 + + # test functionality of one-hot label + top1.update(predicts, labels) + top2.update(predicts, labels) + top3.update(predicts, labels) + self.assertEqual(top1.result(), 0.0) + self.assertEqual(top2.result(), 0.5) + self.assertEqual(top3.result(), 1) + + # test functionality of sparse label + top1.update(predicts, sparse_labels) + top2.update(predicts, sparse_labels) + top3.update(predicts, sparse_labels) + self.assertEqual(top1.result(), 0.25) + self.assertEqual(top2.result(), 0.75) + self.assertEqual(top3.result(), 1) + + # test functionality of single label + top1.update(single_predict, single_label) + top2.update(single_predict, single_label) + top3.update(single_predict, single_label) + self.assertEqual(top1.result(), 0.4) + self.assertEqual(top2.result(), 0.8) + self.assertEqual(top3.result(), 1) + + def test_tensorflow_topk(self): + metrics = METRICS('tensorflow') + top1 = metrics['topk']() + top1.reset() + self.assertEqual(top1.result(), 0) + top2 = metrics['topk'](k=2) + top3 = metrics['topk'](k=3) + + predicts = [[0, 0.2, 0.9, 0.3], [0, 0.9, 0.8, 0]] + single_predict = [0, 0.2, 0.9, 0.3] + + labels = [[0, 1, 0, 0], [0, 0, 1, 0]] + sparse_labels = [2, 2] + single_label = 2 + + # test functionality of one-hot label + top1.update(predicts, labels) + top2.update(predicts, labels) + top3.update(predicts, labels) + self.assertEqual(top1.result(), 0.0) + self.assertEqual(top2.result(), 0.5) + self.assertEqual(top3.result(), 1) + + # test functionality of sparse label + top1.update(predicts, sparse_labels) + top2.update(predicts, sparse_labels) + top3.update(predicts, sparse_labels) + self.assertEqual(top1.result(), 0.25) + self.assertEqual(top2.result(), 0.75) + self.assertEqual(top3.result(), 1) + + # test functionality of single label + top1.update(single_predict, single_label) + top2.update(single_predict, single_label) + top3.update(single_predict, single_label) + self.assertEqual(top1.result(), 0.4) + self.assertEqual(top2.result(), 0.8) + self.assertEqual(top3.result(), 1) + + def test_tensorflow_mAP(self): + import json + import os + metrics = METRICS('tensorflow') + fake_dict = 'dog: 1' + with open('anno.yaml', 'w', encoding="utf-8") as f: + f.write(fake_dict) + mAP = metrics['mAP']('anno.yaml') + self.assertEqual(mAP.category_map_reverse['dog'], 1) + detection = [ + np.array([[5]]), + np.array([[5]]), + np.array([[[0.16117382, 0.59801614, 0.81511605, 0.7858219 ], + [0.5589304 , 0. , 0.98301625, 0.520178 ], + [0.62706745, 0.35748824, 0.6892729 , 0.41513762], + [0.40032804, 0.01218696, 0.6924763 , 0.30341768], + [0.62706745, 0.35748824, 0.6892729 , 0.41513762]]]), + np.array([[0.9267181 , 0.8510787 , 0.60418576, 0.35155892, 0.31158054]]), + np.array([[ 1., 67., 51., 79., 47.]]) + ] + ground_truth = [ + np.array([[[0.5633255 , 0.34003124, 0.69857144, 0.4009531 ], + [0.4763466 , 0.7769531 , 0.54334897, 0.9675937 ]]]), + np.array([['a', 'b']]), + np.array([[]]), + np.array([b'000000397133.jpg']) + ] + self.assertRaises(ValueError, mAP.update, detection, ground_truth) + + detection = [ + np.array([[[0.16117382, 0.59801614, 0.81511605, 0.7858219 ], + [0.62706745, 0.35748824, 0.6892729 , 0.41513762]]]), + np.array([[0.9267181 , 0.8510787]]), + np.array([[ 1., 1.]]) + ] + ground_truth = [ + np.array([[[0.16117382, 0.59801614, 0.81511605, 0.7858219 ], + [0.62706745, 0.35748824, 0.6892729 , 0.41513762]]]), + np.array([[b'dog', b'dog']]), + np.array([[]]), + np.array([b'000000397133.jpg']) + ] + mAP.update(detection, ground_truth) + mAP.result() + self.assertEqual(format(mAP.result(), '.5f'), + '1.00000') + + detection = [ + np.array([[[0.16117382, 0.59801614, 0.81511605, 0.7858219 ], + [0.5589304 , 0. , 0.98301625, 0.520178 ], + [0.62706745, 0.35748824, 0.6892729 , 0.41513762], + [0.40032804, 0.01218696, 0.6924763 , 0.30341768], + [0.62706745, 0.35748824, 0.6892729 , 0.41513762]]]), + np.array([[0.9267181 , 0.8510787 , 0.60418576, 0.35155892, 0.31158054]]), + np.array([[ 1., 67., 51., 79., 47.]]) + ] + detection_2 = [ + np.array([[8]]), + np.array([[[0.82776225, 0.5865939 , 0.8927653 , 0.6302338 ], + [0.8375764 , 0.6424138 , 0.9055594 , 0.6921875 ], + [0.57902956, 0.39394334, 0.8342961 , 0.5577197 ], + [0.7949219 , 0.6513021 , 0.8472295 , 0.68427753], + [0.809729 , 0.5947042 , 0.8539927 , 0.62916476], + [0.7258591 , 0.08907133, 1. , 0.86224866], + [0.43100086, 0.37782395, 0.8384069 , 0.5616918 ], + [0.32005906, 0.84334356, 1. , 1. ]]]), + np.array([[0.86698544, 0.7562499 , 0.66414887, 0.64498234,\ + 0.63083494,0.46618757, 0.3914739 , 0.3094324 ]]), + np.array([[55., 55., 79., 55., 55., 67., 79., 82.]]) + ] + ground_truth = [ + np.array([[[0.5633255 , 0.34003124, 0.69857144, 0.4009531 ], + [0.56262296, 0.0015625 , 1. , 0.5431719 ], + [0.16374707, 0.60728127, 0.813911 , 0.77823436], + [0.5841452 , 0.21182813, 0.65156907, 0.24670312], + [0.8056206 , 0.048875 , 0.90124124, 0.1553125 ], + [0.6729742 , 0.09317187, 0.7696956 , 0.21203125], + [0.3848478 , 0.002125 , 0.61522245, 0.303 ], + [0.61548007, 0. , 0.7015925 , 0.097125 ], + [0.6381967 , 0.1865625 , 0.7184075 , 0.22534375], + [0.6274239 , 0.22104688, 0.71140516, 0.27134374], + [0.39566743, 0.24370313, 0.43578455, 0.284375 ], + [0.2673302 , 0.245625 , 0.3043794 , 0.27353126], + [0.7137705 , 0.15429688, 0.726815 , 0.17114063], + [0.6003747 , 0.25942189, 0.6438876 , 0.27320313], + [0.68845433, 0.13501562, 0.714637 , 0.17245312], + [0.69358313, 0.10959375, 0.7043091 , 0.12409375], + [0.493911 , 0. , 0.72571427, 0.299 ], + [0.69576114, 0.15107812, 0.70714283, 0.16332813], + [0.4763466 , 0.7769531 , 0.54334897, 0.9675937 ]]]), + np.array([[]]), + np.array([[44, 67, 1, 49, 51, 51, 79, 1, 47, 47, 51, 51,\ + 56, 50, 56, 56, 79, 57, 81]]), + np.array([b'000000397133.jpg']) + ] + ground_truth_2 = [ + np.array([[[0.51508695, 0.2911648 , 0.5903478 , 0.31360796], + [0.9358696 , 0.07528409, 0.99891305, 0.25 ], + [0.8242174 , 0.3309659 , 0.93508697, 0.47301137], + [0.77413046, 0.22599432, 0.9858696 , 0.8179261 ], + [0.32582608, 0.8575 , 0.98426086, 0.9984659 ], + [0.77795655, 0.6268466 , 0.89930433, 0.73434657], + [0.5396087 , 0.39053977, 0.8483913 , 0.5615057 ], + [0.58473915, 0.75661933, 0.5998261 , 0.83579546], + [0.80391306, 0.6129829 , 0.8733478 , 0.66201705], + [0.8737391 , 0.6579546 , 0.943 , 0.7053693 ], + [0.775 , 0.6549716 , 0.8227391 , 0.6882955 ], + [0.8130869 , 0.58292615, 0.90526086, 0.62551135], + [0.7844348 , 0.68735796, 0.98182607, 0.83329546], + [0.872 , 0.6190057 , 0.9306522 , 0.6591761 ]]]), + np.array([[]]), + np.array([[64, 62, 62, 67, 82, 52, 79, 81, 55, 55, 55, 55, 62, 55]]), + np.array([b'000000037777.jpg']) + ] + + mAP = metrics['mAP']() + + self.assertEqual(mAP.result(), 0) + + mAP.update(detection, ground_truth) + + mAP.update(detection, ground_truth) + self.assertEqual(format(mAP.result(), '.5f'), + '0.18182') + + mAP.update(detection_2, ground_truth_2) + self.assertEqual(format(mAP.result(), '.5f'), + '0.20347') + mAP.reset() + mAP.update(detection, ground_truth) + self.assertEqual(format(mAP.result(), '.5f'), + '0.18182') + + ground_truth_1 = [ + np.array([[[0.51508695, 0.2911648 , 0.5903478 , 0.31360796], + [0.872 , 0.6190057 , 0.9306522 , 0.6591761 ]]]), + np.array([[]]), + np.array([[[64, 62]]]), + np.array([b'000000037777.jpg']) + ] + self.assertRaises(ValueError, mAP.update, detection, ground_truth_1) + ground_truth_2 = [ + np.array([[[0.51508695, 0.2911648 , 0.5903478 , 0.31360796], + [0.872 , 0.6190057 , 0.9306522 , 0.6591761 ]]]), + np.array([[]]), + np.array([[64]]), + np.array([b'000000037700.jpg']) + ] + self.assertRaises(ValueError, mAP.update, detection, ground_truth_2) + detection_1 = [ + np.array([[[0.16117382, 0.59801614, 0.81511605, 0.7858219 ], + [0.5589304 , 0. , 0.98301625, 0.520178 ]]]), + np.array([[0.9267181 , 0.8510787 , 0.60418576, 0.35155892, 0.31158054]]), + np.array([[ 1., 67., 51., 79., 47.]]) + ] + ground_truth_1 = [ + np.array([[[0.51508695, 0.2911648 , 0.5903478 , 0.31360796], + [0.872 , 0.6190057 , 0.9306522 , 0.6591761 ]]]), + np.array([[]]), + np.array([[64, 62]]), + np.array([b'000000011.jpg']) + ] + self.assertRaises(ValueError, mAP.update, detection_1, ground_truth_1) + ground_truth_2 = [ + np.array([[[0.51508695, 0.2911648 , 0.5903478 , 0.31360796], + [0.872 , 0.6190057 , 0.9306522 , 0.6591761 ]]]), + np.array([[]]), + np.array([[64, 62]]), + np.array([b'000000012.jpg']) + ] + detection_2 = [ + np.array([[[0.16117382, 0.59801614, 0.81511605, 0.7858219 ], + [0.5589304 , 0. , 0.98301625, 0.520178 ]]]), + np.array([[0.9267181 , 0.8510787]]), + np.array([[ 1., 67., 51., 79., 47.]]) + ] + self.assertRaises(ValueError, mAP.update, detection_2, ground_truth_2) + os.remove('anno.yaml') + + + def test_tensorflow_VOCmAP(self): + import os + metrics = METRICS('tensorflow') + fake_dict = 'dog: 1' + with open('anno.yaml', 'w', encoding="utf-8") as f: + f.write(fake_dict) + mAP = metrics['VOCmAP']('anno.yaml') + self.assertEqual(mAP.iou_thrs, 0.5) + self.assertEqual(mAP.map_points, 0) + self.assertEqual(mAP.category_map_reverse['dog'], 1) + detection = [ + np.array([[5]]), + np.array([[5]]), + np.array([[[0.16117382, 0.59801614, 0.81511605, 0.7858219 ], + [0.5589304 , 0. , 0.98301625, 0.520178 ], + [0.62706745, 0.35748824, 0.6892729 , 0.41513762], + [0.40032804, 0.01218696, 0.6924763 , 0.30341768], + [0.62706745, 0.35748824, 0.6892729 , 0.41513762]]]), + np.array([[0.9267181 , 0.8510787 , 0.60418576, 0.35155892, 0.31158054]]), + np.array([[ 1., 67., 51., 79., 47.]]) + ] + ground_truth = [ + np.array([[[0.5633255 , 0.34003124, 0.69857144, 0.4009531 ], + [0.4763466 , 0.7769531 , 0.54334897, 0.9675937 ]]]), + np.array([['a', 'b']]), + np.array([[]]), + np.array([b'000000397133.jpg']) + ] + self.assertRaises(ValueError, mAP.update, detection, ground_truth) + + os.remove('anno.yaml') + + mAP = metrics['VOCmAP']() + detection = [ + np.array([[[0.16117382, 0.59801614, 0.81511605, 0.7858219 ], + [0.5589304 , 0. , 0.98301625, 0.520178 ], + [0.62706745, 0.35748824, 0.6892729 , 0.41513762], + [0.40032804, 0.01218696, 0.6924763 , 0.30341768], + [0.62706745, 0.35748824, 0.6892729 , 0.41513762]]]), + np.array([[0.9267181 , 0.8510787 , 0.60418576, 0.35155892, 0.31158054]]), + np.array([[ 1., 67., 51., 79., 47.]]) + ] + detection_2 = [ + np.array([[8]]), + np.array([[[0.82776225, 0.5865939 , 0.8927653 , 0.6302338 ], + [0.8375764 , 0.6424138 , 0.9055594 , 0.6921875 ], + [0.57902956, 0.39394334, 0.8342961 , 0.5577197 ], + [0.7949219 , 0.6513021 , 0.8472295 , 0.68427753], + [0.809729 , 0.5947042 , 0.8539927 , 0.62916476], + [0.7258591 , 0.08907133, 1. , 0.86224866], + [0.43100086, 0.37782395, 0.8384069 , 0.5616918 ], + [0.32005906, 0.84334356, 1. , 1. ]]]), + np.array([[0.86698544, 0.7562499 , 0.66414887, 0.64498234,\ + 0.63083494,0.46618757, 0.3914739 , 0.3094324 ]]), + np.array([[55., 55., 79., 55., 55., 67., 79., 82.]]) + ] + ground_truth = [ + np.array([[[0.5633255 , 0.34003124, 0.69857144, 0.4009531 ], + [0.56262296, 0.0015625 , 1. , 0.5431719 ], + [0.16374707, 0.60728127, 0.813911 , 0.77823436], + [0.5841452 , 0.21182813, 0.65156907, 0.24670312], + [0.8056206 , 0.048875 , 0.90124124, 0.1553125 ], + [0.6729742 , 0.09317187, 0.7696956 , 0.21203125], + [0.3848478 , 0.002125 , 0.61522245, 0.303 ], + [0.61548007, 0. , 0.7015925 , 0.097125 ], + [0.6381967 , 0.1865625 , 0.7184075 , 0.22534375], + [0.6274239 , 0.22104688, 0.71140516, 0.27134374], + [0.39566743, 0.24370313, 0.43578455, 0.284375 ], + [0.2673302 , 0.245625 , 0.3043794 , 0.27353126], + [0.7137705 , 0.15429688, 0.726815 , 0.17114063], + [0.6003747 , 0.25942189, 0.6438876 , 0.27320313], + [0.68845433, 0.13501562, 0.714637 , 0.17245312], + [0.69358313, 0.10959375, 0.7043091 , 0.12409375], + [0.493911 , 0. , 0.72571427, 0.299 ], + [0.69576114, 0.15107812, 0.70714283, 0.16332813], + [0.4763466 , 0.7769531 , 0.54334897, 0.9675937 ]]]), + np.array([[]]), + np.array([[44, 67, 1, 49, 51, 51, 79, 1, 47, 47, 51, 51,\ + 56, 50, 56, 56, 79, 57, 81]]), + np.array([b'000000397133.jpg']) + ] + ground_truth_2 = [ + np.array([[[0.51508695, 0.2911648 , 0.5903478 , 0.31360796], + [0.9358696 , 0.07528409, 0.99891305, 0.25 ], + [0.8242174 , 0.3309659 , 0.93508697, 0.47301137], + [0.77413046, 0.22599432, 0.9858696 , 0.8179261 ], + [0.32582608, 0.8575 , 0.98426086, 0.9984659 ], + [0.77795655, 0.6268466 , 0.89930433, 0.73434657], + [0.5396087 , 0.39053977, 0.8483913 , 0.5615057 ], + [0.58473915, 0.75661933, 0.5998261 , 0.83579546], + [0.80391306, 0.6129829 , 0.8733478 , 0.66201705], + [0.8737391 , 0.6579546 , 0.943 , 0.7053693 ], + [0.775 , 0.6549716 , 0.8227391 , 0.6882955 ], + [0.8130869 , 0.58292615, 0.90526086, 0.62551135], + [0.7844348 , 0.68735796, 0.98182607, 0.83329546], + [0.872 , 0.6190057 , 0.9306522 , 0.6591761 ]]]), + np.array([[]]), + np.array([[64, 62, 62, 67, 82, 52, 79, 81, 55, 55, 55, 55, 62, 55]]), + np.array([b'000000037777.jpg']) + ] + + self.assertEqual(mAP.result(), 0) + + mAP.update(detection, ground_truth) + + mAP.update(detection, ground_truth) + self.assertEqual(format(mAP.result(), '.5f'), + '0.18182') + + mAP.update(detection_2, ground_truth_2) + self.assertEqual(format(mAP.result(), '.5f'), + '0.20347') + mAP.reset() + mAP.update(detection, ground_truth) + self.assertEqual(format(mAP.result(), '.5f'), + '0.18182') + + ground_truth_1 = [ + np.array([[[0.51508695, 0.2911648 , 0.5903478 , 0.31360796], + [0.872 , 0.6190057 , 0.9306522 , 0.6591761 ]]]), + np.array([[]]), + np.array([[[64, 62]]]), + np.array([b'000000037777.jpg']) + ] + self.assertRaises(ValueError, mAP.update, detection, ground_truth_1) + ground_truth_2 = [ + np.array([[[0.51508695, 0.2911648 , 0.5903478 , 0.31360796], + [0.872 , 0.6190057 , 0.9306522 , 0.6591761 ]]]), + np.array([[]]), + np.array([[64]]), + np.array([b'000000037700.jpg']) + ] + self.assertRaises(ValueError, mAP.update, detection, ground_truth_2) + detection_1 = [ + np.array([[[0.16117382, 0.59801614, 0.81511605, 0.7858219 ], + [0.5589304 , 0. , 0.98301625, 0.520178 ]]]), + np.array([[0.9267181 , 0.8510787 , 0.60418576, 0.35155892, 0.31158054]]), + np.array([[ 1., 67., 51., 79., 47.]]) + ] + ground_truth_1 = [ + np.array([[[0.51508695, 0.2911648 , 0.5903478 , 0.31360796], + [0.872 , 0.6190057 , 0.9306522 , 0.6591761 ]]]), + np.array([[]]), + np.array([[64, 62]]), + np.array([b'000000011.jpg']) + ] + self.assertRaises(ValueError, mAP.update, detection_1, ground_truth_1) + ground_truth_2 = [ + np.array([[[0.51508695, 0.2911648 , 0.5903478 , 0.31360796], + [0.872 , 0.6190057 , 0.9306522 , 0.6591761 ]]]), + np.array([[]]), + np.array([[64, 62]]), + np.array([b'000000012.jpg']) + ] + detection_2 = [ + np.array([[[0.16117382, 0.59801614, 0.81511605, 0.7858219 ], + [0.5589304 , 0. , 0.98301625, 0.520178 ]]]), + np.array([[0.9267181 , 0.8510787]]), + np.array([[ 1., 67., 51., 79., 47.]]) + ] + self.assertRaises(ValueError, mAP.update, detection_2, ground_truth_2) + + + def test_tensorflow_COCOmAP(self): + import os + output_index_mapping = {'num_detections':0, 'boxes':1, 'scores':2, 'classes':3} + metrics = METRICS('tensorflow') + fake_dict = 'dog: 1' + with open('anno.yaml', 'w', encoding="utf-8") as f: + f.write(fake_dict) + mAP = metrics['COCOmAP']('anno.yaml') + mAP2 = metrics['COCOmAPv2']('anno.yaml', output_index_mapping=output_index_mapping) + self.assertEqual(mAP.category_map_reverse['dog'], 1) + self.assertEqual(mAP2.category_map_reverse['dog'], 1) + detection = [ + np.array([[5]]), + np.array([[5]]), + np.array([[[0.16117382, 0.59801614, 0.81511605, 0.7858219 ], + [0.5589304 , 0. , 0.98301625, 0.520178 ], + [0.62706745, 0.35748824, 0.6892729 , 0.41513762], + [0.40032804, 0.01218696, 0.6924763 , 0.30341768], + [0.62706745, 0.35748824, 0.6892729 , 0.41513762]]]), + np.array([[0.9267181 , 0.8510787 , 0.60418576, 0.35155892, 0.31158054]]), + np.array([[ 1., 67., 51., 79., 47.]]) + ] + ground_truth = [ + np.array([[[0.5633255 , 0.34003124, 0.69857144, 0.4009531 ], + [0.4763466 , 0.7769531 , 0.54334897, 0.9675937 ]]]), + np.array([['a', 'b']]), + np.array([[]]), + np.array([b'000000397133.jpg']) + ] + self.assertRaises(ValueError, mAP.update, detection, ground_truth) + + os.remove('anno.yaml') + + mAP = metrics['COCOmAP']() + mAP2 = metrics['COCOmAPv2']() + detection = [ + np.array([[[0.16117382, 0.59801614, 0.81511605, 0.7858219 ], + [0.5589304 , 0. , 0.98301625, 0.520178 ], + [0.62706745, 0.35748824, 0.6892729 , 0.41513762], + [0.40032804, 0.01218696, 0.6924763 , 0.30341768], + [0.62706745, 0.35748824, 0.6892729 , 0.41513762]]]), + np.array([[0.9267181 , 0.8510787 , 0.60418576, 0.35155892, 0.31158054]]), + np.array([[ 1., 67., 51., 79., 47.]]) + ] + detection_2 = [ + np.array([[8]]), + np.array([[[0.82776225, 0.5865939 , 0.8927653 , 0.6302338 ], + [0.8375764 , 0.6424138 , 0.9055594 , 0.6921875 ], + [0.57902956, 0.39394334, 0.8342961 , 0.5577197 ], + [0.7949219 , 0.6513021 , 0.8472295 , 0.68427753], + [0.809729 , 0.5947042 , 0.8539927 , 0.62916476], + [0.7258591 , 0.08907133, 1. , 0.86224866], + [0.43100086, 0.37782395, 0.8384069 , 0.5616918 ], + [0.32005906, 0.84334356, 1. , 1. ]]]), + np.array([[0.86698544, 0.7562499 , 0.66414887, 0.64498234,\ + 0.63083494,0.46618757, 0.3914739 , 0.3094324 ]]), + np.array([[55., 55., 79., 55., 55., 67., 79., 82.]]) + ] + ground_truth = [ + np.array([[[0.5633255 , 0.34003124, 0.69857144, 0.4009531 ], + [0.56262296, 0.0015625 , 1. , 0.5431719 ], + [0.16374707, 0.60728127, 0.813911 , 0.77823436], + [0.5841452 , 0.21182813, 0.65156907, 0.24670312], + [0.8056206 , 0.048875 , 0.90124124, 0.1553125 ], + [0.6729742 , 0.09317187, 0.7696956 , 0.21203125], + [0.3848478 , 0.002125 , 0.61522245, 0.303 ], + [0.61548007, 0. , 0.7015925 , 0.097125 ], + [0.6381967 , 0.1865625 , 0.7184075 , 0.22534375], + [0.6274239 , 0.22104688, 0.71140516, 0.27134374], + [0.39566743, 0.24370313, 0.43578455, 0.284375 ], + [0.2673302 , 0.245625 , 0.3043794 , 0.27353126], + [0.7137705 , 0.15429688, 0.726815 , 0.17114063], + [0.6003747 , 0.25942189, 0.6438876 , 0.27320313], + [0.68845433, 0.13501562, 0.714637 , 0.17245312], + [0.69358313, 0.10959375, 0.7043091 , 0.12409375], + [0.493911 , 0. , 0.72571427, 0.299 ], + [0.69576114, 0.15107812, 0.70714283, 0.16332813], + [0.4763466 , 0.7769531 , 0.54334897, 0.9675937 ]]]), + np.array([[]]), + np.array([[44, 67, 1, 49, 51, 51, 79, 1, 47, 47, 51, 51,\ + 56, 50, 56, 56, 79, 57, 81]]), + np.array([b'000000397133.jpg']) + ] + ground_truth_2 = [ + np.array([[[0.51508695, 0.2911648 , 0.5903478 , 0.31360796], + [0.9358696 , 0.07528409, 0.99891305, 0.25 ], + [0.8242174 , 0.3309659 , 0.93508697, 0.47301137], + [0.77413046, 0.22599432, 0.9858696 , 0.8179261 ], + [0.32582608, 0.8575 , 0.98426086, 0.9984659 ], + [0.77795655, 0.6268466 , 0.89930433, 0.73434657], + [0.5396087 , 0.39053977, 0.8483913 , 0.5615057 ], + [0.58473915, 0.75661933, 0.5998261 , 0.83579546], + [0.80391306, 0.6129829 , 0.8733478 , 0.66201705], + [0.8737391 , 0.6579546 , 0.943 , 0.7053693 ], + [0.775 , 0.6549716 , 0.8227391 , 0.6882955 ], + [0.8130869 , 0.58292615, 0.90526086, 0.62551135], + [0.7844348 , 0.68735796, 0.98182607, 0.83329546], + [0.872 , 0.6190057 , 0.9306522 , 0.6591761 ]]]), + np.array([[]]), + np.array([[64, 62, 62, 67, 82, 52, 79, 81, 55, 55, 55, 55, 62, 55]]), + np.array([b'000000037777.jpg']) + ] + + self.assertEqual(mAP.result(), 0) + self.assertEqual(mAP2.result(), 0) + + mAP.update(detection, ground_truth) + + mAP.update(detection, ground_truth) + self.assertEqual(format(mAP.result(), '.5f'), + '0.14149') + + mAP.update(detection_2, ground_truth_2) + self.assertEqual(format(mAP.result(), '.5f'), + '0.13366') + mAP.reset() + mAP.update(detection, ground_truth) + self.assertEqual(format(mAP.result(), '.5f'), + '0.14149') + + mAP2.update(detection, ground_truth) + + mAP2.update(detection, ground_truth) + self.assertEqual(format(mAP2.result(), '.5f'), + '0.14149') + + mAP2 = metrics['COCOmAPv2'](output_index_mapping=output_index_mapping) + + mAP2.update(detection_2, ground_truth_2) + self.assertEqual(format(mAP2.result(), '.5f'), + '0.20520') + mAP2.reset() + mAP2.update(detection_2, ground_truth_2) + self.assertEqual(format(mAP2.result(), '.5f'), + '0.20520') + + mAP2 = metrics['COCOmAPv2']() + + ground_truth_1 = [ + np.array([[[0.51508695, 0.2911648 , 0.5903478 , 0.31360796], + [0.872 , 0.6190057 , 0.9306522 , 0.6591761 ]]]), + np.array([[]]), + np.array([[[64, 62]]]), + np.array([b'000000037777.jpg']) + ] + self.assertRaises(ValueError, mAP.update, detection, ground_truth_1) + self.assertRaises(ValueError, mAP2.update, detection, ground_truth_1) + + ground_truth_2 = [ + np.array([[[0.51508695, 0.2911648 , 0.5903478 , 0.31360796], + [0.872 , 0.6190057 , 0.9306522 , 0.6591761 ]]]), + np.array([[]]), + np.array([[64]]), + np.array([b'000000037700.jpg']) + ] + self.assertRaises(ValueError, mAP.update, detection, ground_truth_2) + self.assertRaises(ValueError, mAP2.update, detection, ground_truth_2) + + detection_1 = [ + np.array([[[0.16117382, 0.59801614, 0.81511605, 0.7858219 ], + [0.5589304 , 0. , 0.98301625, 0.520178 ]]]), + np.array([[0.9267181 , 0.8510787 , 0.60418576, 0.35155892, 0.31158054]]), + np.array([[ 1., 67., 51., 79., 47.]]) + ] + ground_truth_1 = [ + np.array([[[0.51508695, 0.2911648 , 0.5903478 , 0.31360796], + [0.872 , 0.6190057 , 0.9306522 , 0.6591761 ]]]), + np.array([[]]), + np.array([[64, 62]]), + np.array([b'000000011.jpg']) + ] + self.assertRaises(ValueError, mAP.update, detection_1, ground_truth_1) + self.assertRaises(ValueError, mAP2.update, detection_1, ground_truth_1) + + ground_truth_2 = [ + np.array([[[0.51508695, 0.2911648 , 0.5903478 , 0.31360796], + [0.872 , 0.6190057 , 0.9306522 , 0.6591761 ]]]), + np.array([[]]), + np.array([[64, 62]]), + np.array([b'000000012.jpg']) + ] + detection_2 = [ + np.array([[[0.16117382, 0.59801614, 0.81511605, 0.7858219 ], + [0.5589304 , 0. , 0.98301625, 0.520178 ]]]), + np.array([[0.9267181 , 0.8510787]]), + np.array([[ 1., 67., 51., 79., 47.]]) + ] + self.assertRaises(ValueError, mAP.update, detection_2, ground_truth_2) + self.assertRaises(ValueError, mAP2.update, detection_2, ground_truth_2) + + @unittest.skipIf(platform.system().lower() == "windows", "not support mxnet on windows now") + def test__accuracy(self): + predicts1 = [1, 0, 1, 1] + labels1 = [0, 1, 1, 1] + + predicts2 = [[0, 0], [0, 0]] + labels2 = [[0, 1], [1, 1]] + + predicts3 = [[[0, 1], [0, 0], [0, 1]], [[0, 1], [0, 1], [0, 1]]] + labels3 = [[[0, 1], [0, 1], [1, 0]], [[1, 0], [1, 0], [1, 0]]] + + predicts4 = [[0.2, 0.8], [0.1, 0.9], [0.3, 0.7], [0.4, 0.6]] #1,1,1,1 + labels4 = [0, 1, 0, 0] + + metrics = METRICS('pytorch') + acc = metrics['Accuracy']() + acc.update(predicts1, labels1) + acc_result = acc.result() + self.assertEqual(acc_result, 0.5) + acc.reset() + acc.update(predicts2, labels2) + self.assertEqual(acc.result(), 0.25) + acc.reset() + acc.update(predicts3, labels3) + self.assertEqual(acc.result(), 0.25) + acc.reset() + acc.update(predicts4, labels4) + self.assertEqual(acc.result(), 0.25) + + metrics = METRICS('mxnet') + acc = metrics['Accuracy']() + acc.update(predicts1, labels1) + acc_result = acc.result() + self.assertEqual(acc_result, 0.5) + acc.reset() + acc.update(predicts2, labels2) + self.assertEqual(acc.result(), 0.25) + acc.reset() + acc.update(predicts3, labels3) + self.assertEqual(acc.result(), 0.25) + acc.reset() + acc.update(predicts4, labels4) + self.assertEqual(acc.result(), 0.25) + + metrics = METRICS('onnxrt_qlinearops') + acc = metrics['Accuracy']() + acc.update(predicts1, labels1) + acc_result = acc.result() + self.assertEqual(acc_result, 0.5) + acc.reset() + acc.update(predicts2, labels2) + self.assertEqual(acc.result(), 0.25) + acc.reset() + acc.update(predicts3, labels3) + self.assertEqual(acc.result(), 0.25) + acc.reset() + acc.update(predicts4, labels4) + self.assertEqual(acc.result(), 0.25) + + acc.reset() + acc.update(1, 1) + self.assertEqual(acc.result(), 1.0) + + wrong_predictions = [1, 0, 0] + wrong_labels = [[0, 1, 1]] + self.assertRaises(ValueError, acc.update, wrong_predictions, wrong_labels) + + + @unittest.skipIf(platform.system().lower() == "windows", "not support mxnet on windows yet") + def test_mxnet_accuracy(self): + metrics = METRICS('mxnet') + acc = metrics['Accuracy']() + predicts = [1, 0, 1, 1] + labels = [0, 1, 1, 1] + acc.update(predicts, labels) + acc_result = acc.result() + self.assertEqual(acc_result, 0.5) + + @unittest.skipIf(platform.system().lower() == "windows", "not support mxnet on windows now") + def test_mse(self): + predicts1 = [1, 0, 0, 1] + labels1 = [0, 1, 0, 0] + predicts2 = [1, 1, 1, 1] + labels2 = [0, 1, 1, 0] + + metrics = METRICS('onnxrt_qlinearops') + mse = metrics['MSE'](compare_label=False) + mse.update(predicts1, labels1) + mse_result = mse.result() + self.assertEqual(mse_result, 0.75) + mse.update(predicts2, labels2) + mse_result = mse.result() + self.assertEqual(mse_result, 0.625) + + metrics = METRICS('tensorflow') + mse = metrics['MSE'](compare_label=False) + mse.update(predicts1, labels1) + mse_result = mse.result() + self.assertEqual(mse_result, 0.75) + mse.update(predicts2, labels2) + mse_result = mse.result() + self.assertEqual(mse_result, 0.625) + + + metrics = METRICS('mxnet') + mse = metrics['MSE']() + mse.update(predicts1, labels1) + mse_result = mse.result() + self.assertEqual(mse_result, 0.75) + mse.update(predicts2, labels2) + mse_result = mse.result() + self.assertEqual(mse_result, 0.625) + + metrics = METRICS('pytorch') + mse = metrics['MSE']() + mse.update(predicts1, labels1) + mse_result = mse.result() + self.assertEqual(mse_result, 0.75) + mse.update(predicts2, labels2) + mse_result = mse.result() + self.assertEqual(mse_result, 0.625) + + @unittest.skipIf(platform.system().lower() == "windows", "not support mxnet on windows now") + def test_mae(self): + predicts1 = [1, 0, 0, 1] + labels1 = [0, 1, 0, 0] + predicts2 = [1, 1, 1, 1] + labels2 = [1, 1, 1, 0] + + metrics = METRICS('tensorflow') + mae = metrics['MAE']() + mae.update(predicts1, labels1) + mae_result = mae.result() + self.assertEqual(mae_result, 0.75) + mae.update(0, 1) + mae_result = mae.result() + self.assertEqual(mae_result, 0.8) + mae.reset() + mae.update(predicts2, labels2) + mae_result = mae.result() + self.assertEqual(mae_result, 0.25) + + metrics = METRICS('pytorch') + mae = metrics['MAE']() + mae.update(predicts1, labels1) + mae_result = mae.result() + self.assertEqual(mae_result, 0.75) + mae.update(predicts2, labels2) + mae_result = mae.result() + self.assertEqual(mae_result, 0.5) + + metrics = METRICS('mxnet') + mae = metrics['MAE']() + mae.update(predicts1, labels1) + mae_result = mae.result() + self.assertEqual(mae_result, 0.75) + mae.update(predicts2, labels2) + mae_result = mae.result() + self.assertEqual(mae_result, 0.5) + + metrics = METRICS('onnxrt_qlinearops') + mae = metrics['MAE']() + mae.update(predicts1, labels1) + mae_result = mae.result() + self.assertEqual(mae_result, 0.75) + mae.update(predicts2, labels2) + mae_result = mae.result() + self.assertEqual(mae_result, 0.5) + + self.assertRaises(AssertionError, mae.update, [1], [1, 2]) + self.assertRaises(AssertionError, mae.update, 1, [1,2]) + self.assertRaises(AssertionError, mae.update, [1, 2], [1]) + self.assertRaises(AssertionError, mae.update, 1, np.array([1,2])) + + @unittest.skipIf(platform.system().lower() == "windows", "not support mxnet on windows now") + def test_rmse(self): + predicts1 = [1, 0, 0, 1] + labels1 = [1, 0, 0, 0] + predicts2 = [1, 1, 1, 1] + labels2 = [1, 0, 0, 0] + + metrics = METRICS('tensorflow') + rmse = metrics['RMSE']() + rmse.update(predicts1, labels1) + rmse_result = rmse.result() + self.assertEqual(rmse_result, 0.5) + rmse.reset() + rmse.update(predicts2, labels2) + rmse_result = rmse.result() + self.assertAlmostEqual(rmse_result, np.sqrt(0.75)) + + metrics = METRICS('pytorch') + rmse = metrics['RMSE']() + rmse.update(predicts1, labels1) + rmse_result = rmse.result() + self.assertEqual(rmse_result, 0.5) + rmse.update(predicts2, labels2) + rmse_result = rmse.result() + self.assertAlmostEqual(rmse_result, np.sqrt(0.5)) + + metrics = METRICS('mxnet') + rmse = metrics['RMSE']() + rmse.update(predicts1, labels1) + rmse_result = rmse.result() + self.assertEqual(rmse_result, 0.5) + rmse.update(predicts2, labels2) + rmse_result = rmse.result() + self.assertAlmostEqual(rmse_result, np.sqrt(0.5)) + + metrics = METRICS('onnxrt_qlinearops') + rmse = metrics['RMSE']() + rmse.update(predicts1, labels1) + rmse_result = rmse.result() + self.assertEqual(rmse_result, 0.5) + rmse.update(predicts2, labels2) + rmse_result = rmse.result() + self.assertAlmostEqual(rmse_result, np.sqrt(0.5)) + + def test_loss(self): + metrics = METRICS('pytorch') + loss = metrics['Loss']() + predicts = [1, 0, 0, 1] + labels = [0, 1, 0, 0] + loss.update(predicts, labels) + loss_result = loss.result() + self.assertEqual(loss_result, 0.5) + predicts = [1, 1, 0, 1] + labels = [0, 1, 0, 0] + loss.update(predicts, labels) + loss_result = loss.result() + self.assertEqual(loss_result, 0.625) + loss.reset() + predicts = [1, 0, 0, 1] + labels = [0, 1, 0, 0] + loss.update(predicts, labels) + self.assertEqual(loss.result(), 0.5) + + + metrics = METRICS('onnxrt_qlinearops') + loss = metrics['Loss']() + predicts = [1, 0, 0, 1] + labels = [0, 1, 0, 0] + loss.update(predicts, labels) + loss_result = loss.result() + self.assertEqual(loss_result, 0.5) + predicts = [1, 1, 0, 1] + labels = [0, 1, 0, 0] + loss.update(predicts, labels) + loss_result = loss.result() + self.assertEqual(loss_result, 0.625) + loss.reset() + predicts = [1, 0, 0, 1] + labels = [0, 1, 0, 0] + loss.update(predicts, labels) + self.assertEqual(loss.result(), 0.5) + +if __name__ == "__main__": + unittest.main() diff --git a/test/model/test_model.py b/test/model/test_model.py index 492ada7f844..ae322d92e09 100644 --- a/test/model/test_model.py +++ b/test/model/test_model.py @@ -4,7 +4,8 @@ import os import platform from neural_compressor.model import MODELS -import neural_compressor.model.model as NCModel +from neural_compressor.model.onnx_model import ONNXModel +from neural_compressor.model.mxnet_model import MXNetModel from neural_compressor.model.model import get_model_fwk_name from neural_compressor.experimental.common.model import Model @@ -134,7 +135,7 @@ def test_graph(self): self.assertEqual(True, isinstance(model.graph_def, tf.compat.v1.GraphDef)) def test_validate_graph_node(self): - from neural_compressor.model.model import validate_graph_node + from neural_compressor.model.tensorflow_model import validate_graph_node graph = build_graph() self.assertEqual(False, validate_graph_node(graph.as_graph_def(), [])) self.assertEqual(False, validate_graph_node(graph.as_graph_def(), ['test'])) @@ -253,7 +254,7 @@ def test_tf_qat_model(self): keras_model = build_keras() self.assertEqual('tensorflow', get_model_fwk_name(keras_model)) - from neural_compressor.model.model import TensorflowQATModel + from neural_compressor.model.tensorflow_model import TensorflowQATModel model = TensorflowQATModel(keras_model) assert isinstance(model.model, tf.keras.Model) keras_model.save('./simple_model') @@ -323,7 +324,7 @@ def test_saved_model(self): def test_tensorflow(self): - from neural_compressor.model.model import TensorflowBaseModel + from neural_compressor.model.tensorflow_model import TensorflowBaseModel ori_model = build_graph() self.assertEqual('tensorflow', get_model_fwk_name(ori_model)) self.assertEqual('tensorflow', get_model_fwk_name(TensorflowBaseModel(ori_model))) @@ -367,7 +368,7 @@ def tearDownClass(self): def test_model(self): self.assertEqual('onnxruntime', get_model_fwk_name(self.cnn_export_path)) model = MODELS['onnxruntime'](self.cnn_model) - self.assertEqual(True, isinstance(model, NCModel.ONNXModel)) + self.assertEqual(True, isinstance(model, ONNXModel)) self.assertEqual(True, isinstance(model.model, onnx.ModelProto)) model.save('test.onnx') @@ -377,7 +378,7 @@ def test_model(self): class TestPyTorchModel(unittest.TestCase): def testPyTorch(self): import torchvision - from neural_compressor.model.model import PyTorchModel, PyTorchIpexModel, PyTorchFXModel + from neural_compressor.model.torch_model import PyTorchModel, IPEXModel, PyTorchFXModel ori_model = torchvision.models.mobilenet_v2() self.assertEqual('pytorch', get_model_fwk_name(ori_model)) pt_model = PyTorchModel(ori_model) @@ -386,7 +387,7 @@ def testPyTorch(self): with self.assertRaises(AssertionError): pt_model.workspace_path = './pytorch' - ipex_model = PyTorchIpexModel(ori_model) + ipex_model = IPEXModel(ori_model) self.assertTrue(ipex_model.model) ipex_model.model = ori_model ipex_model = PyTorchModel(torchvision.models.mobilenet_v2()) @@ -395,7 +396,7 @@ def testPyTorch(self): ipex_model.save('./') self.assertEqual('pytorch', get_model_fwk_name(PyTorchModel(ori_model))) - self.assertEqual('pytorch', get_model_fwk_name(PyTorchIpexModel(ori_model))) + self.assertEqual('pytorch', get_model_fwk_name(IPEXModel(ori_model))) self.assertEqual('pytorch', get_model_fwk_name(PyTorchFXModel(ori_model))) def load_mxnet_model(symbol_file, param_file): @@ -438,7 +439,7 @@ def test_model(self): import mxnet as mx self.assertEqual('mxnet', get_model_fwk_name(self.net)) model = MODELS['mxnet'](self.net) - self.assertEqual(True, isinstance(model, NCModel.MXNetModel)) + self.assertEqual(True, isinstance(model, MXNetModel)) self.assertEqual(True, isinstance(model.model, mx.gluon.HybridBlock)) model.save('./test') diff --git a/test/model/test_model_pytorch.py b/test/model/test_model_pytorch.py index a6bae08fa77..f2c62be5714 100644 --- a/test/model/test_model_pytorch.py +++ b/test/model/test_model_pytorch.py @@ -2,7 +2,8 @@ import torchvision import unittest import neural_compressor.adaptor.pytorch as nc_torch -from neural_compressor.model import MODELS +from neural_compressor.model import MODELS, Model +from neural_compressor.model.torch_model import PyTorchModel from packaging.version import Version try: @@ -23,6 +24,11 @@ class TestPytorchModel(unittest.TestCase): model = torchvision.models.quantization.resnet18() lpot_model = MODELS['pytorch'](model) + def test_Model(self): + model = torchvision.models.quantization.resnet18() + inc_model = Model(model) + self.assertTrue(isinstance(inc_model, PyTorchModel)) + def test_get_all_weight_name(self): assert len(list(self.lpot_model.get_all_weight_names())) == 62 diff --git a/test/model/test_tensorflow_auto_input_output.py b/test/model/test_tensorflow_auto_input_output.py index 28eeff06b1e..2e9776a875f 100644 --- a/test/model/test_tensorflow_auto_input_output.py +++ b/test/model/test_tensorflow_auto_input_output.py @@ -6,7 +6,7 @@ import platform from neural_compressor.adaptor.tensorflow import TensorFlowAdaptor from neural_compressor.experimental.common.model import Model as TensorflowModel -from neural_compressor.model.model import validate_graph_node +from neural_compressor.model.tensorflow_model import validate_graph_node class TestTFAutoDetectInputOutput(unittest.TestCase): mb_model_url = 'https://storage.googleapis.com/intel-optimized-tensorflow/models/v1_6/mobilenet_v1_1.0_224_frozen.pb' diff --git a/test/nas/test_nas.py b/test/nas/test_nas.py index 4d22673d578..cdf00275d5e 100644 --- a/test/nas/test_nas.py +++ b/test/nas/test_nas.py @@ -6,7 +6,7 @@ import torch from neural_compressor.conf.config import NASConfig -from neural_compressor.data import DATASETS +from neural_compressor.data import Datasets from neural_compressor.experimental import common, NAS from neural_compressor.experimental.data.dataloaders.pytorch_dataloader import PyTorchDataLoader from neural_compressor.experimental.nas.dynas import DyNAS @@ -143,7 +143,7 @@ def test_basic_nas(self): self.assertTrue(len(best_model_archs) > 0) # Customized train, evaluation - datasets = DATASETS('pytorch') + datasets = Datasets('pytorch') dummy_dataset = datasets['dummy'](shape=(32, 3, 64, 64), low=0., high=1., label=True) dummy_dataloader = PyTorchDataLoader(dummy_dataset) def train_func(model): diff --git a/test/objective/test_objective.py b/test/objective/test_objective.py index e21aa0e57bc..cd4cea35464 100644 --- a/test/objective/test_objective.py +++ b/test/objective/test_objective.py @@ -225,8 +225,8 @@ def tearDownClass(self): shutil.rmtree('./saved', ignore_errors=True) def test_performance(self): - from neural_compressor.data import DATASETS - dataset = DATASETS('tensorflow')['dummy']((100, 256, 256, 1), label=True) + from neural_compressor.data import Datasets + dataset = Datasets('tensorflow')['dummy']((100, 256, 256, 1), label=True) from neural_compressor.experimental import Quantization, common from neural_compressor.utils.utility import get_size @@ -245,8 +245,8 @@ def test_performance(self): def test_model_size(self): from neural_compressor.experimental import Benchmark, common - from neural_compressor.data import DATASETS - dataset = DATASETS('tensorflow')['dummy']((100, 256, 256, 1), label=True) + from neural_compressor.data import Datasets + dataset = Datasets('tensorflow')['dummy']((100, 256, 256, 1), label=True) benchmarker = Benchmark('fake_yaml_model_size.yaml') benchmarker.b_dataloader = common.DataLoader(dataset) @@ -255,8 +255,8 @@ def test_model_size(self): def test_footprint(self): from neural_compressor.experimental import Benchmark, common - from neural_compressor.data import DATASETS - dataset = DATASETS('tensorflow')['dummy']((100, 256, 256, 1), label=True) + from neural_compressor.data import Datasets + dataset = Datasets('tensorflow')['dummy']((100, 256, 256, 1), label=True) benchmarker = Benchmark('fake_yaml_footprint.yaml') benchmarker.b_dataloader = common.DataLoader(dataset) diff --git a/test/pruning/test_gradient_sensitivity.py b/test/pruning/test_gradient_sensitivity.py index 999fdcccc28..6d411e35b0f 100644 --- a/test/pruning/test_gradient_sensitivity.py +++ b/test/pruning/test_gradient_sensitivity.py @@ -2,7 +2,7 @@ import shutil import unittest from neural_compressor.experimental.data.dataloaders.pytorch_dataloader import PyTorchDataLoader -from neural_compressor.data import DATASETS +from neural_compressor.data import Datasets import torch import torchvision @@ -206,7 +206,7 @@ def tearDownClass(cls): def test_unstructured_pruning(self): from neural_compressor.experimental import Pruning, common prune_cv = Pruning('fake_unstructured.yaml') - datasets = DATASETS('pytorch') + datasets = Datasets('pytorch') dummy_dataset = datasets['dummy'](shape=(100, 3, 224, 224), low=0., high=1., label=True) dummy_dataloader = PyTorchDataLoader(dummy_dataset) diff --git a/test/pruning/test_pruning.py b/test/pruning/test_pruning.py index 3e1290e6bb7..5871f6bcc34 100644 --- a/test/pruning/test_pruning.py +++ b/test/pruning/test_pruning.py @@ -7,7 +7,7 @@ import torch.nn as nn from neural_compressor.config import Pruner, PruningConfig -from neural_compressor.data import DATASETS +from neural_compressor.data import Datasets from neural_compressor.experimental.data.dataloaders.pytorch_dataloader import PyTorchDataLoader from neural_compressor.training import prepare_compression @@ -60,7 +60,7 @@ def test_pruning(self): pruner1 = Pruner(start_epoch=1, end_epoch=2, names=['layer1.0.conv1.weight']) pruner2 = Pruner(target_sparsity=0.6, update_frequency=2, names=['layer1.0.conv2.weight']) conf = PruningConfig(pruners=[pruner1, pruner2], end_epoch=2) - datasets = DATASETS('pytorch') + datasets = Datasets('pytorch') dummy_dataset = datasets['dummy'](shape=(100, 3, 224, 224), low=0., high=1., label=True) dummy_dataloader = PyTorchDataLoader(dummy_dataset) compression_manager = prepare_compression(self.model, conf) @@ -98,7 +98,7 @@ def test_pruning_external(self): Pruner(target_sparsity=0.6,update_frequency=2,names=['layer1.0.conv2.weight'])] conf = PruningConfig(pruners) - datasets = DATASETS('pytorch') + datasets = Datasets('pytorch') dummy_dataset = datasets['dummy'](shape=(100, 3, 224, 224), low=0., high=1., label=True) dummy_dataloader = PyTorchDataLoader(dummy_dataset) compression_manager = prepare_compression(self.model, conf) diff --git a/test/pruning/test_pruning_group_lasso.py b/test/pruning/test_pruning_group_lasso.py index 0f129874e95..73c9c4d70b6 100644 --- a/test/pruning/test_pruning_group_lasso.py +++ b/test/pruning/test_pruning_group_lasso.py @@ -6,7 +6,7 @@ import torchvision import torch.nn as nn -from neural_compressor.data import DATASETS +from neural_compressor.data import Datasets from neural_compressor.experimental.data.dataloaders.pytorch_dataloader import PyTorchDataLoader def build_fake_yaml(): diff --git a/test/pruning/test_pruning_pure_yaml.py b/test/pruning/test_pruning_pure_yaml.py index 312b1d8cefb..b8b19dd36db 100644 --- a/test/pruning/test_pruning_pure_yaml.py +++ b/test/pruning/test_pruning_pure_yaml.py @@ -6,7 +6,7 @@ import torchvision import torch.nn as nn -from neural_compressor.data import DATASETS +from neural_compressor.data import Datasets from neural_compressor.experimental.data.dataloaders.pytorch_dataloader import PyTorchDataLoader def build_fake_yaml(): diff --git a/test/pruning/test_pytorch_pruning.py b/test/pruning/test_pytorch_pruning.py index 73739de75ab..5fb2047b7b2 100644 --- a/test/pruning/test_pytorch_pruning.py +++ b/test/pruning/test_pytorch_pruning.py @@ -6,7 +6,7 @@ import torchvision import torch.nn as nn -from neural_compressor.data import DATASETS +from neural_compressor.data import Datasets from neural_compressor.experimental.data.dataloaders.pytorch_dataloader import PyTorchDataLoader @@ -141,7 +141,7 @@ def test_pytorch_pruning_basic(self): prune.model = self.model criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(self.model.parameters(), lr=0.0001) - datasets = DATASETS('pytorch') + datasets = Datasets('pytorch') dummy_dataset = datasets['dummy'](shape=(10, 3, 224, 224), low=0., high=1., label=True) dummy_dataloader = PyTorchDataLoader(dummy_dataset) prune.on_train_begin() @@ -175,7 +175,7 @@ def test_pytorch_pruner_channel_pruning(self): prune.model = self.model criterion = nn.CrossEntropyLoss() optimizer = torch.optim.SGD(self.model.parameters(), lr=0.0001) - datasets = DATASETS('pytorch') + datasets = Datasets('pytorch') dummy_dataset = datasets['dummy'](shape=(10, 3, 224, 224), low=0., high=1., label=True) dummy_dataloader = PyTorchDataLoader(dummy_dataset) prune.on_train_begin() diff --git a/test/scheduler/test_oneshot.py b/test/scheduler/test_oneshot.py index 99462ff670d..484a7449c13 100644 --- a/test/scheduler/test_oneshot.py +++ b/test/scheduler/test_oneshot.py @@ -9,7 +9,7 @@ import neural_compressor.adaptor.pytorch as nc_torch from neural_compressor.conf.config import DistillationConf, PruningConf -from neural_compressor.data import DATASETS +from neural_compressor.data import Datasets from neural_compressor.experimental.data.dataloaders.pytorch_dataloader import PyTorchDataLoader from neural_compressor.experimental.scheduler import Scheduler from neural_compressor.training import prepare_compression @@ -189,7 +189,7 @@ def tearDownClass(cls): def test_prune_qat_oneshot(self): from neural_compressor.experimental import Pruning, Quantization - datasets = DATASETS('pytorch') + datasets = Datasets('pytorch') dummy_dataset = datasets['dummy'](shape=(16, 3, 224, 224), low=0., high=1., label=True) dummy_dataloader = PyTorchDataLoader(dummy_dataset) q_model = copy.deepcopy(self.q_model) @@ -250,7 +250,7 @@ def train_func_for_nc(model): def test_distillation_qat_oneshot(self): from neural_compressor.experimental import Distillation, Quantization - datasets = DATASETS('pytorch') + datasets = Datasets('pytorch') dummy_dataset = datasets['dummy'](shape=(16, 3, 224, 224), low=0., high=1., label=True) dummy_dataloader = PyTorchDataLoader(dummy_dataset) model = copy.deepcopy(self.model) @@ -303,7 +303,7 @@ def train_func_for_nc(model): def test_distillation_prune_oneshot_with_new_API(self): from neural_compressor.config import DistillationConfig, KnowledgeDistillationLossConfig from neural_compressor.config import Pruner, PruningConfig - datasets = DATASETS('pytorch') + datasets = Datasets('pytorch') dummy_dataset = datasets['dummy'](shape=(16, 3, 224, 224), low=0., high=1., label=True) dummy_dataloader = PyTorchDataLoader(dummy_dataset) model = copy.deepcopy(self.model) @@ -363,7 +363,7 @@ def train_func_for_nc(model): def test_prune_qat_distillation_oneshot(self): from neural_compressor.experimental import Pruning, Quantization, Distillation - datasets = DATASETS('pytorch') + datasets = Datasets('pytorch') dummy_dataset = datasets['dummy'](shape=(16, 3, 224, 224), low=0., high=1., label=True) dummy_dataloader = PyTorchDataLoader(dummy_dataset) model = copy.deepcopy(self.model) @@ -421,7 +421,7 @@ def train_func_for_nc(model): def test_prune_qat_oneshot_fx(self): from neural_compressor.experimental import Pruning, Quantization - datasets = DATASETS('pytorch_fx') + datasets = Datasets('pytorch_fx') dummy_dataset = datasets['dummy'](shape=(16, 3, 224, 224), low=0., high=1., label=True) dummy_dataloader = PyTorchDataLoader(dummy_dataset) prune = Pruning('./fx_fake.yaml') @@ -479,7 +479,7 @@ def train_func_for_nc(model): "requires higher version of torch than 1.9.0") def test_distillation_qat_oneshot_fx(self): from neural_compressor.experimental import Distillation, Quantization - datasets = DATASETS('pytorch_fx') + datasets = Datasets('pytorch_fx') dummy_dataset = datasets['dummy'](shape=(16, 3, 224, 224), low=0., high=1., label=True) dummy_dataloader = PyTorchDataLoader(dummy_dataset) model = DynamicControlModel() @@ -532,7 +532,7 @@ def train_func_for_nc(model): def test_distillation_prune_oneshot_fx(self): from neural_compressor.experimental import Distillation, Pruning - datasets = DATASETS('pytorch_fx') + datasets = Datasets('pytorch_fx') dummy_dataset = datasets['dummy'](shape=(16, 3, 224, 224), low=0., high=1., label=True) dummy_dataloader = PyTorchDataLoader(dummy_dataset) distiller = Distillation('./fx_fake3.yaml') @@ -590,7 +590,7 @@ def train_func_for_nc(model): "requires higher version of torch than 1.9.0") def test_prune_qat_distillation_oneshot_fx(self): from neural_compressor.experimental import Pruning, Quantization, Distillation - datasets = DATASETS('pytorch_fx') + datasets = Datasets('pytorch_fx') dummy_dataset = datasets['dummy'](shape=(16, 3, 224, 224), low=0., high=1., label=True) dummy_dataloader = PyTorchDataLoader(dummy_dataset) model = copy.deepcopy(self.model) diff --git a/test/scheduler/test_scheduler.py b/test/scheduler/test_scheduler.py index 7656cfb7c5a..ed1f0698f5a 100644 --- a/test/scheduler/test_scheduler.py +++ b/test/scheduler/test_scheduler.py @@ -7,7 +7,7 @@ import torch.nn as nn import neural_compressor.adaptor.pytorch as nc_torch -from neural_compressor.data import DATASETS +from neural_compressor.data import Datasets from neural_compressor.experimental.data.dataloaders.pytorch_dataloader import PyTorchDataLoader from neural_compressor.experimental.scheduler import Scheduler from packaging.version import Version @@ -325,7 +325,7 @@ def test_pruning(self): prune = Pruning('fake.yaml') scheduler = Scheduler() scheduler.model = self.model - datasets = DATASETS('pytorch') + datasets = Datasets('pytorch') dummy_dataset = datasets['dummy'](shape=(16, 3, 224, 224), low=0., high=1., label=True) dummy_dataloader = PyTorchDataLoader(dummy_dataset) diff --git a/test/strategy/test_basic.py b/test/strategy/test_basic.py index 436f56c5240..15b6c4ff1cd 100644 --- a/test/strategy/test_basic.py +++ b/test/strategy/test_basic.py @@ -221,10 +221,10 @@ def test_run_basic_max_trials_multimetric_weight(self): def test_run_basic_one_trial_new_api(self): from neural_compressor.quantization import fit from neural_compressor.config import PostTrainingQuantConfig - from neural_compressor.data import DATASETS, DATALOADERS + from neural_compressor.data import Datasets, DATALOADERS # dataset and dataloader - dataset = DATASETS("tensorflow")["dummy"](((100, 3, 3, 1))) + dataset = Datasets("tensorflow")["dummy"](((100, 3, 3, 1))) dataloader = DATALOADERS["tensorflow"](dataset) # tuning and accuracy criterion diff --git a/test/strategy/test_hawq_v2_2.x.py b/test/strategy/test_hawq_v2_2.x.py index 19b52e07826..c3858063a1a 100644 --- a/test/strategy/test_hawq_v2_2.x.py +++ b/test/strategy/test_hawq_v2_2.x.py @@ -28,7 +28,7 @@ def test_hawq_v2_pipeline(self): logger.info("*** Test: HAWQ v2 with pytorch model.") from neural_compressor.quantization import fit from neural_compressor.config import PostTrainingQuantConfig, TuningCriterion - from neural_compressor.data import DATASETS, DATALOADERS + from neural_compressor.data import Datasets, DATALOADERS # model model = copy.deepcopy(self.model) @@ -40,7 +40,7 @@ def _fake_eval(model): return self.test_hawq_v2_pipeline_fake_acc # dataset and dataloader - dataset = DATASETS("pytorch")["dummy"](((1, 3, 224, 224))) + dataset = Datasets("pytorch")["dummy"](((1, 3, 224, 224))) dataloader = DATALOADERS["pytorch"](dataset) #tuning and accuracy criterion @@ -57,4 +57,4 @@ def _fake_eval(model): self.assertIsNone(q_model) if __name__ == "__main__": - unittest.main() \ No newline at end of file + unittest.main() diff --git a/test/strategy/test_mse_v2_2.x.py b/test/strategy/test_mse_v2_2.x.py index a6b0219c62a..6f33aa14e1d 100644 --- a/test/strategy/test_mse_v2_2.x.py +++ b/test/strategy/test_mse_v2_2.x.py @@ -63,8 +63,8 @@ def fake_eval_func(_): from neural_compressor.quantization import fit from neural_compressor.config import TuningCriterion, PostTrainingQuantConfig - from neural_compressor.data import DATASETS, DATALOADERS - dataset = DATASETS("tensorflow")["dummy"](((100, 3, 3, 1))) + from neural_compressor.data import Datasets, DATALOADERS + dataset = Datasets("tensorflow")["dummy"](((100, 3, 3, 1))) dataloader = DATALOADERS['tensorflow'](dataset) conf = PostTrainingQuantConfig( @@ -90,8 +90,8 @@ def fake_eval_func(_): from neural_compressor.quantization import fit from neural_compressor.config import TuningCriterion, PostTrainingQuantConfig - from neural_compressor.data import DATASETS, DATALOADERS - dataset = DATASETS("tensorflow")["dummy"](((100, 3, 3, 1))) + from neural_compressor.data import Datasets, DATALOADERS + dataset = Datasets("tensorflow")["dummy"](((100, 3, 3, 1))) dataloader = DATALOADERS['tensorflow'](dataset) conf = PostTrainingQuantConfig( @@ -120,8 +120,8 @@ def fake_eval_func(model): from neural_compressor.quantization import fit from neural_compressor.config import TuningCriterion, PostTrainingQuantConfig - from neural_compressor.data import DATASETS, DATALOADERS - dataset = DATASETS("pytorch")["dummy"](((1, 3, 224, 224))) + from neural_compressor.data import Datasets, DATALOADERS + dataset = Datasets("pytorch")["dummy"](((1, 3, 224, 224))) dataloader = DATALOADERS['pytorch'](dataset) conf = PostTrainingQuantConfig( diff --git a/test/strategy/test_optimization_level_2.x.py b/test/strategy/test_optimization_level_2.x.py index 31fd31de6c8..d3077bd45d3 100644 --- a/test/strategy/test_optimization_level_2.x.py +++ b/test/strategy/test_optimization_level_2.x.py @@ -61,14 +61,14 @@ def test_tf_opt_level_0(self): logger.info("*** Test: optimization level 0 with tensorflow model.") from neural_compressor.quantization import fit from neural_compressor.config import PostTrainingQuantConfig - from neural_compressor.data import DATASETS, DATALOADERS + from neural_compressor.data import Datasets, DATALOADERS # fake evaluation function def _fake_eval(model): return 1 # dataset and dataloader - dataset = DATASETS("tensorflow")["dummy"](((100, 3, 3, 1))) + dataset = Datasets("tensorflow")["dummy"](((100, 3, 3, 1))) dataloader = DATALOADERS["tensorflow"](dataset) # tuning and accuracy criterion @@ -87,7 +87,7 @@ def test_tf_opt_level_1(self): logger.info("*** Test: optimization level 1 with tensorflow model.") from neural_compressor.quantization import fit from neural_compressor.config import PostTrainingQuantConfig - from neural_compressor.data import DATASETS, DATALOADERS + from neural_compressor.data import Datasets, DATALOADERS # fake evaluation function self._fake_acc = 10 @@ -96,7 +96,7 @@ def _fake_eval(model): return self._fake_acc # dataset and dataloader - dataset = DATASETS("tensorflow")["dummy"](((100, 3, 3, 1))) + dataset = Datasets("tensorflow")["dummy"](((100, 3, 3, 1))) dataloader = DATALOADERS["tensorflow"](dataset) # tuning and accuracy criterion @@ -114,7 +114,7 @@ def test_pt_opt_level_0(self): logger.info("*** Test: optimization level 0 with pytorch model.") from neural_compressor.quantization import fit from neural_compressor.config import PostTrainingQuantConfig - from neural_compressor.data import DATASETS, DATALOADERS + from neural_compressor.data import Datasets, DATALOADERS import torchvision # model @@ -131,7 +131,7 @@ def _fake_eval(model): return acc_lst[self.test_pt_opt_level_0_index] # dataset and dataloader - dataset = DATASETS("pytorch")["dummy"](((100, 3, 3, 1))) + dataset = Datasets("pytorch")["dummy"](((100, 3, 3, 1))) dataloader = DATALOADERS["pytorch"](dataset) # tuning and accuracy criterion diff --git a/test/strategy/test_sigopt.py b/test/strategy/test_sigopt.py index eb7fa77e25b..1b78167edb7 100644 --- a/test/strategy/test_sigopt.py +++ b/test/strategy/test_sigopt.py @@ -144,10 +144,10 @@ def test_run_basic_max_trials(self): def test_run_sigopt_one_trial_new_api(self): from neural_compressor.quantization import fit from neural_compressor.config import AccuracyCriterion, PostTrainingQuantConfig, TuningCriterion - from neural_compressor.data import DATASETS, DATALOADERS + from neural_compressor.data import Datasets, DATALOADERS # dataset and dataloader - dataset = DATASETS("tensorflow")["dummy"](((100, 3, 3, 1))) + dataset = Datasets("tensorflow")["dummy"](((100, 3, 3, 1))) dataloader = DATALOADERS["tensorflow"](dataset) # tuning and accuracy criterion diff --git a/test/tfnewapi/test_tf_spr_base_distributed_tf_dataloader.py b/test/tfnewapi/test_tf_spr_base_distributed_tf_dataloader.py index 9f4b1d69ba6..88f249ba522 100644 --- a/test/tfnewapi/test_tf_spr_base_distributed_tf_dataloader.py +++ b/test/tfnewapi/test_tf_spr_base_distributed_tf_dataloader.py @@ -11,7 +11,7 @@ from neural_compressor import data from neural_compressor.utils.create_obj_from_config import create_dataset, create_dataloader from neural_compressor.data.dataloaders.dataloader import DataLoader -from neural_compressor.data import DATASETS, DATALOADERS, TRANSFORMS +from neural_compressor.data import Datasets, DATALOADERS, TRANSFORMS from neural_compressor.utils import logger from neural_compressor.adaptor.tf_utils.util import version1_lt_version2