Remove some redundant classes for quantization, benchmark and mixed precision

docs/source/benchmark.md

@@ -49,7 +49,7 @@ And please make sure `cores_per_instance * num_of_instance` must be less than CP
 from neural_compressor.config import BenchmarkConfig
 from neural_compressor.benchmark import fit
 conf = BenchmarkConfig(warmup=10, iteration=100, cores_per_instance=4, num_of_instance=7)
-fit(model='./int8.pb', config=conf, b_dataloader=eval_dataloader)
+fit(model='./int8.pb', conf=conf, b_dataloader=eval_dataloader)
 ```
 
 ## Examples

docs/source/mixed_precision.md

@@ -42,8 +42,8 @@ Supported precisions for mix precision include bf16 and fp16. If users want to g
 from neural_compressor import mix_precision
 from neural_compressor.config import MixedPrecisionConfig
 
-conf = MixedPrecisionConfig(precision='bf16')
-converted_model = mix_precision.fit(model, config=conf)
+conf = MixedPrecisionConfig() # default precision is bf16
+converted_model = mix_precision.fit(model, conf=conf)
 converted_model.save('./path/to/save/')
 ```
 
@@ -56,8 +56,8 @@ from neural_compressor.config import MixedPrecisionConfig
 conf = MixedPrecisionConfig(
         backend='onnxrt_cuda_ep',
         device='gpu',
-        precision='fp16')
-converted_model = mix_precision.fit(model, config=conf)
+        precisions='fp16')
+converted_model = mix_precision.fit(model, conf=conf)
 converted_model.save('./path/to/save/')
 ```
 

docs/source/pruning.md

@@ -301,7 +301,8 @@ The following section exemplifies how to use hooks in user pass-in training func
      [**Experimental option** ]Modify model and optimizer.
 
     ```python
-        from neural_compressor.training import prepare_pruning, WeightPruningConfig 
+        from neural_compressor import WeightPruningConfig
+        from neural_compressor.experimental.compression import prepare_pruning
         config = WeightPruningConfig(configs)
         prepare_pruning(config, model, optimizer) # modify model and optimizer
         for epoch in range(num_train_epochs):

examples/helloworld/tf_example5/README.md

@@ -52,7 +52,7 @@ python test.py --benchmark --dataset_location=/path/to/imagenet/
 ```python
     from neural_compressor.benchmark import fit
     conf = BenchmarkConfig(iteration=100, cores_per_instance=4, num_of_instance=1)
-    fit(model='./int8.pb', config=conf, b_dataloader=eval_dataloader)
+    fit(model='./int8.pb', conf=conf, b_dataloader=eval_dataloader)
 
 ```
 

examples/keras/image_recognition/inception_resnet_v2/quantization/ptq/main.py

@@ -119,7 +119,7 @@ def main(_):
     if FLAGS.tune:
         from neural_compressor.quantization import fit
         from neural_compressor.config import PostTrainingQuantConfig
-        from neural_compressor.utils.utility import set_random_seed
+        from neural_compressor import set_random_seed
         set_random_seed(9527)
         config = PostTrainingQuantConfig(backend='itex', 
             calibration_sampling_size=[20, 150])
@@ -139,7 +139,7 @@ def main(_):
             fit(FLAGS.input_model, conf, b_func=evaluate)
         else:
             from neural_compressor.model.model import Model
-            accuracy = evaluate(Model(FLAGS.input_model, backend='keras').model)
+            accuracy = evaluate(Model(FLAGS.input_model, backend='itex').model)
             logger.info('Batch size = %d' % FLAGS.batch_size)
             logger.info("Accuracy: %.5f" % accuracy)
 

examples/keras/image_recognition/inception_v3/quantization/ptq/main.py

@@ -119,7 +119,7 @@ def main(_):
     if FLAGS.tune:
         from neural_compressor.quantization import fit
         from neural_compressor.config import PostTrainingQuantConfig
-        from neural_compressor.utils.utility import set_random_seed
+        from neural_compressor import set_random_seed
         set_random_seed(9527)
         config = PostTrainingQuantConfig(backend='itex', 
             calibration_sampling_size=[50, 100])
@@ -139,9 +139,10 @@ def main(_):
             fit(FLAGS.input_model, conf, b_func=evaluate)
         else:
             from neural_compressor.model.model import Model
-            accuracy = evaluate(Model(FLAGS.input_model, backend='keras').model)
+            accuracy = evaluate(Model(FLAGS.input_model, backend='itex').model)
             logger.info('Batch size = %d' % FLAGS.batch_size)
             logger.info("Accuracy: %.5f" % accuracy)
 
+
 if __name__ == "__main__":
     tf.compat.v1.app.run()

examples/keras/image_recognition/mobilenet_v2/quantization/ptq/main.py

@@ -118,7 +118,7 @@ def main(_):
     if FLAGS.tune:
         from neural_compressor.quantization import fit
         from neural_compressor.config import PostTrainingQuantConfig
-        from neural_compressor.utils.utility import set_random_seed
+        from neural_compressor import set_random_seed
         set_random_seed(9527)
         config = PostTrainingQuantConfig(backend='itex', 
             calibration_sampling_size=[50, 100])
@@ -138,7 +138,7 @@ def main(_):
             fit(FLAGS.input_model, conf, b_func=evaluate)
         else:
             from neural_compressor.model.model import Model
-            accuracy = evaluate(Model(FLAGS.input_model, backend='keras').model)
+            accuracy = evaluate(Model(FLAGS.input_model, backend='itex').model)
             logger.info('Batch size = %d' % FLAGS.batch_size)
             logger.info("Accuracy: %.5f" % accuracy)
 

examples/keras/image_recognition/resnet101/quantization/ptq/main.py

@@ -125,7 +125,7 @@ def main(_):
     if FLAGS.tune:
         from neural_compressor.quantization import fit
         from neural_compressor.config import PostTrainingQuantConfig
-        from neural_compressor.utils.utility import set_random_seed
+        from neural_compressor import set_random_seed
         set_random_seed(9524)
         config = PostTrainingQuantConfig(backend='itex', 
             calibration_sampling_size=[10, 15])
@@ -145,7 +145,7 @@ def main(_):
             fit(FLAGS.input_model, conf, b_func=evaluate)
         else:
             from neural_compressor.model.model import Model
-            accuracy = evaluate(Model(FLAGS.input_model, backend='keras').model)
+            accuracy = evaluate(Model(FLAGS.input_model, backend='itex').model)
             logger.info('Batch size = %d' % FLAGS.batch_size)
             logger.info("Accuracy: %.5f" % accuracy)
 

examples/keras/image_recognition/resnet50/quantization/ptq/main.py

@@ -110,7 +110,7 @@ def eval_func(dataloader, metric):
     return acc
 
 def main(_):
-    from neural_compressor.utils import set_random_seed
+    from neural_compressor import set_random_seed
     set_random_seed(9527)
     if FLAGS.tune:
         from neural_compressor import quantization
@@ -130,7 +130,7 @@ def main(_):
             fit(FLAGS.input_model, conf, b_func=evaluate)
         else:
             from neural_compressor.model import Model
-            model = Model(FLAGS.input_model, backend='keras').model
+            model = Model(FLAGS.input_model, backend='itex').model
             accuracy = evaluate(model)
             print('Batch size = %d' % FLAGS.batch_size)
             print("Accuracy: %.5f" % accuracy)

examples/keras/image_recognition/resnet50_fashion/quantization/ptq/main.py

@@ -100,7 +100,7 @@ def eval_func(data_loader, metric):
     return acc
 
 def main(_):
-    from neural_compressor.utils import set_random_seed
+    from neural_compressor import set_random_seed
     set_random_seed(9527)
     if FLAGS.tune:
         from neural_compressor import quantization
@@ -128,7 +128,7 @@ def main(_):
             fit(FLAGS.input_model, conf, b_func=evaluate)
         else:
             from neural_compressor.model import Model
-            model = Model(FLAGS.input_model, backend='keras').model
+            model = Model(FLAGS.input_model, backend='itex').model
             accuracy = evaluate(model)
             print('Batch size = %d' % FLAGS.batch_size)
             print("Accuracy: %.5f" % accuracy)

