Skip to content

Commit

Permalink
[Torch Fx] PTQ Quantizer config test (#2937)
Browse files Browse the repository at this point in the history
### Changes

New test file added in tests/torch/fx. Implemented
TemplateTestQuantizerConfig for the ptq quantizer for TorchFX backend.

### Related tickets

Closes issue #2874
  • Loading branch information
rk119 authored Sep 23, 2024
1 parent 393c2dd commit d9b3f38
Show file tree
Hide file tree
Showing 2 changed files with 71 additions and 0 deletions.
33 changes: 33 additions & 0 deletions tests/torch/fx/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,14 @@

from nncf.experimental.torch.fx.transformations import apply_quantization_transformations
from nncf.torch.dynamic_graph.patch_pytorch import disable_patching
from nncf.torch.graph.graph import PTNNCFGraph
from nncf.torch.graph.operator_metatypes import PTConstNoopMetatype
from nncf.torch.graph.operator_metatypes import PTModuleConv2dMetatype
from nncf.torch.graph.operator_metatypes import PTModuleDepthwiseConv2dSubtype
from nncf.torch.graph.operator_metatypes import PTSumMetatype
from tests.cross_fw.test_templates.models import NNCFGraphToTest
from tests.cross_fw.test_templates.models import NNCFGraphToTestDepthwiseConv
from tests.cross_fw.test_templates.models import NNCFGraphToTestSumAggregation


class TinyImagenetDatasetManager:
Expand Down Expand Up @@ -126,3 +134,28 @@ def get_torch_fx_model(model: torch.nn.Module) -> torch.fx.GraphModule:
fx_model = capture_pre_autograd_graph(model, args=(ex_input,))
apply_quantization_transformations(fx_model)
return fx_model


def get_single_conv_nncf_graph() -> NNCFGraphToTest:
return NNCFGraphToTest(
conv_metatype=PTModuleConv2dMetatype,
nncf_graph_cls=PTNNCFGraph,
const_metatype=PTConstNoopMetatype,
)


def get_depthwise_conv_nncf_graph() -> NNCFGraphToTestDepthwiseConv:
return NNCFGraphToTestDepthwiseConv(
depthwise_conv_metatype=PTModuleDepthwiseConv2dSubtype,
nncf_graph_cls=PTNNCFGraph,
const_metatype=PTConstNoopMetatype,
)


def get_sum_aggregation_nncf_graph() -> NNCFGraphToTestSumAggregation:
return NNCFGraphToTestSumAggregation(
conv_metatype=PTModuleConv2dMetatype,
sum_metatype=PTSumMetatype,
nncf_graph_cls=PTNNCFGraph,
const_metatype=PTConstNoopMetatype,
)
38 changes: 38 additions & 0 deletions tests/torch/fx/test_quantizer_config.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
# Copyright (c) 2024 Intel Corporation
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pytest

from nncf.quantization.algorithms.min_max.torch_fx_backend import FXMinMaxAlgoBackend
from tests.cross_fw.test_templates.models import NNCFGraphToTest
from tests.cross_fw.test_templates.models import NNCFGraphToTestDepthwiseConv
from tests.cross_fw.test_templates.models import NNCFGraphToTestSumAggregation
from tests.cross_fw.test_templates.test_quantizer_config import TemplateTestQuantizerConfig
from tests.torch.fx.helpers import get_depthwise_conv_nncf_graph
from tests.torch.fx.helpers import get_single_conv_nncf_graph
from tests.torch.fx.helpers import get_sum_aggregation_nncf_graph


class TestQuantizerConfig(TemplateTestQuantizerConfig):
def get_algo_backend(self):
return FXMinMaxAlgoBackend()

@pytest.fixture
def single_conv_nncf_graph(self) -> NNCFGraphToTest:
return get_single_conv_nncf_graph()

@pytest.fixture
def depthwise_conv_nncf_graph(self) -> NNCFGraphToTestDepthwiseConv:
return get_depthwise_conv_nncf_graph()

@pytest.fixture
def conv_sum_aggregation_nncf_graph(self) -> NNCFGraphToTestSumAggregation:
return get_sum_aggregation_nncf_graph()

0 comments on commit d9b3f38

Please sign in to comment.