diff --git a/backends/qualcomm/_passes/__init__.py b/backends/qualcomm/_passes/__init__.py index 9c884d7ab93..6d50bbb51b2 100644 --- a/backends/qualcomm/_passes/__init__.py +++ b/backends/qualcomm/_passes/__init__.py @@ -9,6 +9,7 @@ from .annotate_unbind import AnnotateUnbind from .convert_bmm_to_matmul import ConvertBmmToMatmul from .convert_conv1d_to_conv2d import ConvertConv1dToConv2d +from .convert_upsample_bicubic2d import ConvertUpsampleBicubicWithBilinear from .decompose_any import DecomposeAny from .decompose_einsum import DecomposeEinsum from .decompose_expm1 import DecomposeExpM1 @@ -40,6 +41,7 @@ ConvertBmmToMatmul, ConvertConv1dToConv2d, DecomposeAny, + ConvertUpsampleBicubicWithBilinear, DecomposeEinsum, DecomposeExpM1, DecomposeLinalgVectorNorm, diff --git a/backends/qualcomm/_passes/convert_upsample_bicubic2d.py b/backends/qualcomm/_passes/convert_upsample_bicubic2d.py new file mode 100644 index 00000000000..367e9155c77 --- /dev/null +++ b/backends/qualcomm/_passes/convert_upsample_bicubic2d.py @@ -0,0 +1,27 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from executorch.exir.dialects._ops import ops as exir_ops +from executorch.exir.pass_base import ExportPass + + +class ConvertUpsampleBicubicWithBilinear(ExportPass): + """ + Qnn does not support bicubic interpolation, so we need to convert it to bilinear. + This pass will convert bicubic interpolation to bilinear interpolation. + """ + + bicubic_op_targets = { + exir_ops.edge.aten.upsample_bicubic2d.vec, + } + upsample_bilinear_op = exir_ops.edge.aten.upsample_bilinear2d.default + + def __init__(self): + super(ConvertUpsampleBicubicWithBilinear, self).__init__() + + def call_operator(self, op, args, kwargs, meta): + if op not in self.bicubic_op_targets: + return super().call_operator(op, args, kwargs, meta) + return super().call_operator(self.upsample_bilinear_op, args[:-1], kwargs, meta) diff --git a/backends/qualcomm/_passes/layout_transform.py b/backends/qualcomm/_passes/layout_transform.py index 4d47c38bc03..19c5417f8f8 100644 --- a/backends/qualcomm/_passes/layout_transform.py +++ b/backends/qualcomm/_passes/layout_transform.py @@ -55,6 +55,7 @@ class LayoutTransform(ExportPass): exir_ops.edge.aten.ceil.default, exir_ops.edge.aten.clamp.default, exir_ops.edge.aten.constant_pad_nd.default, + exir_ops.edge.aten.cumsum.default, exir_ops.edge.aten.div.Tensor, exir_ops.edge.aten.elu.default, exir_ops.edge.aten.eq.Tensor, diff --git a/backends/qualcomm/_passes/qnn_pass_manager.py b/backends/qualcomm/_passes/qnn_pass_manager.py index ab2c86102df..ed3c8eb217e 100644 --- a/backends/qualcomm/_passes/qnn_pass_manager.py +++ b/backends/qualcomm/_passes/qnn_pass_manager.py @@ -14,6 +14,7 @@ AnnotateUnbind, ConvertBmmToMatmul, ConvertConv1dToConv2d, + ConvertUpsampleBicubicWithBilinear, DecomposeAny, DecomposeEinsum, DecomposeExpM1, @@ -74,6 +75,7 @@ def get_capture_program_passes(): (AnnotateUnbind, True), (ConvertBmmToMatmul, True), (ConvertConv1dToConv2d, True), + (ConvertUpsampleBicubicWithBilinear, False), (DecomposeAny, True), (ExpandBroadcastTensorShape, False), (FixedLinearKeepDim, True), diff --git a/backends/qualcomm/_passes/recompose_pixel_unshuffle.py b/backends/qualcomm/_passes/recompose_pixel_unshuffle.py index 7aac4fb823e..81214facc3a 100644 --- a/backends/qualcomm/_passes/recompose_pixel_unshuffle.py +++ b/backends/qualcomm/_passes/recompose_pixel_unshuffle.py @@ -45,13 +45,11 @@ def call(self, graph_module: torch.fx.GraphModule): continue view_node = premute_node.args[0] - if any( - [ - view_node.op != "call_function", - view_node.target != self.view_target, - len(view_node.args[1]) != 6, - len(premute_node.args[1]) != 6, - ] + if ( + view_node.op != "call_function" + or view_node.target != self.view_target + or len(view_node.args[1]) != 6 + or len(premute_node.args[1]) != 6 ): continue diff --git a/backends/qualcomm/_passes/utils.py b/backends/qualcomm/_passes/utils.py index a8eb6b192ee..6b7dc8c16fa 100755 --- a/backends/qualcomm/_passes/utils.py +++ b/backends/qualcomm/_passes/utils.py @@ -78,6 +78,7 @@ def get_passes_dependency_for_capture_program(): AnnotateUnbind, ConvertBmmToMatmul, ConvertConv1dToConv2d, + ConvertUpsampleBicubicWithBilinear, DecomposeAny, DecomposeLinalgVectorNorm, ExpandBroadcastTensorShape, @@ -96,18 +97,20 @@ def get_passes_dependency_for_capture_program(): AnnotateQuantAttrs: [ RecomposePixelUnshuffle, ConvertBmmToMatmul, + ConvertUpsampleBicubicWithBilinear, RemoveRedundancy, ], AnnotateStack: [RemoveRedundancy], AnnotateUnbind: [RemoveRedundancy], ConvertBmmToMatmul: [RecomposePixelUnshuffle], ConvertConv1dToConv2d: [FoldQDQ], + ConvertUpsampleBicubicWithBilinear: [RemoveRedundancy], DecomposeAny: [RemoveRedundancy], DecomposeLinalgVectorNorm: [RemoveRedundancy], ExpandBroadcastTensorShape: [RemoveRedundancy], FixedLinearKeepDim: [FoldQDQ], FoldQDQ: [AnnotateQuantAttrs, AnnotateStack, AnnotateUnbind], - I64toI32: [RemoveRedundancy], + I64toI32: [ConvertUpsampleBicubicWithBilinear, RemoveRedundancy], LayoutTransform: [ AnnotateQuantAttrs, ConvertConv1dToConv2d, diff --git a/backends/qualcomm/builders/__init__.py b/backends/qualcomm/builders/__init__.py index 645b823d0e5..705d5d163cd 100644 --- a/backends/qualcomm/builders/__init__.py +++ b/backends/qualcomm/builders/__init__.py @@ -21,6 +21,7 @@ op_clamp, op_conv2d, op_cos, + op_cum_sum, op_depth_to_space, op_dequantize, op_div, @@ -108,6 +109,7 @@ op_clamp, op_conv2d, op_cos, + op_cum_sum, op_depth_to_space, op_dequantize, op_div, diff --git a/backends/qualcomm/builders/op_cos.py b/backends/qualcomm/builders/op_cos.py index 3858a947d93..589bf3ef88e 100644 --- a/backends/qualcomm/builders/op_cos.py +++ b/backends/qualcomm/builders/op_cos.py @@ -3,7 +3,6 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. - from typing import Dict import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper diff --git a/backends/qualcomm/builders/op_cum_sum.py b/backends/qualcomm/builders/op_cum_sum.py new file mode 100644 index 00000000000..f62485bc519 --- /dev/null +++ b/backends/qualcomm/builders/op_cum_sum.py @@ -0,0 +1,84 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from typing import cast, Dict + +import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper + +import numpy as np +import torch +from executorch.backends.qualcomm.utils.constants import QCOM_AXIS_ORDER, QCOM_DATA + +from .node_visitor import NodeVisitor, register_node_visitor +from .qnn_constants import OpCumulativeSum, QNN_OP_PACKAGE_NAME_QTI_AISW + + +@register_node_visitor +class CumulativeSum(NodeVisitor): + target = ["aten.cumsum.default"] + + def __init__(self, *args) -> None: + super().__init__(*args) + + def get_param(self, node, input_tensor): + dim = node.args[1] + + if dim < 0: + dim = dim % len(input_tensor.shape) + if QCOM_AXIS_ORDER in node.meta: + dim = node.meta[QCOM_AXIS_ORDER].index(dim) + + return cast(np.uint32, dim) + + def define_node( + self, + node: torch.fx.Node, + nodes_to_wrappers: Dict[torch.fx.Node, PyQnnWrapper.TensorWrapper], + ) -> PyQnnWrapper.PyQnnOpWrapper: + input_node = node.args[0] + input_tensor = self.get_tensor(input_node, node) + input_tensor_wrapper = self.define_tensor( + input_node, + node, + input_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + + dim = self.get_param(node, input_tensor) + + output_tensor = self.get_tensor(node, node) + output_tensor_wrapper = self.define_tensor( + node, + node, + output_tensor, + PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE, + nodes_to_wrappers, + ) + + cumsum_op = PyQnnWrapper.PyQnnOpWrapper( + node.name, + QNN_OP_PACKAGE_NAME_QTI_AISW, + OpCumulativeSum.op_name, + ) + cumsum_op.AddInputTensors([input_tensor_wrapper]) + cumsum_op.AddOutputTensors([output_tensor_wrapper]) + cumsum_op.AddScalarParam( + OpCumulativeSum.param_axis, + PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_UINT_32, + {QCOM_DATA: dim}, + ) + cumsum_op.AddScalarParam( + OpCumulativeSum.param_exclusive, + PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_BOOL_8, + {QCOM_DATA: False}, + ) + cumsum_op.AddScalarParam( + OpCumulativeSum.param_reverse, + PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_BOOL_8, + {QCOM_DATA: False}, + ) + + return cumsum_op diff --git a/backends/qualcomm/builders/op_sin.py b/backends/qualcomm/builders/op_sin.py index 89fce6bee9c..8828685ac9e 100644 --- a/backends/qualcomm/builders/op_sin.py +++ b/backends/qualcomm/builders/op_sin.py @@ -3,7 +3,6 @@ # # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. - from typing import Dict import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper diff --git a/backends/qualcomm/builders/op_stack.py b/backends/qualcomm/builders/op_stack.py index 616d0ee0ccc..fdef148ad4d 100644 --- a/backends/qualcomm/builders/op_stack.py +++ b/backends/qualcomm/builders/op_stack.py @@ -51,7 +51,7 @@ def define_node( dim = 0 if len(node.args) == 1 else cast(int, node.args[1]) if dim < 0: - dim = dim % len(input_tensor.shape) + dim = dim % len(output_tensor.shape) if QCOM_AXIS_ORDER in node.meta: dim = node.meta[QCOM_AXIS_ORDER].index(dim) stack_op = PyQnnWrapper.PyQnnOpWrapper( diff --git a/backends/qualcomm/builders/qnn_constants.py b/backends/qualcomm/builders/qnn_constants.py index 31822a174b9..6398fbba1ed 100644 --- a/backends/qualcomm/builders/qnn_constants.py +++ b/backends/qualcomm/builders/qnn_constants.py @@ -57,6 +57,14 @@ class OpConvert: op_name: str = "Convert" +@dataclass(init=False, frozen=True) +class OpCumulativeSum: + op_name = "CumulativeSum" + param_axis = "axis" + param_exclusive = "exclusive" + param_reverse = "reverse" + + @dataclass(init=False, frozen=True) class OpDepthToSpace: op_name: str = "DepthToSpace" diff --git a/backends/qualcomm/partition/common_defs.py b/backends/qualcomm/partition/common_defs.py index b427c59ce07..6326f4d1210 100644 --- a/backends/qualcomm/partition/common_defs.py +++ b/backends/qualcomm/partition/common_defs.py @@ -13,6 +13,7 @@ exir_ops.edge.aten.clone.default, exir_ops.edge.aten.slice_scatter.default, exir_ops.edge.aten.copy.default, + exir_ops.edge.aten.upsample_bicubic2d.vec, exir_ops.edge.quantized_decomposed.embedding_4bit.dtype, ] diff --git a/backends/qualcomm/partition/utils.py b/backends/qualcomm/partition/utils.py index 1e2b17b2a69..6931e35e6e3 100644 --- a/backends/qualcomm/partition/utils.py +++ b/backends/qualcomm/partition/utils.py @@ -39,6 +39,7 @@ def get_skip_decomp_table() -> List[torch._ops.OperatorBase]: torch.ops.aten.rms_norm.default, torch.ops.aten._safe_softmax.default, torch.ops.aten.stack.default, + torch.ops.aten.upsample_bicubic2d.vec, # This request is ignored because it is in a blocklist. Refer to exir/program/_program.py # torch.ops.aten.unbind.int, torch.ops.pt2e_quant.quantize_affine.default, diff --git a/backends/qualcomm/quantizer/annotators.py b/backends/qualcomm/quantizer/annotators.py index 52662202795..f59df46b3c1 100644 --- a/backends/qualcomm/quantizer/annotators.py +++ b/backends/qualcomm/quantizer/annotators.py @@ -976,6 +976,11 @@ def annotate_conv2d(node: Node, quantization_config: QuantizationConfig) -> None ) +@register_annotator([torch.ops.aten.cumsum.default]) +def annotate_cumsum(node: Node, quantization_config: QuantizationConfig) -> None: + annotate_single_in_single_out(node, quantization_config) + + @register_annotator([torch.ops.aten.linear.default]) def annotate_linear(node: Node, quantization_config: QuantizationConfig) -> None: act_node = node.args[0] diff --git a/backends/qualcomm/tests/models.py b/backends/qualcomm/tests/models.py index f1171b129e6..6570b4befcb 100644 --- a/backends/qualcomm/tests/models.py +++ b/backends/qualcomm/tests/models.py @@ -568,6 +568,14 @@ def forward(self, x): return torch.cos(x) +class CumSum(torch.nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + return x.cumsum(dim=0) + + class Div(torch.nn.Module): def __init__(self): super().__init__() diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py index 9aba5a059e0..895326c743a 100644 --- a/backends/qualcomm/tests/test_qnn_delegate.py +++ b/backends/qualcomm/tests/test_qnn_delegate.py @@ -233,6 +233,11 @@ def test_qnn_backend_cos(self): sample_input = (torch.randn(2, 5, 1, 3),) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_cumsum(self): + module = CumSum() # noqa: F405 + sample_input = (torch.randn(4),) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_einsum_outer_product(self): module = EinsumOuterProduct() # noqa: F405 x = torch.randn(5) @@ -1297,6 +1302,12 @@ def test_qnn_backend_cos(self): module = self.get_qdq_module(module, sample_input) self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_cumsum(self): + module = CumSum() # noqa: F405 + sample_input = (torch.randn(4),) + module = self.get_qdq_module(module, sample_input) + self.lower_module_and_test_output(module, sample_input) + def test_qnn_backend_einsum_outer_product(self): module = EinsumOuterProduct() # noqa: F405 x = torch.randn(5) @@ -3537,7 +3548,6 @@ def test_conv_former(self): self.assertGreaterEqual(msg["top_1"], 60) self.assertGreaterEqual(msg["top_5"], 80) - @unittest.skip("bicubic resize is not supported") def test_dino_v2(self): if not self.required_envs([self.image_dataset]): self.skipTest("missing required envs") @@ -3573,6 +3583,46 @@ def test_dino_v2(self): self.assertGreaterEqual(msg["top_1"], 70) self.assertGreaterEqual(msg["top_5"], 85) + def test_efficientSAM(self): + if not self.required_envs( + [self.image_dataset, self.pretrained_weight, self.oss_repo] + ): + self.skipTest("missing required envs") + cmds = [ + "python", + f"{self.executorch_root}/examples/qualcomm/oss_scripts/efficientSAM.py", + "--dataset", + self.image_dataset, + "--artifact", + self.artifact_dir, + "--build_folder", + self.build_folder, + "--device", + self.device, + "--model", + self.model, + "--oss_repo", + self.oss_repo, + "--pretrained_weight", + self.pretrained_weight, + "--ip", + self.ip, + "--port", + str(self.port), + ] + if self.host: + cmds.extend(["--host", self.host]) + + p = subprocess.Popen(cmds, stdout=subprocess.DEVNULL) + with Listener((self.ip, self.port)) as listener: + conn = listener.accept() + p.communicate() + msg = json.loads(conn.recv()) + if "Error" in msg: + self.fail(msg["Error"]) + else: + self.assertGreaterEqual(msg["MIoU"], 0.55) + def test_esrgan(self): if not self.required_envs(): self.skipTest("missing required envs") diff --git a/backends/qualcomm/tests/utils.py b/backends/qualcomm/tests/utils.py index 42eec15891c..a3006eff8f6 100644 --- a/backends/qualcomm/tests/utils.py +++ b/backends/qualcomm/tests/utils.py @@ -438,12 +438,14 @@ def lower_module_and_test_output( skip_node_id_set: set = None, skip_node_op_set: set = None, dynamic_shapes: Dict = None, + passes_job: collections.OrderedDict = None, ): delegated_program = to_edge_transform_and_lower_to_qnn( module, sample_inputs, self.compiler_specs, dynamic_shapes=dynamic_shapes, + passes_job=passes_job, skip_node_id_set=skip_node_id_set, skip_node_op_set=skip_node_op_set, ) diff --git a/examples/qualcomm/oss_scripts/dino_v2.py b/examples/qualcomm/oss_scripts/dino_v2.py index 2eb26e6cece..18b5ade8b35 100644 --- a/examples/qualcomm/oss_scripts/dino_v2.py +++ b/examples/qualcomm/oss_scripts/dino_v2.py @@ -10,7 +10,12 @@ import numpy as np import torch +from executorch.backends.qualcomm._passes import ConvertUpsampleBicubicWithBilinear +from executorch.backends.qualcomm._passes.qnn_pass_manager import ( + get_capture_program_passes, +) from executorch.backends.qualcomm.quantizer.quantizer import QuantDtype +from executorch.backends.qualcomm.utils.constants import QCOM_PASS_ACTIVATE_KEY from executorch.examples.qualcomm.utils import ( build_executorch_binary, @@ -56,6 +61,8 @@ def main(args): pte_filename = "dino_v2" instance = get_instance() + passes_job = get_capture_program_passes() + passes_job[ConvertUpsampleBicubicWithBilinear][QCOM_PASS_ACTIVATE_KEY] = True build_executorch_binary( instance, sample_input, @@ -65,6 +72,7 @@ def main(args): skip_node_id_set=skip_node_id_set, skip_node_op_set=skip_node_op_set, quant_dtype=QuantDtype.use_8a8w, + passes_job=passes_job, shared_buffer=args.shared_buffer, ) diff --git a/examples/qualcomm/oss_scripts/efficientSAM/efficientSAM.py b/examples/qualcomm/oss_scripts/efficientSAM/efficientSAM.py new file mode 100644 index 00000000000..ea65917dcd9 --- /dev/null +++ b/examples/qualcomm/oss_scripts/efficientSAM/efficientSAM.py @@ -0,0 +1,357 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. + +import getpass +import json +import os +import zipfile +from multiprocessing.connection import Client +from typing import Callable, List + +import numpy as np +import torch +from executorch.backends.qualcomm._passes import ( + ConvertUpsampleBicubicWithBilinear, + ExpandBroadcastTensorShape, +) +from executorch.backends.qualcomm._passes.qnn_pass_manager import ( + get_capture_program_passes, +) +from executorch.backends.qualcomm.utils.constants import QCOM_PASS_ACTIVATE_KEY +from executorch.examples.qualcomm.oss_scripts.efficientSAM.source_transformation import ( + replace_maskdecoder_with_custom_op, + replace_pos_emb_with_custom_op, +) + +from executorch.examples.qualcomm.utils import ( + build_executorch_binary, + class_agnostic_mIoU, + make_output_dir, + parse_skip_delegation_node, + setup_common_args_and_variables, + SimpleADB, +) +from PIL import Image, ImageDraw +from scipy.ndimage import label +from torch.utils.data import DataLoader, Dataset +from torchvision import datasets, transforms + + +def load_dataset(dataset_path): + image_shape = (224, 224) + preprocess = transforms.Compose( + [ + transforms.Resize(image_shape), + transforms.ToTensor(), + ] + ) + imagenet_data = datasets.ImageFolder(dataset_path, transform=preprocess) + + return list(imagenet_data) + + +class EfficientSAMDataset(Dataset): + def __init__(self, dataset_path, data_size=1) -> None: + self.to_tensor = transforms.ToTensor() + dataset = load_dataset(dataset_path) + self.inputs = self.get_val_dataset(dataset, data_size) + self.data_size = data_size + + def get_val_dataset(self, dataset, data_size): + imgs, pt_prompts, pt_labels = [], [], [] + for i, data in enumerate(dataset): + if i >= data_size: + break + img = data[0] + h, w = img.shape[-2:] + + # Assuming the main object usually appears in the middle of the image, this default value is set for better demo visualization. + # Users can modify/add the point prompt here. + pt_prompt = torch.tensor([[w / 2, (h * 2 / 3)]], dtype=torch.float32)[ + None, ... + ] + # Users can increase the tensor size by adding more labels (0 for negative samples, 1 for positive samples) to label the corresponding points. + # The default label is [[1]], indicating that the point is a positive sample. + pt_label = torch.tensor([[1]], dtype=torch.float32) + + imgs.append(img) + pt_prompts.append(pt_prompt) + pt_labels.append(pt_label) + + imgs = torch.stack(imgs) + pt_prompts = torch.stack(pt_prompts) + pt_labels = torch.stack(pt_labels) + inputs = (imgs, pt_prompts, pt_labels) + return inputs + + def __getitem__(self, idx): + return self.inputs[0][idx], self.inputs[1][idx], self.inputs[2][idx] + + def __len__(self): + return self.data_size + + +def get_dataset(dataset_path, data_size=1): + + dataset = EfficientSAMDataset(dataset_path, data_size=data_size) + dataloader = DataLoader(dataset) + + # prepare input data + inputs, input_list = [], "" + for index, data in enumerate(dataloader): + if index >= data_size: + break + inputs.append(tuple(data)) + num_feature = len(data) + for idx, _ in enumerate(data): + input_name = f"input_{index}_{idx}.raw" + input_list += input_name + " " if idx < num_feature - 1 else input_name + + input_list = input_list + "\n" + + return inputs, input_list + + +def source_transform( + model, transforms: List[Callable[[torch.nn.Module], torch.nn.Module]] +): + for transform in transforms: + model = transform(model) + return model + + +def get_instance(args): + import sys + + sys.path.insert(0, args.oss_repo) + from efficient_sam.efficient_sam import build_efficient_sam + + ckpt = args.pretrained_weight + file_path, file_extension = os.path.splitext(ckpt) + file_dir, filename = os.path.split(file_path) + + if file_extension == ".zip": + with zipfile.ZipFile(ckpt, "r") as zip_ref: + zip_ref.extractall(file_dir) + ckpt = file_path + filename = os.path.splitext(filename)[0] + + model_arch = filename.split("_")[-1] + + if model_arch == "vitt": + encoder_patch_embed_dim, encoder_num_heads = (192, 3) + elif model_arch == "vits": + encoder_patch_embed_dim, encoder_num_heads = (384, 6) + else: + raise ValueError(f"Unsupported model architecture: {model_arch}") + + model = build_efficient_sam( + encoder_patch_embed_dim=encoder_patch_embed_dim, + encoder_num_heads=encoder_num_heads, + checkpoint=ckpt, + ).eval() + + return model + + +def generate_mask(predicted_logits, predicted_iou): + sorted_ids = torch.argsort(predicted_iou, dim=-1, descending=True) + predicted_iou = torch.take_along_dim(predicted_iou, sorted_ids, dim=2) + predicted_logits = torch.take_along_dim( + predicted_logits, sorted_ids[..., None, None], dim=2 + ) + + # The masks are already sorted by their predicted IOUs. + # We use the first mask. + mask = torch.ge(predicted_logits[0, 0, 0, :, :], 0).cpu().detach().numpy() + return mask + + +def save_mask(mask, input, save_path): + image, prompt, pt_label = input + original_image_tensor = image[0] + + # Convert tensor to numpy array if necessary + if not isinstance(original_image_tensor, np.ndarray): + original_image_tensor = original_image_tensor.detach().numpy() + + # Transpose if the image has 3 channels + if original_image_tensor.shape[0] == 3: + original_image_tensor = original_image_tensor.transpose(1, 2, 0) + + original_img = Image.fromarray( + (original_image_tensor * 255).astype(np.uint8) + ).convert("RGBA") + + # Create an empty RGBA image for the mask + mask_img = np.ones((mask.shape[0], mask.shape[1], 4)) + mask_img[:, :, 3] = 0 + + colors = [ + [1, 0, 0, 0.5], + [0, 1, 0, 0.5], + [0, 0, 1, 0.5], + [1, 1, 0, 0.5], + [1, 0, 1, 0.5], + [0, 1, 1, 0.5], + ] + + # Apply mask + labeled_mask, num_feature = label(mask) + for i in range(1, num_feature + 1): + mask_img[labeled_mask == i] = colors[(i - 1) % len(colors)] + + mask_img = Image.fromarray((mask_img * 255).astype(np.uint8), "RGBA") + + # Combine original image with mask + combined_img = Image.alpha_composite(original_img, mask_img) + + # Draw prompts point ("green" for positive samples, "red" for negative samples) + draw = ImageDraw.Draw(combined_img) + for pt, l in zip(prompt[0][0], pt_label[0][0]): + color = "green" if l else "red" + point_size = 3 + x1, y1 = max(0, int(pt[0]) - point_size), max(0, int(pt[1]) - point_size) + x2, y2 = min(combined_img.size[0], int(pt[0]) + point_size), min( + combined_img.size[1], int(pt[1]) + point_size + ) + draw.ellipse((x1, y1, x2, y2), fill=color, outline=color) + + combined_img.save(save_path) + + +def main(args): + skip_node_id_set, skip_node_op_set = parse_skip_delegation_node(args) + + os.makedirs(args.artifact, exist_ok=True) + + data_size = 1 + inputs, input_list = get_dataset(args.dataset, data_size) + assert args.pretrained_weight, "Checkpoint params can't be empty" + + # Get the EfficientSAM model. + model = get_instance(args) + model = source_transform( + model, + [ + replace_maskdecoder_with_custom_op, + replace_pos_emb_with_custom_op, + ], + ) + + pte_filename = "efficientSAM_qnn" + + # lower to QNN + passes_job = get_capture_program_passes() + passes_job[ConvertUpsampleBicubicWithBilinear][QCOM_PASS_ACTIVATE_KEY] = True + passes_job[ExpandBroadcastTensorShape][QCOM_PASS_ACTIVATE_KEY] = True + build_executorch_binary( + model, + inputs[0], + args.model, + f"{args.artifact}/{pte_filename}", + dataset=inputs, + skip_node_id_set=skip_node_id_set, + skip_node_op_set=skip_node_op_set, + passes_job=passes_job, + shared_buffer=args.shared_buffer, + ) + + if args.compile_only: + return + + workspace = f"/data/local/tmp/{getpass.getuser()}/executorch/{pte_filename}" + pte_path = f"{args.artifact}/{pte_filename}.pte" + + adb = SimpleADB( + qnn_sdk=os.getenv("QNN_SDK_ROOT"), + build_path=f"{args.build_folder}", + pte_path=pte_path, + workspace=workspace, + device_id=args.device, + host_id=args.host, + soc_model=args.model, + ) + adb.push(inputs=inputs, input_list=input_list) + adb.execute() + + # collect output data + output_data_folder = f"{args.artifact}/outputs" + make_output_dir(output_data_folder) + outputs = [] + + def post_process(): + for i, f in enumerate(sorted(os.listdir(output_data_folder))): + filename = os.path.join(output_data_folder, f) + output = np.fromfile(filename, dtype=np.float32) + output_shape = [1, 1, 3] if i % 2 else [1, 1, 3, 224, 224] + output = torch.from_numpy(output).reshape(output_shape) + outputs.append(output) + + adb.pull(output_path=args.artifact, callback=post_process) + + # MIoU analysis + miou = 0 + targets = [model(img, pt, pt_label) for img, pt, pt_label in inputs] + for i in range(data_size): + pred_mask = generate_mask(outputs[i * 2], outputs[i * 2 + 1]) + save_mask(pred_mask, inputs[i], f"{args.artifact}/output_{i}.png") + target_mask = generate_mask(targets[i][0], targets[i][1]) + miou += class_agnostic_mIoU([pred_mask], [target_mask]) + miou /= data_size + + if args.ip and args.port != -1: + with Client((args.ip, args.port)) as conn: + conn.send(json.dumps({"MIoU": miou})) + else: + print(f"MIoU->{miou}") + + +if __name__ == "__main__": + parser = setup_common_args_and_variables() + parser.add_argument( + "-a", + "--artifact", + help="path for storing generated artifacts and output by this example. Default ./EfficientSAM_qnn", + default="./EfficientSAM_qnn", + type=str, + ) + + parser.add_argument( + "--pretrained_weight", + help="Path to ESAM checkpoint, such as ./efficient_sam_vitt.pt or ./efficient_sam_vits.pt.zip", + type=str, + required=True, + ) + + parser.add_argument( + "-d", + "--dataset", + help=( + "path to the validation folder of ImageNet dataset. " + "e.g. --dataset imagenet-mini/val " + "for https://www.kaggle.com/datasets/ifigotin/imagenetmini-1000)" + ), + type=str, + required=True, + ) + + parser.add_argument( + "--oss_repo", + help="Path to clone https://github.com/yformer/EfficientSAM", + type=str, + required=True, + ) + + args = parser.parse_args() + try: + main(args) + except Exception as e: + if args.ip and args.port != -1: + with Client((args.ip, args.port)) as conn: + conn.send(json.dumps({"Error": str(e)})) + else: + raise Exception(e) diff --git a/examples/qualcomm/oss_scripts/efficientSAM/source_transformation/__init__.py b/examples/qualcomm/oss_scripts/efficientSAM/source_transformation/__init__.py new file mode 100644 index 00000000000..fd54a727136 --- /dev/null +++ b/examples/qualcomm/oss_scripts/efficientSAM/source_transformation/__init__.py @@ -0,0 +1,17 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from executorch.examples.qualcomm.oss_scripts.efficientSAM.source_transformation.mask_decoder import ( + replace_maskdecoder_with_custom_op, +) +from executorch.examples.qualcomm.oss_scripts.efficientSAM.source_transformation.pos_emb import ( + replace_pos_emb_with_custom_op, +) + + +__all__ = [ + replace_maskdecoder_with_custom_op, + replace_pos_emb_with_custom_op, +] diff --git a/examples/qualcomm/oss_scripts/efficientSAM/source_transformation/mask_decoder.py b/examples/qualcomm/oss_scripts/efficientSAM/source_transformation/mask_decoder.py new file mode 100644 index 00000000000..c70d51a48fe --- /dev/null +++ b/examples/qualcomm/oss_scripts/efficientSAM/source_transformation/mask_decoder.py @@ -0,0 +1,125 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from typing import List, Tuple + +import torch +import torch.nn as nn + + +class MaskDecoderCustom(nn.Module): + def forward( + self, + image_embeddings: torch.Tensor, + image_pe: torch.Tensor, + sparse_prompt_embeddings: torch.Tensor, + multimask_output: bool, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """ + Predict masks given image and prompt embeddings. + + Arguments: + image_embeddings: A tensor of shape [B, C, H, W] or [B*max_num_queries, C, H, W] + image_pe (torch.Tensor): positional encoding with the shape of image_embeddings (the batch dimension is broadcastable). + sparse_prompt_embeddings (torch.Tensor): the embeddings of the points and boxes + multimask_output (bool): Whether to return multiple masks or a single + mask. + + Returns: + torch.Tensor: batched predicted masks + torch.Tensor: batched predictions of mask quality + """ + + ( + batch_size, + max_num_queries, + sparse_embed_dim_1, + sparse_embed_dim_2, + ) = sparse_prompt_embeddings.shape + + ( + _, + image_embed_dim_c, + image_embed_dim_h, + image_embed_dim_w, + ) = image_embeddings.shape + + # QNN don't support dim greater than 4 + image_embeddings_expanded = image_embeddings.expand(max_num_queries, -1, -1, -1) + image_embeddings_tiled = image_embeddings_expanded.contiguous().view( + batch_size * max_num_queries, + image_embed_dim_c, + image_embed_dim_h, + image_embed_dim_w, + ) + sparse_prompt_embeddings = sparse_prompt_embeddings.reshape( + batch_size * max_num_queries, sparse_embed_dim_1, sparse_embed_dim_2 + ) + masks, iou_pred = self.predict_masks( + image_embeddings=image_embeddings_tiled, + image_pe=image_pe, + sparse_prompt_embeddings=sparse_prompt_embeddings, + ) + if multimask_output and self.num_multimask_outputs > 1: + return masks[:, 1:, :], iou_pred[:, 1:] + else: + return masks[:, :1, :], iou_pred[:, :1] + + def predict_masks( + self, + image_embeddings: torch.Tensor, + image_pe: torch.Tensor, + sparse_prompt_embeddings: torch.Tensor, + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Predicts masks. See 'forward' for more details.""" + # Concatenate output tokens + output_tokens = torch.cat( + [self.iou_token.weight, self.mask_tokens.weight], dim=0 + ) + output_tokens = output_tokens.unsqueeze(0).expand( + sparse_prompt_embeddings.size(0), -1, -1 + ) + tokens = torch.cat((output_tokens, sparse_prompt_embeddings), dim=1) + # Expand per-image data in batch direction to be per-mask + # QNN don't support dim greater than 4, + pos_src = image_pe.expand([tokens.shape[0]] + [*image_pe.shape[1:]]) + b, c, h, w = image_embeddings.shape + hs, src = self.transformer(image_embeddings, pos_src, tokens) + iou_token_out = hs[:, 0, :] + mask_tokens_out = hs[:, 1 : (1 + self.num_mask_tokens), :] + + # Upscale mask embeddings and predict masks using the mask tokens + upscaled_embedding = src.transpose(1, 2).view(b, c, h, w) + + for upscaling_layer in self.final_output_upscaling_layers: + upscaled_embedding = upscaling_layer(upscaled_embedding) + hyper_in_list: List[torch.Tensor] = [] + for i, output_hypernetworks_mlp in enumerate(self.output_hypernetworks_mlps): + hyper_in_list.append(output_hypernetworks_mlp(mask_tokens_out[:, i, :])) + hyper_in = torch.stack(hyper_in_list, dim=1) + b, c, h, w = upscaled_embedding.shape + masks = (hyper_in @ upscaled_embedding.view(b, c, h * w)).view(b, -1, h, w) + # Generate mask quality predictions + iou_pred = self.iou_prediction_head(iou_token_out) + return masks, iou_pred + + +def _replace_maskdecoder_with_custom_op(module: torch.nn.Module): + from efficient_sam.efficient_sam_decoder import MaskDecoder # B007 + + for _, child in module.named_children(): + if isinstance(child, MaskDecoder): + child.forward = MaskDecoderCustom.forward.__get__(child, MaskDecoder) + child.predict_masks = MaskDecoderCustom.predict_masks.__get__( + child, MaskDecoder + ) + else: + _replace_maskdecoder_with_custom_op(child) + + +def replace_maskdecoder_with_custom_op(module: torch.nn.Module) -> torch.nn.Module: + + _replace_maskdecoder_with_custom_op(module) + return module diff --git a/examples/qualcomm/oss_scripts/efficientSAM/source_transformation/pos_emb.py b/examples/qualcomm/oss_scripts/efficientSAM/source_transformation/pos_emb.py new file mode 100644 index 00000000000..7a3a91c7607 --- /dev/null +++ b/examples/qualcomm/oss_scripts/efficientSAM/source_transformation/pos_emb.py @@ -0,0 +1,64 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# All rights reserved +# +# This source code is licensed under the BSD-style license found in the +# LICENSE file in the root directory of this source tree. +from typing import Tuple + +import numpy as np + +import torch +import torch.nn as nn + + +class PositionEmbeddingRandomCustom(nn.Module): + """ + Positional encoding using random spatial frequencies. + """ + + def _pe_encoding(self, coords: torch.Tensor) -> torch.Tensor: + """Positionally encode points that are normalized to [0,1].""" + # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape + coords = 2 * coords - 1 + coords = coords.unsqueeze(0) + coords = torch.matmul( + coords, self.positional_encoding_gaussian_matrix.unsqueeze(0) + ) + coords = coords.squeeze(0) + coords = 2 * np.pi * coords + # outputs d_1 x ... x d_n x C shape + return torch.cat([torch.sin(coords), torch.cos(coords)], dim=-1) + + def forward_with_coords( + self, coords_input: torch.Tensor, image_size: Tuple[int, int] + ) -> torch.Tensor: + """Positionally encode points that are not normalized to [0,1].""" + coords = coords_input.clone() + coords_0 = coords[:, :, 0] / image_size[1] + coords_1 = coords[:, :, 1] / image_size[0] + coords = torch.stack((coords_0, coords_1), dim=-1) + + return self._pe_encoding(coords.to(torch.float)) # B x N x C + + +def _replace_pos_emb_with_custom_op(module: torch.nn.Module): + from efficient_sam.efficient_sam_decoder import PositionEmbeddingRandom # B007 + + for _, child in module.named_children(): + if isinstance(child, PositionEmbeddingRandom): + child._pe_encoding = PositionEmbeddingRandomCustom._pe_encoding.__get__( + child, PositionEmbeddingRandom + ) + child.forward_with_coords = ( + PositionEmbeddingRandomCustom.forward_with_coords.__get__( + child, PositionEmbeddingRandom + ) + ) + else: + _replace_pos_emb_with_custom_op(child) + + +def replace_pos_emb_with_custom_op(module: torch.nn.Module) -> torch.nn.Module: + + _replace_pos_emb_with_custom_op(module) + return module diff --git a/examples/qualcomm/oss_scripts/fastvit.py b/examples/qualcomm/oss_scripts/fastvit.py index 501ea522acd..b15e7d7267f 100644 --- a/examples/qualcomm/oss_scripts/fastvit.py +++ b/examples/qualcomm/oss_scripts/fastvit.py @@ -107,14 +107,16 @@ def main(args): weight=weight_qspec, bias=_derived_bias_quant_spec, ) + # rewrite default ptq config - q_config = quantizer.bit8_quant_config - quantizer.bit8_quant_config = QuantizationConfig( + q_config = quantizer.quant_config + quantizer.quant_config = QuantizationConfig( input_activation=act_qspec, output_activation=act_qspec, weight=q_config.weight, bias=q_config.bias, ) + # lower to QNN passes_job = get_capture_program_passes() passes_job[ExpandBroadcastTensorShape][QCOM_PASS_ACTIVATE_KEY] = True diff --git a/examples/qualcomm/utils.py b/examples/qualcomm/utils.py index b17bc8f98bd..3aecf405fff 100755 --- a/examples/qualcomm/utils.py +++ b/examples/qualcomm/utils.py @@ -428,6 +428,15 @@ def histogram(golden, predict): return (pa, mpa, miou, cls_iou) +def class_agnostic_mIoU(predictions, targets): + total_iou = 0 + for pred, tar in zip(predictions, targets): + inter = np.count_nonzero(pred & tar) + union = np.count_nonzero(pred | tar) + total_iou += inter / (union + 1e-10) + return total_iou / len(predictions) + + def get_imagenet_dataset( dataset_path, data_size, image_shape, crop_size=None, shuffle=True ):