examples/keras/image_recognition/resnetv2_101/quantization/ptq/main.py

@@ -118,7 +118,7 @@ def main(_):
     if FLAGS.tune:
         from neural_compressor.quantization import fit
         from neural_compressor.config import PostTrainingQuantConfig
-        from neural_compressor.utils.utility import set_random_seed
+        from neural_compressor import set_random_seed
         set_random_seed(9527)
         config = PostTrainingQuantConfig(backend='itex', 
             calibration_sampling_size=[50, 100])
@@ -138,7 +138,7 @@ def main(_):
             fit(FLAGS.input_model, conf, b_func=evaluate)
         else:
             from neural_compressor.model.model import Model
-            accuracy = evaluate(Model(FLAGS.input_model, backend='keras').model)
+            accuracy = evaluate(Model(FLAGS.input_model, backend='itex').model)
             logger.info('Batch size = %d' % FLAGS.batch_size)
             logger.info("Accuracy: %.5f" % accuracy)
 

examples/keras/image_recognition/resnetv2_50/quantization/ptq/main.py

@@ -118,7 +118,7 @@ def main(_):
     if FLAGS.tune:
         from neural_compressor.quantization import fit
         from neural_compressor.config import PostTrainingQuantConfig
-        from neural_compressor.utils.utility import set_random_seed
+        from neural_compressor import set_random_seed
         set_random_seed(9527)
         config = PostTrainingQuantConfig(backend='itex', 
             calibration_sampling_size=[50, 100])
@@ -138,7 +138,7 @@ def main(_):
             fit(FLAGS.input_model, conf, b_func=evaluate)
         else:
             from neural_compressor.model.model import Model
-            accuracy = evaluate(Model(FLAGS.input_model, backend='keras').model)
+            accuracy = evaluate(Model(FLAGS.input_model, backend='itex').model)
             logger.info('Batch size = %d' % FLAGS.batch_size)
             logger.info("Accuracy: %.5f" % accuracy)
 

examples/keras/image_recognition/vgg16/quantization/ptq/main.py

@@ -110,7 +110,7 @@ def eval_func(dataloader, metric):
     return acc
 
 def main(_):
-    from neural_compressor.utils import set_random_seed
+    from neural_compressor import set_random_seed
     set_random_seed(9527)
     if FLAGS.tune:
         from neural_compressor import quantization
@@ -130,7 +130,7 @@ def main(_):
             fit(FLAGS.input_model, conf, b_func=evaluate)
         else:
             from neural_compressor.model import Model
-            model = Model(FLAGS.input_model, backend='keras').model
+            model = Model(FLAGS.input_model, backend='itex').model
             accuracy = evaluate(model)
             print('Batch size = %d' % FLAGS.batch_size)
             print("Accuracy: %.5f" % accuracy)

examples/keras/image_recognition/vgg19/quantization/ptq/main.py

@@ -110,7 +110,7 @@ def eval_func(dataloader, metric):
     return acc
 
 def main(_):
-    from neural_compressor.utils import set_random_seed
+    from neural_compressor import set_random_seed
     set_random_seed(9527)
     if FLAGS.tune:
         from neural_compressor import quantization
@@ -130,7 +130,7 @@ def main(_):
             fit(FLAGS.input_model, conf, b_func=evaluate)
         else:
             from neural_compressor.model import Model
-            model = Model(FLAGS.input_model, backend='keras').model
+            model = Model(FLAGS.input_model, backend='itex').model
             accuracy = evaluate(model)
             print('Batch size = %d' % FLAGS.batch_size)
             print("Accuracy: %.5f" % accuracy)

examples/keras/image_recognition/xception/quantization/ptq/main.py

@@ -118,7 +118,7 @@ def main(_):
     if FLAGS.tune:
         from neural_compressor.quantization import fit
         from neural_compressor.config import PostTrainingQuantConfig
