replace UT DATASETS

Signed-off-by: changwa1 <[email protected]>

Author: changwangss

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')

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

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)

test/adaptor/pytorch_adaptor/test_adaptor_pytorch_2.x.py

@@ -10,11 +10,11 @@
 import os
 from neural_compressor import PostTrainingQuantConfig, QuantizationAwareTrainingConfig
 from neural_compressor.config import set_tensorboard, set_workspace
-from neural_compressor.data import DATASETS, DATALOADERS
+from neural_compressor.data import Datasets, DATALOADERS
 from neural_compressor.adaptor import FRAMEWORKS
 from neural_compressor.model import MODELS
 from neural_compressor.experimental import Quantization, common
-from neural_compressor.experimental.data.datasets.dataset import DATASETS
+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
@@ -419,7 +419,7 @@ def test_quantization_new_API(self):
                 compression_manager = prepare_compression(copy.deepcopy(model), quant_conf)
                 q_model = train_func(compression_manager, compression_manager.model)
             else:
-                dataset = DATASETS("pytorch")["dummy"]((100, 3, 224, 224))
+                dataset = Datasets("pytorch")["dummy"]((100, 3, 224, 224))
                 dataloader = DATALOADERS["pytorch"](dataset)
                 if fake_yaml == "dynamic":
                     quant_conf = PostTrainingQuantConfig(approach="dynamic",
@@ -449,7 +449,7 @@ def eval_func(model):
         # run fx_quant in neural_compressor and save the quantized GraphModule
         quant_conf = PostTrainingQuantConfig(approach="auto")
         set_workspace("./saved")
-        dataset = DATASETS("pytorch")["dummy"]((100, 3, 224, 224))
+        dataset = Datasets("pytorch")["dummy"]((100, 3, 224, 224))
         dataloader = common.DataLoader(dataset)
         q_model = quantization.fit(model_origin,
                                    quant_conf,
@@ -469,7 +469,7 @@ def test_tensorboard(self):
         quant_conf = PostTrainingQuantConfig(approach="static",
                                         backend="pytorch")
         set_tensorboard(True)
-        dataset = DATASETS("pytorch")["dummy"]((100, 3, 224, 224))
+        dataset = Datasets("pytorch")["dummy"]((100, 3, 224, 224))
         dataloader = common.DataLoader(dataset)
         quantization.fit(
           model, quant_conf, calib_dataloader=dataloader, eval_func=eval_func
@@ -493,7 +493,7 @@ def tearDownClass(self):
     def test_fx_quant(self):
         for fake_yaml 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 fake_yaml == "qat":
                 conf = QuantizationAwareTrainingConfig(
@@ -526,7 +526,7 @@ def test_fx_quant(self):
         for fake_yaml 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 fake_yaml == "qat":
                 conf = QuantizationAwareTrainingConfig(
@@ -600,7 +600,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 fake_yaml == "qat":
@@ -622,7 +622,7 @@ def eval_func(model):
     def test_fx_sub_module_quant(self):
         for fake_yaml 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 fake_yaml == "qat":
@@ -658,7 +658,7 @@ 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))
+        dataset = Datasets("pytorch")["dummy"]((100, 3, 224, 224))
         dataloader = DATALOADERS["pytorch"](dataset)
         set_workspace=("./saved")
         conf = PostTrainingQuantConfig(op_name_list=ptq_fx_op_name_list, backend="pytorch_fx")

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
 
@@ -212,7 +212,7 @@ def test_fx_quant(self):
                 approach=fake_yaml,
                 backend="pytorch_fx"
             )
-            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,

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",

test/benchmark/test_benchmark_2.x.py

@@ -18,9 +18,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)
@@ -33,8 +33,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)
@@ -64,8 +64,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",

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
@@ -202,7 +202,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):
@@ -245,7 +245,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):

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',

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: