Improve dynamic batch size solution when calibration and evaluation stage (#530)

PenghuiCheng · web-flow · commit 26b53b1b103e · 2021-11-29T14:40:34.000+08:00
diff --git a/examples/pytorch/eager/image_recognition/efficientnet/conf_efficientnet_b0.yaml b/examples/pytorch/eager/image_recognition/efficientnet/conf_efficientnet_b0.yaml
@@ -26,9 +26,10 @@ quantization:                                        # optional. tuning constrai
         ImageFolder:
           root: /path/to/calibration/dataset         # NOTE: modify to calibration dataset location if needed
       transform:
-        RandomResizedCrop:
-            size: 224
-        RandomHorizontalFlip: {}
+        Resize:
+          size: 256
+        CenterCrop:
+          size: 224
         ToTensor: {}
         Normalize:
             mean: [0.485, 0.456, 0.406]
diff --git a/examples/pytorch/eager/image_recognition/efficientnet/conf_mobilenetv3_rw.yaml b/examples/pytorch/eager/image_recognition/efficientnet/conf_mobilenetv3_rw.yaml
@@ -26,9 +26,10 @@ quantization:                                        # optional. tuning constrai
         ImageFolder:
           root: /path/to/calibration/dataset         # NOTE: modify to calibration dataset location if needed
       transform:
-        RandomResizedCrop:
-            size: 224
-        RandomHorizontalFlip: {}
+        Resize:
+          size: 256
+        CenterCrop:
+          size: 224
         ToTensor: {}
         Normalize:
             mean: [0.485, 0.456, 0.406]
diff --git a/examples/pytorch/eager/image_recognition/imagenet/cpu/prune_and_ptq/ptq_conf.yaml b/examples/pytorch/eager/image_recognition/imagenet/cpu/prune_and_ptq/ptq_conf.yaml
@@ -26,9 +26,10 @@ quantization:                                        # optional. tuning constrai
         ImageFolder:
           root: /path/to/calibration/dataset         # NOTE: modify to calibration dataset location if needed
       transform:
-        RandomResizedCrop:
+        Resize:
+          size: 256
+        CenterCrop:
           size: 224
-        RandomHorizontalFlip: {}
         ToTensor: {}
         Normalize:
           mean: [0.485, 0.456, 0.406]
diff --git a/examples/pytorch/eager/image_recognition/imagenet/cpu/ptq/conf.yaml b/examples/pytorch/eager/image_recognition/imagenet/cpu/ptq/conf.yaml
@@ -26,9 +26,10 @@ quantization:                                        # optional. tuning constrai
         ImageFolder:
           root: /path/to/calibration/dataset         # NOTE: modify to calibration dataset location if needed
       transform:
-        RandomResizedCrop:
+        Resize:
+          size: 256
+        CenterCrop:
           size: 224
-        RandomHorizontalFlip: {}
         ToTensor: {}
         Normalize:
           mean: [0.485, 0.456, 0.406]
diff --git a/examples/pytorch/eager/image_recognition/imagenet/cpu/ptq/conf_dump_tensors.yaml b/examples/pytorch/eager/image_recognition/imagenet/cpu/ptq/conf_dump_tensors.yaml
@@ -26,9 +26,10 @@ quantization:                                        # optional. tuning constrai
         ImageFolder:
           root: /path/to/calibration/dataset         # NOTE: modify to calibration dataset location if needed
       transform:
-        RandomResizedCrop:
+        Resize:
+          size: 256
+        CenterCrop:
           size: 224
-        RandomHorizontalFlip: {}
         ToTensor: {}
         Normalize:
           mean: [0.485, 0.456, 0.406]
diff --git a/examples/pytorch/eager/image_recognition/imagenet/gpu/conf.yaml b/examples/pytorch/eager/image_recognition/imagenet/gpu/conf.yaml
@@ -28,9 +28,10 @@ quantization:                                        # optional. tuning constrai
         ImageFolder:
           root: /path/to/calibration/dataset         # NOTE: modify to calibration dataset location if needed
       transform:
-        RandomResizedCrop:
+        Resize:
+          size: 256
+        CenterCrop:
           size: 224
-        RandomHorizontalFlip: {}
         ToTensor: {}
         Normalize:
           mean: [0.485, 0.456, 0.406]
diff --git a/examples/pytorch/eager/image_recognition/peleenet/conf.yaml b/examples/pytorch/eager/image_recognition/peleenet/conf.yaml
@@ -26,9 +26,10 @@ quantization:                                        # optional. tuning constrai
         ImageFolder:
           root: /path/to/calibration/dataset         # NOTE: modify to calibration dataset location if needed
       transform:
-        RandomResizedCrop:
-            size: 224
-        RandomHorizontalFlip: {}
+        Resize:
+          size: 256
+        CenterCrop:
+          size: 224
         ToTensor: {}
         Normalize:
             mean: [0.485, 0.456, 0.406]
diff --git a/examples/pytorch/eager/image_recognition/resnest/conf.yaml b/examples/pytorch/eager/image_recognition/resnest/conf.yaml
@@ -26,9 +26,10 @@ quantization:                                        # optional. tuning constrai
         ImageFolder:
           root: /path/to/calibration/dataset         # NOTE: modify to calibration dataset location if needed
       transform:
-        RandomResizedCrop:
-            size: 224
-        RandomHorizontalFlip: {}
+        Resize:
+          size: 256
+        CenterCrop:
+          size: 224
         ToTensor: {}
         Normalize:
             mean: [0.485, 0.456, 0.406]
diff --git a/examples/pytorch/eager/image_recognition/se_resnext/conf.yaml b/examples/pytorch/eager/image_recognition/se_resnext/conf.yaml
@@ -26,9 +26,10 @@ quantization:                                        # optional. tuning constrai
         ImageFolder:
           root: /path/to/calibration/dataset         # NOTE: modify to calibration dataset location if needed
       transform:
-        RandomResizedCrop:
-            size: 224
-        RandomHorizontalFlip: {}
+        Resize:
+          size: 256
+        CenterCrop:
+          size: 224
         ToTensor: {}
         Normalize:
             mean: [0.485, 0.456, 0.406]
diff --git a/examples/pytorch/fx/image_recognition/imagenet/cpu/ptq/conf.yaml b/examples/pytorch/fx/image_recognition/imagenet/cpu/ptq/conf.yaml
@@ -26,9 +26,10 @@ quantization:                                        # optional. tuning constrai
         ImageFolder:
           root: /path/to/calibration/dataset         # NOTE: modify to calibration dataset location if needed
       transform:
-        RandomResizedCrop:
+        Resize:
+          size: 256
+        CenterCrop:
           size: 224
-        RandomHorizontalFlip: {}
         ToTensor: {}
         Normalize:
           mean: [0.485, 0.456, 0.406]
diff --git a/examples/pytorch/ipex/image_recognition/imagenet/cpu/ptq/conf_ipex.yaml b/examples/pytorch/ipex/image_recognition/imagenet/cpu/ptq/conf_ipex.yaml
@@ -26,9 +26,10 @@ quantization:                                        # optional. tuning constrai
         ImageFolder:
           root: /path/to/calibration/dataset        # NOTE: modify to calibration dataset location if needed
       transform:
-        RandomResizedCrop:
+        Resize:
+          size: 256
+        CenterCrop:
           size: 224
-        RandomHorizontalFlip: {}
         ToTensor: {}
         Normalize:
           mean: [0.485, 0.456, 0.406]
diff --git a/neural_compressor/adaptor/mxnet.py b/neural_compressor/adaptor/mxnet.py
@@ -22,7 +22,7 @@
 from neural_compressor.adaptor.adaptor import adaptor_registry, Adaptor
 from neural_compressor.adaptor.query import QueryBackendCapability
 from neural_compressor.utils.utility import dump_elapsed_time, LazyImport, singleton, \
-                                            CpuInfo, GLOBAL_STATE, MODE
+                                            GLOBAL_STATE, MODE
 from collections import OrderedDict
 from neural_compressor.adaptor.mxnet_utils.util import *
 from collections import OrderedDict
@@ -57,7 +57,6 @@ def __init__(self, framework_specific_info):
 
         self.ctx = mx.cpu() if framework_specific_info['device'] == 'cpu' else None
         self.benchmark = (GLOBAL_STATE.STATE == MODE.BENCHMARK)
-        self.num_cores_per_sock = CpuInfo().cores_per_socket
         assert self.ctx is not None, 'Unsupported device'
 
     @dump_elapsed_time("Pass quantize model")
@@ -102,23 +101,42 @@ def calib_func(tmp_tune_cfg, dataloader):
                 'framework_specific_info': {'device': self.ctx.device_type}}
             return q_nc_model
 
+        calib_sampling_size = tune_cfg.get('calib_sampling_size', 1)
         if isinstance(dataloader, BaseDataLoader):
             batch_size = dataloader.batch_size
             try:
-                calib_sampling_size = tune_cfg.get('calib_sampling_size', 1)
-                for i in range(self.num_cores_per_sock):
-                    if calib_sampling_size % (self.num_cores_per_sock - i) == 0:
-                        calib_batch_size = self.num_cores_per_sock - i
+                for i in range(batch_size):
+                    if calib_sampling_size % (batch_size - i) == 0:
+                        calib_batch_size = batch_size - i
+                        if i != 0:  # pragma: no cover
+                            logger.warning("Reset `calibration.dataloader.batch_size` field "
+                                           "to {}".format(calib_batch_size) +
+                                           " to make sure the sampling_size is "
+                                           "divisible exactly by batch size")
                         break
                 tmp_iterations = int(math.ceil(calib_sampling_size / calib_batch_size))
                 tmp_tune_cfg = deepcopy(tune_cfg)
                 tmp_tune_cfg['calib_iteration'] = tmp_iterations
                 dataloader.batch(calib_batch_size)
                 return calib_func(tmp_tune_cfg, dataloader)
             except Exception:  # pragma: no cover
-                dataloader.batch(batch_size)
-                return calib_func(tune_cfg, dataloader)
+                logger.warning(
+                    "Fail to forward with batch size={}, set to {} now.".
+                    format(batch_size, 1))
+                tmp_tune_cfg = deepcopy(tune_cfg)
+                tmp_tune_cfg['calib_iteration'] = calib_sampling_size
+                dataloader.batch(1)
+                return calib_func(tmp_tune_cfg, dataloader)
         else:  # pragma: no cover
