Add AutoML-based mixed-precision initialization mode - AutoQ (openvin…

…otoolkit#250)

* Adaptation of MIT HAN Lab's HAQ: Hardware-Aware Automated Quantization with Mixed Precision

* Introduce a Deep Reinforcement Learning algorithm (DDPG) to learn and
  initialize layer-wise quantization bitwidth, prior to NNCF quantize-aware fine-tuning

* The mixed-precision initialization is optimized towards minimal accuracy drop given
  a user-specified model size constraint

* Supported precision depends on target HW (VPU 8/4/2) or user-specified precision space

examples/classification/configs/mixed_precision/mobilenet_v2_imagenet_mixed_int_autoq.json

@@ -0,0 +1,50 @@
+{
+    "model": "mobilenet_v2",
+    "pretrained": true,
+    "input_info": {
+      "sample_size": [2, 3, 224, 224]
+    },
+    "num_classes": 1000,
+    "batch_size" : 128,
+    "epochs": 5,
+    "optimizer": {
+        "type": "Adam",
+        "base_lr": 0.00001,
+        "schedule_type": "multistep",
+        "steps": [
+            5
+        ]
+    },
+    "target_device": "TRIAL",
+    "compression": {
+        "algorithm": "quantization",
+        "weights": {
+            "mode": "asymmetric",
+            "bits": 8,
+            "per_channel": true
+        },
+        "activations": {
+            "mode": "asymmetric",
+            "bits": 8,
+            "per_channel": false
+        },
+        "initializer": {
+            "batchnorm_adaptation": {
+                "num_bn_adaptation_samples": 4000
+            },
+            "range":
+            {
+                "type": "mean_min_max",
+                "num_init_samples": 1500
+            },
+            "precision": {
+                "type": "autoq",
+                "bits": [2, 4, 8],
+                "iter_number": 600,
+                "compression_ratio": 0.15,
+                "eval_subset_ratio": 0.20,
+                "dump_init_precision_data": true
+            }
+        }
+    }
+}

examples/classification/configs/mixed_precision/resnet50_imagenet_mixed_int_autoq.json

@@ -0,0 +1,57 @@
+{
+    "model": "resnet50",
+    "pretrained": true,
+    "input_info": {
+        "sample_size": [
+            1,
+            3,
+            224,
+            224
+        ]
+    },
+    "num_classes": 1000,
+    "batch_size" : 128,
+    "epochs": 30,
+    "optimizer": {
+        "base_lr": 0.00031,
+        "schedule_type": "plateau",
+        "type": "Adam",
+        "scheduler_params": {
+            "threshold": 0.1,
+            "cooldown": 3
+        },
+        "weight_decay": 1e-05
+    },
+    "target_device": "TRIAL",
+    "compression": {
+        "algorithm": "quantization",
+        "weights": {
+            "mode": "symmetric",
+            "bits": 8,
+            "per_channel": true
+        },
+        "activations": {
+            "mode": "asymmetric",
+            "bits": 8,
+            "per_channel": false
+        },
+        "initializer": {
+            "batchnorm_adaptation": {
+                "num_bn_adaptation_samples": 4000
+            },
+            "range":
+            {
+                "type": "mean_min_max",
+                "num_init_samples": 1500
+            },
+            "precision": {
+                "type": "autoq",
+                "bits": [2, 4, 8],
+                "iter_number": 600,
+                "compression_ratio": 0.15,
+                "eval_subset_ratio": 0.20,
+                "dump_init_precision_data": true
+            }
+        }
+    }
+}

examples/classification/main.py

@@ -140,7 +140,13 @@ def main_worker(current_gpu, config: SampleConfig):
         train_dataset, val_dataset = create_datasets(config)
         train_loader, train_sampler, val_loader, init_loader = create_data_loaders(config, train_dataset, val_dataset)
 
-        nncf_config = register_default_init_args(nncf_config, init_loader, criterion, train_criterion_fn, config.device)
+        def autoq_eval_fn(model, eval_loader):
+            _, top5 = validate(eval_loader, model, criterion, config)
+            return top5
+
+        nncf_config = register_default_init_args(
+            nncf_config, init_loader, criterion, train_criterion_fn,
+            autoq_eval_fn, val_loader, config.device)
 
     # create model
     model = load_model(model_name,

examples/classification/staged_quantization_worker.py

@@ -133,7 +133,14 @@ def staged_quantization_main_worker(current_gpu, config):
         # Data loading code
         train_dataset, val_dataset = create_datasets(config)
         train_loader, train_sampler, val_loader, init_loader = create_data_loaders(config, train_dataset, val_dataset)
-        nncf_config = register_default_init_args(nncf_config, init_loader, criterion, train_criterion_fn, config.device)
+
+        def autoq_eval_fn(model, eval_loader):
+            _, top5 = validate(eval_loader, model, criterion, config)
+            return top5
+
+        nncf_config = register_default_init_args(
+            nncf_config, init_loader, criterion, train_criterion_fn,
+            autoq_eval_fn, val_loader, config.device)
 
     # create model
     model_name = config['model']

examples/common/utils.py

@@ -33,6 +33,7 @@
 from torch.utils.tensorboard import SummaryWriter
 from texttable import Texttable
 import mlflow
+import torch
 
 from examples.common.example_logger import logger as default_logger
 from nncf.utils import is_main_process
@@ -150,6 +151,9 @@ def configure_device(current_gpu, config: SampleConfig):
 
     config.device = get_device(config)
 
+    if config.execution_mode == ExecutionMode.SINGLE_GPU:
+        torch.cuda.set_device(config.current_gpu)
+
 
 def configure_logging(sample_logger, config):
     config.tb = SummaryWriter(config.log_dir)

examples/object_detection/main.py

@@ -140,7 +140,14 @@ def criterion_fn(model_outputs, target, criterion):
             loss_l, loss_c = criterion(model_outputs, target)
             return loss_l + loss_c
 
-        nncf_config = register_default_init_args(nncf_config, init_data_loader, criterion, criterion_fn, config.device)
+        def autoq_test_fn(model, eval_loader):
+            # RL is maximization, change the loss polarity
+            return -1 * test_net(model, config.device, eval_loader, distributed=config.distributed,
+                                 loss_inference=True, criterion=criterion)
+
+        nncf_config = register_default_init_args(
+            nncf_config, init_data_loader, criterion, criterion_fn,
+            autoq_test_fn, test_data_loader, config.device)
 
     ##################
     # Prepare model

examples/semantic_segmentation/main.py

@@ -50,15 +50,15 @@
 from nncf.utils import is_main_process
 
 
-def get_arguments(args):
+def get_arguments_parser():
     parser = get_common_argument_parser()
     parser.add_argument(
         "--dataset",
         help="Dataset to use.",
         choices=["camvid", "cityscapes", "mapillary"],
         default=None
     )
-    return parser.parse_args(args=args)
+    return parser
 
 
 def get_preprocessing_transforms(config):
@@ -436,6 +436,7 @@ def test(model, test_loader, criterion, class_encoding, config):
             gt_labels = center_crop(gt_labels, outputs_size_hw).contiguous()
         data_utils.show_ground_truth_vs_prediction(images, gt_labels, color_predictions, class_encoding)
 
+    return miou
 
 def predict(model, images, class_encoding, config):
     images = images.to(config.device)
@@ -489,7 +490,13 @@ def criterion_fn(model_outputs, target, criterion_):
         loaders, w_class = load_dataset(dataset, config)
         train_loader, val_loader, init_loader = loaders
         criterion = get_criterion(w_class, config)
-        nncf_config = register_default_init_args(nncf_config, init_loader, criterion, criterion_fn, config.device)
+
+        def autoq_test_fn(model, eval_loader):
+            return test(model, eval_loader, criterion, color_encoding, config)
+
+        nncf_config = register_default_init_args(
+            nncf_config, init_loader, criterion, criterion_fn,
+            autoq_test_fn, val_loader, config.device)
 
     model = load_model(config.model,
                        pretrained=pretrained,
@@ -537,7 +544,7 @@ def criterion_fn(model_outputs, target, criterion_):
 
 
 def main(argv):
-    parser = get_common_argument_parser()
+    parser = get_arguments_parser()
     arguments = parser.parse_args(args=argv)
     config = create_sample_config(arguments, parser)
     if arguments.dist_url == "env://":