Skip to content

expose all stages in init #302

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
wants to merge 3 commits into from
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
13 changes: 6 additions & 7 deletions backends/xnnpack/test/models/mobilenet_v2.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,16 +11,15 @@
from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
XnnpackQuantizedPartitioner,
)
from executorch.backends.xnnpack.test.tester import Partition, Tester
from executorch.backends.xnnpack.test.tester.tester import Export
from executorch.backends.xnnpack.test.tester import Export, Partition, Tester
from executorch.exir import CaptureConfig
from torchvision.models.mobilenetv2 import MobileNet_V2_Weights


class TestXNNPACKMobileNetV2(unittest.TestCase):
class TestMobileNetV2(unittest.TestCase):
mv2 = models.mobilenetv2.mobilenet_v2(weights=MobileNet_V2_Weights)
mv2 = mv2.eval()
model_inputs = (torch.ones(1, 3, 224, 244),)
model_inputs = (torch.ones(1, 3, 224, 224),)

all_operators = {
"executorch_exir_dialects_edge__ops_aten__native_batch_norm_legit_no_training_default",
Expand All @@ -32,7 +31,7 @@ class TestXNNPACKMobileNetV2(unittest.TestCase):
"executorch_exir_dialects_edge__ops_aten_convolution_default",
}

def test_mv2_fp32(self):
def test_fp32_mv2(self):

(
Tester(self.mv2, self.model_inputs)
Expand All @@ -48,15 +47,15 @@ def test_mv2_fp32(self):
.compare_outputs()
)

def test_mv2_qs8_pt2e(self):
def test_qs8_mv2(self):
# Quantization fuses away batchnorm, so it is no longer in the graph
ops_after_quantization = self.all_operators - {
"executorch_exir_dialects_edge__ops_aten__native_batch_norm_legit_no_training_default",
}

(
Tester(self.mv2, self.model_inputs)
.quantize2()
.quantize()
.export(Export(CaptureConfig(enable_aot=True)))
.to_edge()
.check(list(ops_after_quantization))
Expand Down
13 changes: 6 additions & 7 deletions backends/xnnpack/test/models/mobilenet_v3.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,15 +11,14 @@
from executorch.backends.xnnpack.partition.xnnpack_partitioner import (
XnnpackQuantizedPartitioner,
)
from executorch.backends.xnnpack.test.tester import Partition, Tester
from executorch.backends.xnnpack.test.tester.tester import Export
from executorch.backends.xnnpack.test.tester import Export, Partition, Tester
from executorch.exir import CaptureConfig


class TestXNNPACKMobileNetV3(unittest.TestCase):
class TestMobileNetV3(unittest.TestCase):
mv3 = models.mobilenetv3.mobilenet_v3_small(pretrained=True)
mv3 = mv3.eval()
model_inputs = (torch.ones(1, 3, 224, 244),)
model_inputs = (torch.ones(1, 3, 224, 224),)

all_operators = {
"executorch_exir_dialects_edge__ops_aten__native_batch_norm_legit_no_training_default",
Expand All @@ -35,7 +34,7 @@ class TestXNNPACKMobileNetV3(unittest.TestCase):
"executorch_exir_dialects_edge__ops_aten_mean_dim",
}

def test_mv3_fp32(self):
def test_fp32_mv3(self):
(
Tester(self.mv3, self.model_inputs)
.export(Export(CaptureConfig(enable_aot=True)))
Expand All @@ -50,7 +49,7 @@ def test_mv3_fp32(self):
.compare_outputs()
)

def test_mv3_qs8_pt2e(self):
def test_qs8_mv3(self):
ops_after_quantization = self.all_operators - {
"executorch_exir_dialects_edge__ops_aten__native_batch_norm_legit_no_training_default",
}
Expand All @@ -66,7 +65,7 @@ def test_mv3_qs8_pt2e(self):

(
Tester(self.mv3, self.model_inputs)
.quantize2()
.quantize()
.export(Export(CaptureConfig(enable_aot=True)))
.to_edge()
.check(list(ops_after_quantization))
Expand Down
4 changes: 2 additions & 2 deletions backends/xnnpack/test/ops/add.py
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,7 @@ def test_qs8_add(self):
inputs = (torch.ones(1), torch.ones(1))
(
Tester(self.Add(), inputs)
.quantize2()
.quantize()
.export()
.check_count({"torch.ops.aten.add.Tensor": 4})
.check(["torch.ops.quantized_decomposed"])
Expand Down Expand Up @@ -91,7 +91,7 @@ def test_qs8_add_relu(self):
inputs = (torch.randn(1, 1, 4, 4), torch.randn(1, 1, 4, 4))
(
Tester(self.AddRelu(), inputs)
.quantize2()
.quantize()
.export()
.check_count({"torch.ops.aten.add.Tensor": 1})
.check_count({"torch.ops.aten.relu.default": 1})
Expand Down
20 changes: 18 additions & 2 deletions backends/xnnpack/test/tester/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,22 @@
# LICENSE file in the root directory of this source tree.

# TODO: Be more delibrate on module structure
from executorch.backends.xnnpack.test.tester.tester import Partition, Tester
from executorch.backends.xnnpack.test.tester.tester import (
Export,
Partition,
Quantize,
Serialize,
Tester,
ToEdge,
ToExecutorch,
)

__all__ = [Tester, Partition]
__all__ = [
Tester,
Partition,
Quantize,
Export,
ToEdge,
ToExecutorch,
Serialize,
]
49 changes: 1 addition & 48 deletions backends/xnnpack/test/tester/tester.py
Original file line number Diff line number Diff line change
Expand Up @@ -30,22 +30,9 @@
from executorch.exir.backend.partitioner import Partitioner
from executorch.exir.passes.spec_prop_pass import SpecPropPass

from executorch.extension.pybindings.portable_lib import (
from executorch.extension.pybindings.portable_lib import ( # @manual
_load_for_executorch_from_buffer,
)
from torch.ao.quantization.backend_config import BackendConfig
from torch.ao.quantization.backend_config.executorch import (
get_executorch_backend_config,
)
from torch.ao.quantization.qconfig_mapping import (
_get_symmetric_qnnpack_qconfig_mapping,
QConfigMapping,
)

from torch.ao.quantization.quantize_fx import (
_convert_to_reference_decomposed_fx,
prepare_fx,
)
from torch.ao.quantization.quantize_pt2e import convert_pt2e, prepare_pt2e
from torch.ao.quantization.quantizer.quantizer import Quantizer
from torch.ao.quantization.quantizer.xnnpack_quantizer import (
Expand Down Expand Up @@ -103,36 +90,6 @@ def register_stage(stage: Stage):

@register_stage
class Quantize(Stage):
def __init__(
self,
qconfig_mapping: Optional[QConfigMapping] = None,
backend_config: Optional[BackendConfig] = None,
):
self.qconfig_mapping = (
qconfig_mapping or _get_symmetric_qnnpack_qconfig_mapping()
)
self.backend_config = backend_config or get_executorch_backend_config()
self.converted = None

def run(self, artifact: torch.nn.Module, inputs: Tuple[torch.Tensor]) -> None:
prepared = prepare_fx(
artifact, self.qconfig_mapping, inputs, backend_config=self.backend_config
)
self.converted = _convert_to_reference_decomposed_fx(
prepared, backend_config=self.backend_config
)

@property
def artifact(self) -> torch.fx.GraphModule:
return self.converted

@property
def graph_module(self) -> str:
return self.converted


@register_stage
class Quantize2(Stage):
def __init__(
self,
quantizer: Optional[Quantizer] = None,
Expand Down Expand Up @@ -278,7 +235,6 @@ def __init__(
self.inputs = inputs
self.stages: Dict[str, Stage] = OrderedDict.fromkeys(list(_stages_.keys()))
self.pipeline = {
self._stage_name(Quantize2): [self._stage_name(Export)],
self._stage_name(Quantize): [self._stage_name(Export)],
self._stage_name(Export): [
self._stage_name(ToEdge),
Expand Down Expand Up @@ -339,9 +295,6 @@ def quantize(self, quantize_stage: Optional[Quantize] = None):
def export(self, export_stage: Optional[Export] = None):
return self._run_stage(export_stage or Export(), self.inputs)

def quantize2(self, quantize_stage: Optional[Quantize2] = None):
return self._run_stage(quantize_stage or Quantize2(), self.inputs)

def to_edge(self, to_edge_stage: Optional[ToEdge] = None):
return self._run_stage(to_edge_stage or ToEdge())

Expand Down