Added Mixed Precision Quantization for decoder

Author: Shiva Chilukamari

sam2.1-hiera-small/QNN/README.md

@@ -19,7 +19,7 @@ Model compilation using QNN Execution Provider requires a Python environment wit
 pip install olive-ai onnxruntime-qnn torch torchvision transformers
 ```
 
-Replace `/path/to/qnn/env/bin` in [sam21_vision_encoder_qnn_ctx.json](sam21_vision_encoder_qnn_ctx.json) and [sam21_mask_decoder_qnn_ctx.json](sam21_mask_decoder_qnn_ctx.json) with the path to the directory containing your QNN environment's Python executable. This path can be found by running the following command in the environment:
+Replace `/path/to/qnn/env/bin` in [sam21_vision_encoder_qnn_ctx.json](sam21_vision_encoder_qnn_ctx.json), [sam21_mask_decoder_qnn_fp16_ctx.json](sam21_mask_decoder_qnn_fp16_ctx.json) and  [sam21_mask_decoder_qnn_mp_ctx.json](sam21_mask_decoder_qnn_mp_ctx.json) with the path to the directory containing your QNN environment's Python executable. This path can be found by running the following command in the environment:
 
 ```bash
 # Linux
@@ -45,9 +45,14 @@ For Encoder Model:
 olive run --config sam21_vision_encoder_qnn_ctx.json
 ```
 
-For Decoder Model:
+For Decoder Model with FP16 quantization:
 ```bash
