From 9d65011a0c5513f702bcca8b849efdd3e6d6f0e3 Mon Sep 17 00:00:00 2001 From: dlyakhov Date: Mon, 18 May 2026 14:44:09 +0200 Subject: [PATCH] [FX] Redundant sanity test is removed --- tests/torch/fx/helpers.py | 90 ---------------------- tests/torch/fx/test_models.py | 13 +++- tests/torch/fx/test_sanity.py | 139 ---------------------------------- 3 files changed, 12 insertions(+), 230 deletions(-) delete mode 100644 tests/torch/fx/test_sanity.py diff --git a/tests/torch/fx/helpers.py b/tests/torch/fx/helpers.py index 18363cff888..24d451b8b9a 100644 --- a/tests/torch/fx/helpers.py +++ b/tests/torch/fx/helpers.py @@ -9,14 +9,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -from pathlib import Path import torch.fx -import torch.nn.parallel -import torch.utils.data -import torchvision.datasets as datasets -import torchvision.transforms as transforms -from fastdownload import FastDownload from torch.fx.passes.graph_drawer import FxGraphDrawer from nncf.experimental.torch.fx.transformations import apply_quantization_transformations @@ -30,90 +24,6 @@ from tests.cross_fw.test_templates.models import NNCFGraphToTestSumAggregation -class TinyImagenetDatasetManager: - DATASET_URL = "http://cs231n.stanford.edu/tiny-imagenet-200.zip" - DATASET_PATH = "~/.cache/nncf/tests/datasets" - - def __init__(self, image_size: int, batch_size: int) -> None: - self.image_size = image_size - self.batch_size = batch_size - - @staticmethod - def download_dataset() -> Path: - downloader = FastDownload(base=TinyImagenetDatasetManager.DATASET_PATH, archive="downloaded", data="extracted") - return downloader.get(TinyImagenetDatasetManager.DATASET_URL) - - @staticmethod - def prepare_tiny_imagenet_200(dataset_dir: Path): - # Format validation set the same way as train set is formatted. - val_data_dir = dataset_dir / "val" - val_images_dir = val_data_dir / "images" - if not val_images_dir.exists(): - return - - val_annotations_file = val_data_dir / "val_annotations.txt" - with open(val_annotations_file) as f: - val_annotation_data = map(lambda line: line.split("\t")[:2], f.readlines()) - for image_filename, image_label in val_annotation_data: - from_image_filepath = val_images_dir / image_filename - to_image_dir = val_data_dir / image_label - if not to_image_dir.exists(): - to_image_dir.mkdir() - to_image_filepath = to_image_dir / image_filename - from_image_filepath.rename(to_image_filepath) - val_annotations_file.unlink() - val_images_dir.rmdir() - - def create_data_loaders(self): - dataset_path = TinyImagenetDatasetManager.download_dataset() - - TinyImagenetDatasetManager.prepare_tiny_imagenet_200(dataset_path) - print(f"Successfully downloaded and prepared dataset at: {dataset_path}") - - train_dir = dataset_path / "train" - val_dir = dataset_path / "val" - - normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) - - train_dataset = datasets.ImageFolder( - train_dir, - transforms.Compose( - [ - transforms.Resize(self.image_size), - transforms.RandomHorizontalFlip(), - transforms.ToTensor(), - normalize, - ] - ), - ) - val_dataset = datasets.ImageFolder( - val_dir, - transforms.Compose( - [ - transforms.Resize(self.image_size), - transforms.ToTensor(), - normalize, - ] - ), - ) - - train_loader = torch.utils.data.DataLoader( - train_dataset, batch_size=self.batch_size, shuffle=True, num_workers=0, pin_memory=True, sampler=None - ) - - val_loader = torch.utils.data.DataLoader( - val_dataset, batch_size=self.batch_size, shuffle=False, num_workers=0, pin_memory=True - ) - - # Creating separate dataloader with batch size = 1 - # as dataloaders with batches > 1 are not supported yet. - calibration_dataset = torch.utils.data.DataLoader( - val_dataset, batch_size=1, shuffle=False, num_workers=0, pin_memory=True - ) - - return train_loader, val_loader, calibration_dataset - - def visualize_fx_model(model: torch.fx.GraphModule, output_svg_path: str): g = FxGraphDrawer(model, output_svg_path) g.get_dot_graph().write_svg(output_svg_path) diff --git a/tests/torch/fx/test_models.py b/tests/torch/fx/test_models.py index 6b4b2bc2142..eefbb47b6d5 100644 --- a/tests/torch/fx/test_models.py +++ b/tests/torch/fx/test_models.py @@ -38,7 +38,6 @@ from tests.cross_fw.test_templates.helpers import YOLO26AttentionBlock from tests.torch import test_models from tests.torch.fx.helpers import get_torch_fx_model -from tests.torch.fx.test_sanity import count_q_dq from tests.torch.test_models.synthetic import ConcatSameTensorModel from tests.torch.test_models.synthetic import ConvReluBranchModel from tests.torch.test_models.synthetic import EmbeddingSumModel @@ -198,6 +197,18 @@ def test_model(test_case: ModelCase, regen_ref_data: bool): ) +def count_q_dq(model: torch.fx.GraphModule): + q, dq = 0, 0 + for node in model.graph.nodes: + if node.op == "call_function" and hasattr(node.target, "overloadpacket"): + node_type = str(node.target.overloadpacket).split(".")[1] + if node_type in ["quantize_per_tensor", "quantize_per_channel"]: + q += 1 + elif node_type in ["dequantize_per_tensor", "dequantize_per_channel"]: + dq += 1 + return q, dq + + @pytest.mark.parametrize("enable_dynamic_shapes", [True, False]) @pytest.mark.parametrize("compress_weights", [True, False]) @pytest.mark.parametrize( diff --git a/tests/torch/fx/test_sanity.py b/tests/torch/fx/test_sanity.py deleted file mode 100644 index 0029179069a..00000000000 --- a/tests/torch/fx/test_sanity.py +++ /dev/null @@ -1,139 +0,0 @@ -# Copyright (c) 2026 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. - - -from dataclasses import dataclass - -import numpy as np -import openvino.torch # noqa -import pytest -import torch -import torch.nn as nn -import torch.nn.parallel -import torch.utils.data -import torchvision.models as models - -import nncf -from nncf.common.logging.track_progress import track -from tests.torch.fx.helpers import TinyImagenetDatasetManager - -IMAGE_SIZE = 64 -BATCH_SIZE = 128 - - -@pytest.fixture(name="tiny_imagenet_dataset", scope="module") -def tiny_imagenet_dataset_fixture(): - return TinyImagenetDatasetManager(IMAGE_SIZE, BATCH_SIZE).create_data_loaders() - - -@dataclass -class SanitySampleCase: - model_id: str - checkpoint_url: str - top1_int8_ref: float - ref_num_q: int - ref_num_dq: int - - -MODELS = ( - SanitySampleCase( - "resnet18", - "https://storage.openvinotoolkit.org/repositories/nncf/openvino_notebook_ckpts/302_resnet18_fp32_v1.pth", - 55.30, - 30, - 37, - ), -) - - -def get_model(model_id: str, checkpoint_url: str, device: torch.device) -> torch.nn.Module: - num_classes = 200 # 200 is for Tiny ImageNet, default is 1000 for ImageNet - model = getattr(models, model_id)(weights=None) - # Update the last FC layer for Tiny ImageNet number of classes. - model.fc = nn.Linear(in_features=512, out_features=num_classes, bias=True) - model.to(device) - checkpoint = torch.hub.load_state_dict_from_url(checkpoint_url, map_location=torch.device("cpu"), progress=False) - model.load_state_dict(checkpoint["state_dict"]) - return model - - -def validate(val_loader: torch.utils.data.DataLoader, model: torch.nn.Module, device: torch.device) -> float: - top1_sum = 0.0 - with torch.no_grad(): - for images, target in track(val_loader, total=len(val_loader), description="Validation:"): - images = images.to(device) - target = target.to(device) - - # Compute output. - output = model(images) - - # Measure accuracy and record loss. - [acc1] = accuracy(output, target, topk=(1,)) - top1_sum += acc1.item() - - num_samples = len(val_loader) - top1_avg = top1_sum / num_samples - return top1_avg - - -def accuracy(output: torch.Tensor, target: torch.tensor, topk: tuple[int, ...] = (1,)): - with torch.no_grad(): - maxk = max(topk) - batch_size = target.size(0) - - _, pred = output.topk(maxk, 1, True, True) - pred = pred.t() - correct = pred.eq(target.view(1, -1).expand_as(pred)) - - res = [] - for k in topk: - correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True) - res.append(correct_k.mul_(100.0 / batch_size)) - return res - - -def count_q_dq(model: torch.fx.GraphModule): - q, dq = 0, 0 - for node in model.graph.nodes: - if node.op == "call_function" and hasattr(node.target, "overloadpacket"): - node_type = str(node.target.overloadpacket).split(".")[1] - if node_type in ["quantize_per_tensor", "quantize_per_channel"]: - q += 1 - elif node_type in ["dequantize_per_tensor", "dequantize_per_channel"]: - dq += 1 - return q, dq - - -@pytest.mark.parametrize("test_case", MODELS) -def test_sanity(test_case: SanitySampleCase, tiny_imagenet_dataset): - torch.manual_seed(42) - device = torch.device("cpu") - model = get_model(test_case.model_id, test_case.checkpoint_url, device) - _, val_dataloader, calibration_dataset = tiny_imagenet_dataset - - def transform_fn(data_item): - return data_item[0].to(device) - - calibration_dataset = nncf.Dataset(calibration_dataset, transform_fn) - - with torch.no_grad(): - ex_input = next(iter(calibration_dataset.get_inference_data())) - model.eval() - exported_model = torch.export.export(model, args=(ex_input,), strict=True).module(check_guards=False) - quantized_model = nncf.quantize(exported_model, calibration_dataset) - quantized_model = torch.compile(quantized_model, dynamic=True, backend="openvino") - - top1_int8 = validate(val_dataloader, quantized_model, device) - assert np.isclose(top1_int8, test_case.top1_int8_ref, atol=0.1) - - num_q, num_dq = count_q_dq(quantized_model) - assert num_q == test_case.ref_num_q - assert num_dq == test_case.ref_num_dq