+            if hasattr(dataloader, 'batch_size') and \
+              calib_sampling_size % dataloader.batch_size != 0:
+                    iter = tune_cfg['calib_iteration']
+                    logger.warning(
+                        "Please Note that calibration sampling size {} \
+                        isn't divisible exactly by batch size {}. \
+                        So the real sampling size is {}.".
+                        format(calib_sampling_size, dataloader.batch_size,
+                               dataloader.batch_size * iter))
             return calib_func(tune_cfg, dataloader)
 
     def _collect_thresholds(self, sym_model, calib_data, calib_tensors, calib_cfg, calib_cache):
@@ -216,14 +234,15 @@ def post_batch(net, batch, outs):
                 metric.update(out, label)
 
         if isinstance(data_x, BaseDataLoader) and not self.benchmark:
-            batch_size = data_x.batch_size
             try:
-                data_x.batch(self.num_cores_per_sock)
                 sym_model, dataloader = prepare_model_data(nc_model, self.ctx, data_x)
                 run_forward(sym_model, self.ctx, dataloader, b_filter(),
                             pre_batch=pre_batch, post_batch=post_batch)
             except Exception:  # pragma: no cover
-                data_x.batch(batch_size)
+                logger.warning(
+                    "Fail to forward with batch size={}, set to {} now.".
+                    format(data_x.batch_size, 1))
+                data_x.batch(1)
                 sym_model, dataloader = prepare_model_data(nc_model, self.ctx, data_x)
                 run_forward(sym_model, self.ctx, dataloader, b_filter(),
                             pre_batch=pre_batch, post_batch=post_batch)