-olive run --config sam21_mask_decoder_qnn_ctx.json
+olive run --config sam21_mask_decoder_qnn_fp16_ctx.json
+```
+
+For Decoder Model with Mixed Precision quantization:
+```bash
+olive run --config sam21_mask_decoder_qnn_mp_ctx.json
 ```
 
 > ⚠️ If optimization fails during context binary generation, rerun the command. The process will resume from the last completed step.

sam2.1-hiera-small/QNN/config.py

@@ -10,5 +10,5 @@ class ModelConfig:
     ve_input_name = "pixel_values"
     ve_channel_size = 3
     ve_sample_size = 1024
-    md_input_names = ("image_embeddings", "high_res_features1", "high_res_features2", "coords.1", "labels")
-    ms_input_shapes = ((1, 256, 64, 64), (1, 32, 256, 256), (1, 64, 128, 128), (1, 5, 2), (1, 5))
+    md_input_names = ("image_embeddings", "high_res_features1", "high_res_features2", "point_coords", "point_labels", "mask_input", "has_mask_input")
+    ms_input_shapes = ((1, 256, 64, 64), (1, 32, 256, 256), (1, 64, 128, 128), (1, 5, 2), (1, 5), (1, 1, 256, 256), [1])

sam2.1-hiera-small/QNN/info.yml

@@ -3,7 +3,12 @@ keywords:
   qnn
 arch: hiera
 recipes:
-  - file: "sam21_mask_decoder_qnn_ctx.json"
+  - file: "sam21_mask_decoder_qnn_fp16_ctx.json"
+    device:
+    - npu
+    ep: QNNExecutionProvider
+    
+  - file: "sam21_mask_decoder_qnn_mp_ctx.json"
     device:
     - npu
     ep: QNNExecutionProvider

…mall/QNN/sam21_mask_decoder_qnn_ctx.json …QNN/sam21_mask_decoder_qnn_fp16_ctx.jsonsam2.1-hiera-small/QNN/sam21_mask_decoder_qnn_ctx.json renamed to sam2.1-hiera-small/QNN/sam21_mask_decoder_qnn_fp16_ctx.json sam2.1-hiera-small/QNN/sam21_mask_decoder_qnn_ctx.json renamed to sam2.1-hiera-small/QNN/sam21_mask_decoder_qnn_fp16_ctx.json

@@ -21,23 +21,14 @@
     },
     "data_configs": [
         {
-            "name": "dummy_data_config",
-            "type": "DummyDataContainer",
+            "name": "latency_data_config",
+            "user_script": "user_script.py",
             "load_dataset_config": {
-                "input_shapes": [
-                    [
-                        1,
-                        3,
-                        1024,
-                        1024
-                    ]
-                ],
-                "input_names": [
-                    "input.1"
-                ],
-                "input_types": [
-                    "float32"
-                ]
+                "type": "local_dataset"
+            },
+            "dataloader_config": {
+                "type": "md_data_loader",
+                "batch_size": 1
             }
         }
     ],
@@ -62,7 +53,7 @@
     "log_severity_level": 0,
     "ort_log_severity_level": 1,
     "ort_py_log_severity_level": 1,
-    "cache_dir": "cache_decoder",
-    "output_dir": "model/decoder",
+    "cache_dir": "cache_decoder_fp16",
+    "output_dir": "model/decoder_fp16",
     "no_artifacts": true
 }

sam2.1-hiera-small/QNN/sam21_mask_decoder_qnn_mp_ctx.json

@@ -0,0 +1,110 @@
+{
+    "input_model": {
+        "type": "ONNXModel",
+        "model_path": "sam21_mask_decoder.onnx"
+    },
+    "systems": {
+        "local_system": {
+            "type": "LocalSystem",
+            "accelerators": [
+                {
+                    "device": "gpu",
+                    "execution_providers": [
+                        "CUDAExecutionProvider"
+                    ]
+                }
+            ]
+        },
+        "target_system": {
+            "type": "PythonEnvironment",
+            "python_environment_path": "/path/to/qnn/env/bin",
+            "accelerators": [
+                {
+                    "execution_providers": [
+                        "QNNExecutionProvider"
+                    ]
+                }
+            ]
+        }
+    },
+    "data_configs": [
+        {
+            "name": "latency_data_config",
+            "user_script": "user_script.py",
+            "load_dataset_config": {
+                "type": "local_dataset"
+            },
+            "dataloader_config": {
+                "type": "md_data_loader",
+                "batch_size": 1
+            }
+        },
+        {
+            "name": "quantize_data_config",
+            "user_script": "user_script.py",
+            "load_dataset_config": {
+                "type": "local_dataset"
+            },
+            "dataloader_config": {
+                "type": "md_quantize_data_loader",
+                "data_num": 200,
+                "point_p": 0.3,
+                "mask_p": 0.2
+            }
+        }
+    ],
+    "passes": {
+        "f16": {
+            "type": "OnnxFloatToFloat16",
+            "keep_io_types": true,
+            "save_as_external_data": true,
+            "node_include_list": ["/Add", "/Concat", "/Concat_1", "/Gather", 
+                "/Div", "/Reshape", "/ScatterND", "/Gather_2", 
+                "/Div_1", "/Reshape_1", "/ScatterND_1", "/Mul_6", 
+                "/Sub", "/MatMul", "/Mul_7", "/Sin", "/Cos", 
+                "/Concat_6", "/Equal_6", "/Unsqueeze_6", "/Where_6", 
+                "/Equal_7", "/Unsqueeze_7", "/Add_1", "/Where_7", 
+                "/Equal_8", "/Unsqueeze_8", "/Add_2", "/Where_8", 
+                "/Equal_9", "/Unsqueeze_9", "/Add_3", "/Where_9", 
+                "/Equal_10", "/Unsqueeze_10", "/Add_4", "/Where_10"]
+        },
+        "sq": {
+            "type": "OnnxStaticQuantization",
+            "data_config": "quantize_data_config",
+            "activation_type": "uint8",
+            "precision": "uint8",
+            "calibration_providers": [ "CUDAExecutionProvider" ],
+            "calibrate_method": "MinMax",
+            "quant_preprocess": true,
+            "op_types_to_exclude": ["Cast"],
+            "save_as_external_data": true,
+            "nodes_to_exclude":["/Add", "/Concat", "/Concat_1", "/Gather", 
+                "/Div", "/Reshape", "/ScatterND", "/Gather_2", 
+                "/Div_1", "/Reshape_1", "/ScatterND_1", "/Mul_6", 
+                "/Sub", "/MatMul", "/Mul_7", "/Sin", "/Cos", 
+                "/Concat_6", "/Equal_6", "/Unsqueeze_6", "/Where_6", 
+                "/Equal_7", "/Unsqueeze_7", "/Add_1", "/Where_7", 
+                "/Equal_8", "/Unsqueeze_8", "/Add_2", "/Where_8", 
+                "/Equal_9", "/Unsqueeze_9", "/Add_3", "/Where_9", 
+                "/Equal_10", "/Unsqueeze_10", "/Add_4", "/Where_10"]
+        },
+        "cb": {
+            "type": "EPContextBinaryGenerator",
+            "provider_options": {
+                "htp_performance_mode": "burst",
+                "htp_graph_finalization_optimization_mode": "3",
+                "offload_graph_io_quantization": "0",
+                "soc_model": "60"
+            },
+            "weight_sharing": false
+        }
+    },
+    "host": "local_system",
+    "target": "target_system",
+    "log_severity_level": 0,
+    "ort_log_severity_level": 1,
+    "ort_py_log_severity_level": 1,
+    "cache_dir": "cache_decoder_mp",
+    "output_dir": "model/decoder_mp",
+    "no_artifacts": true
+}

sam2.1-hiera-small/QNN/user_script.py

@@ -52,6 +52,21 @@ def __init__(self, total):
         self.finish_load()
 
 
+class MdDecoderGeneratedDataLoader(BaseDataLoader):
+    def __init__(self, total, point_p, mask_p):
+        super().__init__(total)
+        md_generate_quant_data(total, point_p, mask_p)
+        self.data_files = [
+            os.path.join(ModelConfig.data_dir, f.name)
+            for f in os.scandir(ModelConfig.data_dir)
+            if "points.npz" in f.name
+        ]
+        self.data_files.sort()
+        for f in self.data_files:
+            self.load(f)
+        self.finish_load()
+
+
 class RandomDataLoader:
     def __init__(self, create_inputs_func, batch_size, torch_dtype):
         self.create_input_func = create_inputs_func
@@ -102,6 +117,11 @@ def md_data_loader(dataset, batch_size, *args, **kwargs):
     return RandomDataLoader(md_inputs, batch_size, torch.float32)
 
 
+@Registry.register_dataloader()
+def md_quantize_data_loader(dataset, data_num, point_p, mask_p, *args, **kwargs):
+    return MdDecoderGeneratedDataLoader(data_num, point_p, mask_p)
+
+    
 def ve_generate_quant_data(num_samples):
     p = Path(ModelConfig.data_dir)
     if p.is_dir() and (len([f for f in p.glob("*images.npz")]) >= num_samples):
@@ -118,3 +138,79 @@ def ve_generate_quant_data(num_samples):
         inputs = processor(image, return_tensors="pt")
         pixel_values = inputs["pixel_values"].detach().cpu().numpy()
         np.savez(f"{ModelConfig.data_dir}/input_{i}_images.npz", input=pixel_values)
+
+
+def get_inputs(sample, point_p = 0.0, mask_p = 0.0):
+    inputs = {}
+    segments_info = sample['segments_info']
+    segment_info = np.random.choice(segments_info)
+    box = segment_info['bbox']
+    p1 = [box[0], box[1]]
+    p2 = [p1[0] + box[2], p1[1] + box[3]]
+    p = np.mean([p1, p2], axis = 0)
+    
+    if np.random.random() < point_p:
+        inputs['point_coords'] = np.concatenate([[p], np.zeros((4, 2))])[None, :]
+        inputs['point_labels'] = np.concatenate([[1], -np.ones(4)])[None, :]
+    else:
+        inputs['point_coords'] = np.concatenate([[p1, p2], np.zeros((3, 2))])[None, :]
+        inputs['point_labels'] = np.concatenate([[2, 3], -np.ones(3)])[None, :]
+
+    if np.random.random() < mask_p:
+        pil_mask = sample['label']
+        w, h = pil_mask.size
+        masks = np.array(pil_mask.resize((256, 256)))
+        mask_point = masks[int(p[1]*256/h), int(p[0]*256/w)]
+        mask = (masks == mask_point).all(axis = -1)
+        inputs['mask_input'] = mask[None, None, :]
+        inputs['has_mask_input'] = np.array([1])
+    else:
+        inputs['mask_input'] = np.zeros((1, 1, 256, 256))
+        inputs['has_mask_input'] = np.array([0])
+    return inputs
+
+
+def md_generate_quant_data(num_samples, point_p, mask_p):
+    p = Path(ModelConfig.data_dir)
+    if p.is_dir() and (len([f for f in p.glob("*points.npz")]) >= num_samples):
+        return
+    from hydra import initialize
+    from hydra.core.global_hydra import GlobalHydra
+    from sam2.build_sam import build_sam2
+    from generate_model import model_weights_url, checkpoint, model_config_url, model_cfg
+    from generate_model import download_file, SAM2Encoder
+
+    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
+    
+
+    processor = Sam2Processor.from_pretrained(ModelConfig.model_name)
+    dataset = load_dataset("nielsr/coco-panoptic-val2017")
+    dataset = dataset["train"]
+    os.makedirs(ModelConfig.data_dir, exist_ok=True)
+    
+    download_file(model_weights_url, checkpoint)
+    download_file(model_config_url, model_cfg)
+    
+    GlobalHydra.instance().clear()
+    initialize(config_path="./", job_name="sam2_inference", version_base=None)
+    sam2_model = build_sam2(model_cfg, checkpoint, device="cpu")
+    
+    encoder = SAM2Encoder(sam2_model).to(device)
+    
+    for i, sample in enumerate(dataset):
+        if i >= num_samples:
+            break
+        image = sample['image']
+        inputs = get_inputs(sample, point_p, mask_p)
+        process_inputs = processor(image, input_points = [inputs['point_coords']], input_labels = [inputs['point_labels']], return_tensors="pt")
+        image_embed, high_feat_1, high_feat2 = encoder(input = process_inputs['pixel_values'].to(device))
+        quant_inputs = {}
+        quant_inputs['image_embeddings'] = image_embed.detach().cpu().numpy().astype(np.float32)
+        quant_inputs['high_res_features1'] = high_feat_1.detach().cpu().numpy().astype(np.float32)
+        quant_inputs['high_res_features2'] = high_feat2.detach().cpu().numpy().astype(np.float32)
+        quant_inputs['point_coords'] = process_inputs['input_points'].detach().cpu().numpy()[0].astype(np.float32)
+        quant_inputs['point_labels'] = process_inputs['input_labels'].detach().cpu().numpy()[0].astype(np.float32)
+        quant_inputs['mask_input'] = inputs['mask_input'].astype(np.float32)
+        quant_inputs['has_mask_input'] = inputs['has_mask_input'].astype(np.float32)
+
+        np.savez(f"{ModelConfig.data_dir}/input_{i}_points.npz", **quant_inputs)