-        from neural_compressor.utils.utility import set_random_seed
+        from neural_compressor import set_random_seed
         set_random_seed(9527)
         config = PostTrainingQuantConfig(backend='itex', 
             calibration_sampling_size=[50, 100])
@@ -138,7 +138,7 @@ def main(_):
             fit(FLAGS.input_model, conf, b_func=evaluate)
         else:
             from neural_compressor.model.model import Model
-            accuracy = evaluate(Model(FLAGS.input_model, backend='keras').model)
+            accuracy = evaluate(Model(FLAGS.input_model, backend='itex').model)
             logger.info('Batch size = %d' % FLAGS.batch_size)
             logger.info("Accuracy: %.5f" % accuracy)
 

examples/onnxrt/nlp/mobilebert/quantization/ptq_static/main.py

@@ -414,7 +414,7 @@ def eval_func(model):
         config = PostTrainingQuantConfig(approach="static",
                                          quant_format=args.quant_format,
                                          recipes={"optypes_to_exclude_output_quant": ["MatMul"]})
-        q_model = quantization.fit(model, 
+        q_model = quantization.fit(model,
                                    config,
                                    eval_func=eval_func,
                                    calib_dataloader=dataloader)

examples/onnxrt/nlp/onnx_model_zoo/gpt2/quantization/ptq_dynamic/gpt2.py

@@ -272,7 +272,7 @@ def eval_func(model):
         accuracy_criterion.relative = 0.11
         config = PostTrainingQuantConfig(approach='dynamic',
                                          accuracy_criterion=accuracy_criterion)
-        q_model = quantization.fit(model, 
+        q_model = quantization.fit(model,
                                    config,
                                    eval_func=eval_func)
         q_model.save(args.output_model)

examples/pytorch/nlp/huggingface_models/text-classification/pruning/eager/run_glue_no_trainer.py

@@ -523,7 +523,7 @@ def preprocess_function(examples):
     # pruner = Pruning(config)
     # pruner.model = model
     # pruner.on_train_begin()
-    from neural_compressor.training import prepare_pruning
+    from neural_compressor.experimental.compression import prepare_pruning
     prepare_pruning(configs, model, optimizer)
 
 

examples/pytorch/nlp/huggingface_models/text-classification/pruning/eager/run_glue_no_trainer_mixed.py

@@ -46,7 +46,8 @@
 )
 from transformers.file_utils import get_full_repo_name
 from transformers.utils.versions import require_version
-from neural_compressor.training import prepare_pruning,WeightPruningConfig
+from neural_compressor import WeightPruningConfig
+from neural_compressor.experimental.compression import prepare_pruning
 
 logger = logging.getLogger(__name__)
 

examples/pytorch/nlp/huggingface_models/text-to-image/quantization/run_diffusion.py

@@ -322,7 +322,7 @@ def b_func(model):
             from neural_compressor.config import BenchmarkConfig
 
             b_conf = BenchmarkConfig(cores_per_instance=4, num_of_instance=1)
-            fit(model, config=b_conf, b_func=b_func)
+            fit(model, conf=b_conf, b_func=b_func)
         if args.accuracy_only:
             setattr(pipe, "unet", model)
             accuracy(pipe, generator, _rows, args)

examples/pytorch/object_detection/ssd_resnet34/quantization/ptq/fx/python/main.py

@@ -623,15 +623,15 @@ def benchmark_func(model):
             from neural_compressor.config import BenchmarkConfig
             from neural_compressor import benchmark
             b_conf = BenchmarkConfig(cores_per_instance=4, num_of_instance=1)
-            benchmark.fit(int8_model, config=b_conf, b_func=benchmark_func)
+            benchmark.fit(int8_model, conf=b_conf, b_func=benchmark_func)
     else:
         if args.accuracy:
             eval_func(raw_model)
         elif args.benchmark:
             from neural_compressor.config import BenchmarkConfig
             from neural_compressor import benchmark
             b_conf = BenchmarkConfig(cores_per_instance=4, num_of_instance=1)
-            benchmark.fit(raw_model, config=b_conf, b_func=benchmark_func)
+            benchmark.fit(raw_model, conf=b_conf, b_func=benchmark_func)
 
     runner.finish()
     lg.DestroyQSL(qsl)

examples/pytorch/object_detection/ssd_resnet34/quantization/ptq/ipex/infer.py

@@ -658,7 +658,7 @@ def coco_eval(model):
         from neural_compressor.config import BenchmarkConfig
         from neural_compressor import benchmark
         b_conf = BenchmarkConfig(cores_per_instance=4, num_of_instance=1)
-        benchmark.fit(ssd_r34, config=b_conf, b_func=coco_eval)
+        benchmark.fit(ssd_r34, conf=b_conf, b_func=coco_eval)
         return
 
 

examples/pytorch/object_detection/ssd_resnet34/quantization/qat/fx/ssd/main.py

@@ -424,7 +424,7 @@ def training_func_for_nc(model, dataloader=None):
             from neural_compressor.config import BenchmarkConfig
             from neural_compressor import benchmark
             b_conf = BenchmarkConfig(cores_per_instance=4, num_of_instance=1)
-            benchmark.fit(new_model, config=b_conf, b_func=eval_func)
+            benchmark.fit(new_model, conf=b_conf, b_func=eval_func)
         return
 
     return False

examples/tensorflow/image_recognition/SavedModel/efficientnet_v2_b0/quantization/ptq/main.py

@@ -100,7 +100,7 @@ def eval_func(dataloader, metric):
 
 class eval_object_detection_optimized_graph(object):
     def run(self):
-        from neural_compressor.utils import set_random_seed
+        from neural_compressor import set_random_seed
         set_random_seed(9527)
         if args.tune:
             from neural_compressor import quantization
@@ -127,7 +127,7 @@ def run(self):
             from neural_compressor.config import BenchmarkConfig
             if args.mode == 'performance':
                 conf = BenchmarkConfig(cores_per_instance=4, num_of_instance=1)
-                fit(model=args.input_graph, config=conf, b_func=evaluate)
+                fit(model=args.input_graph, conf=conf, b_func=evaluate)
             else:
                 from neural_compressor.model import Model
                 model = Model(args.input_graph).model

examples/tensorflow/image_recognition/SavedModel/mobilenet_v1/quantization/ptq/main.py

@@ -93,7 +93,7 @@ def eval_func(dataloader, metric):
 
 class eval_object_detection_optimized_graph(object):
     def run(self):
-        from neural_compressor.utils import set_random_seed
+        from neural_compressor import set_random_seed
         set_random_seed(9527)
         if args.tune:
             from neural_compressor import quantization
@@ -125,7 +125,7 @@ def run(self):
                     'filter': None
                 }
                 eval_dataloader = create_dataloader('tensorflow', dataloader_args)