diff --git a/neural_compressor/adaptor/onnxrt.py b/neural_compressor/adaptor/onnxrt.py
@@ -27,7 +27,7 @@
 from neural_compressor.adaptor.adaptor import adaptor_registry, Adaptor
 from neural_compressor.adaptor.query import QueryBackendCapability
 from neural_compressor.utils.utility import LazyImport, dump_elapsed_time, \
-                                            CpuInfo, GLOBAL_STATE, MODE
+                                            GLOBAL_STATE, MODE
 from ..utils.utility import OpPrecisionStatistics
 from ..experimental.data.dataloaders.base_dataloader import BaseDataLoader
 import math
@@ -63,7 +63,6 @@ def __init__(self, framework_specific_info):
         self.fp32_preds_as_label = False
         self.quantize_config = {} # adaptor should know current configs at any time
         self.quantize_params = {} # adaptor should know current params at any time
-        self.num_cores_per_sock = CpuInfo().cores_per_socket
 
     @dump_elapsed_time("Pass quantize model")
     def quantize(self, tune_cfg, model, data_loader, q_func=None):
@@ -98,24 +97,40 @@ def quantize(self, tune_cfg, model, data_loader, q_func=None):
 
         quantize_config = self._cfg_to_quantize_config(tune_cfg)
         iterations = tune_cfg.get('calib_iteration', 1)
+        calib_sampling_size = tune_cfg.get('calib_sampling_size', 1)
         if self.static:
             if isinstance(data_loader, BaseDataLoader):
                 batch_size = data_loader.batch_size
                 try:
-                    calib_sampling_size = tune_cfg.get('calib_sampling_size', 1)
-                    for i in range(self.num_cores_per_sock):
-                        if calib_sampling_size % (self.num_cores_per_sock - i) == 0:
-                            calib_batch_size = self.num_cores_per_sock - i
+                    for i in range(batch_size):
+                        if calib_sampling_size % (batch_size - i) == 0:
+                            calib_batch_size = batch_size - i
+                            if i != 0:  # pragma: no cover
+                                logger.warning("Reset `calibration.dataloader.batch_size` field "
+                                               "to {}".format(calib_batch_size) +
+                                               " to make sure the sampling_size is "
+                                               "divisible exactly by batch size")
                             break
                     tmp_iterations = int(math.ceil(calib_sampling_size / calib_batch_size))
                     data_loader.batch(calib_batch_size)
                     quantize_params = self._get_quantize_params(tmp_model.model, data_loader, \
                                                                 quantize_config, tmp_iterations)
                 except Exception:  # pragma: no cover
-                    data_loader.batch(batch_size)
+                    logger.warning(
+                        "Fail to forward with batch size={}, set to {} now.".
+                        format(batch_size, 1))
+                    data_loader.batch(1)
                     quantize_params = self._get_quantize_params(tmp_model.model, data_loader, \
-                                                                quantize_config, iterations)
+                                                                quantize_config, calib_sampling_size)
             else:  # pragma: no cover
+                if hasattr(data_loader, 'batch_size') and \
+                  calib_sampling_size % data_loader.batch_size != 0:
+                    logger.warning(
+                        "Please Note that calibration sampling size {} \
+                        isn't divisible exactly by batch size {}. \
+                        So the real sampling size is {}.".
+                        format(calib_sampling_size, data_loader.batch_size,
+                               data_loader.batch_size * iterations))
                 quantize_params = self._get_quantize_params(tmp_model.model, data_loader, \
                                                             quantize_config, iterations)
         else:
@@ -633,12 +648,13 @@ def eval_func(dataloader):
                     break
 
         if isinstance(dataloader, BaseDataLoader) and not self.benchmark:
-            batch_size = dataloader.batch_size
             try:
-                dataloader.batch(self.num_cores_per_sock)
                 eval_func(dataloader)
             except Exception:  # pragma: no cover
-                dataloader.batch(batch_size)
+                logger.warning(
+                    "Fail to forward with batch size={}, set to {} now.".
+                    format(dataloader.batch_size, 1))
+                dataloader.batch(1)
                 eval_func(dataloader)
         else:  # pragma: no cover
             eval_func(dataloader)
diff --git a/neural_compressor/adaptor/pytorch.py b/neural_compressor/adaptor/pytorch.py
@@ -675,7 +675,6 @@ def __init__(self, framework_specific_info):
 
         self.fp32_results = []
         self.fp32_preds_as_label = False
-        self.num_cores_per_sock = CpuInfo().cores_per_socket
 
     def calib_func(self, model, dataloader, tmp_iterations, conf=None):
         try:
@@ -695,24 +694,38 @@ def model_calibration(self, q_model, dataloader, iterations=1, conf=None,
                           calib_sampling_size=1):
         assert iterations > 0
 
-        batch_size = dataloader.batch_size
         with torch.no_grad():
             if isinstance(dataloader, BaseDataLoader):
+                batch_size = dataloader.batch_size
                 try:
-                    for i in range(self.num_cores_per_sock):
-                        if calib_sampling_size % (self.num_cores_per_sock - i) == 0:
-                            calib_batch_size = self.num_cores_per_sock - i
+                    for i in range(batch_size):
+                        if calib_sampling_size % (batch_size - i) == 0:
+                            calib_batch_size = batch_size - i
+                            if i != 0:
+                                logger.warning("Reset `calibration.dataloader.batch_size` field "
+                                               "to {}".format(calib_batch_size) +
+                                               " to make sure the sampling_size is "
+                                               "divisible exactly by batch size")
                             break
                     tmp_iterations = int(math.ceil(calib_sampling_size / calib_batch_size))
                     dataloader.batch(calib_batch_size)
                     self.calib_func(q_model, dataloader, tmp_iterations, conf)
                 except Exception:  # pragma: no cover
                     logger.warning(
                         "Fail to forward with batch size={}, set to {} now.".
-                        format(dataloader.batch_size, batch_size))
-                    dataloader.batch(batch_size)
-                    self.calib_func(q_model, dataloader, iterations, conf)
+                        format(batch_size, 1))
+                    dataloader.batch(1)
+                    self.calib_func(q_model, dataloader, calib_sampling_size, conf)
             else:  # pragma: no cover
+                if hasattr(dataloader, 'batch_size') and \
+                  calib_sampling_size % dataloader.batch_size != 0:
+                    logger.warning(
+                        "Please Note that calibration sampling size {} \
+                        isn't divisible exactly by batch size {}. \
+                        So the real sampling size is {}.".
+                        format(calib_sampling_size, dataloader.batch_size,
+                               dataloader.batch_size * iterations))
+
                 self.calib_func(q_model, dataloader, iterations, conf)
 
     def eval_func(self, model, dataloader, postprocess, metric, measurer, iteration, conf=None):
@@ -746,17 +759,18 @@ def eval_func(self, model, dataloader, postprocess, metric, measurer, iteration,
 
     def model_eval(self, model, dataloader, postprocess=None,
                    metric=None, measurer=None, iteration=-1, conf=None):
-        batch_size = dataloader.batch_size
         with torch.no_grad():
             if metric:
                 metric.reset()
             if isinstance(dataloader, BaseDataLoader) and not self.benchmark:
                 try:
-                    dataloader.batch(self.num_cores_per_sock)
                     results = self.eval_func(
                         model, dataloader, postprocess, metric, measurer, iteration, conf)
                 except Exception:  # pragma: no cover
-                    dataloader.batch(batch_size)
+                    logger.warning(
+                        "Fail to forward with batch size={}, set to {} now.".
+                        format(dataloader.batch_size, 1))
+                    dataloader.batch(1)
                     results = self.eval_func(
                         model, dataloader, postprocess, metric, measurer, iteration, conf)
             else:  # pragma: no cover
diff --git a/neural_compressor/adaptor/tensorflow.py b/neural_compressor/adaptor/tensorflow.py
diff --git a/neural_compressor/conf/config.py b/neural_compressor/conf/config.py
diff --git a/neural_compressor/strategy/strategy.py b/neural_compressor/strategy/strategy.py
diff --git a/neural_compressor/utils/create_obj_from_config.py b/neural_compressor/utils/create_obj_from_config.py