-                fit(model=args.input_graph, config=conf, b_dataloader=eval_dataloader)
+                fit(model=args.input_graph, conf=conf, b_dataloader=eval_dataloader)
             else:
                 from neural_compressor.model import Model
                 model = Model(args.input_graph).model

examples/tensorflow/image_recognition/SavedModel/mobilenet_v2/quantization/ptq/main.py

@@ -93,7 +93,7 @@ def eval_func(dataloader, metric):
 
 class eval_object_detection_optimized_graph(object):
     def run(self):
-        from neural_compressor.utils import set_random_seed
+        from neural_compressor import set_random_seed
         set_random_seed(9527)
         if args.tune:
             from neural_compressor import quantization
@@ -125,7 +125,7 @@ def run(self):
                     'filter': None
                 }
                 eval_dataloader = create_dataloader('tensorflow', dataloader_args)
-                fit(model=args.input_graph, config=conf, b_dataloader=eval_dataloader)
+                fit(model=args.input_graph, conf=conf, b_dataloader=eval_dataloader)
             else:
                 from neural_compressor.model import Model
                 model = Model(args.input_graph).model

examples/tensorflow/image_recognition/resnet_v2/quantization/qat/main.py

@@ -335,7 +335,7 @@ def evaluate(model):
 
     Args:
         model (tensorflow.Graph_def): The input model graph
-        
+
     Returns:
         accuracy (float): evaluation result, the larger is better.
     """