From 3fb62df12278f6f7572eaddf5d8c95ec45801c9b Mon Sep 17 00:00:00 2001
From: dlyakhov <daniil.lyakhov@intel.com>
Date: Wed, 26 Jun 2024 17:02:50 +0200
Subject: [PATCH] torch-fx tests are added to pre-commit

---
 .github/workflows/precommit.yml | 44 +++++++++++++++++++++++++++++++++
 Makefile                        | 11 +++++++++
 tests/torch_fx/requirements.txt |  6 +++++
 tests/torch_fx/test_sanity.py   | 16 +++++++-----
 4 files changed, 71 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/precommit.yml b/.github/workflows/precommit.yml
index 6f7ec2906d..b94256b642 100644
--- a/.github/workflows/precommit.yml
+++ b/.github/workflows/precommit.yml
@@ -95,6 +95,50 @@ jobs:
           token: ${{ secrets.CODECOV_TOKEN }}
           name: coverage_openvino
           flags: OPENVINO
+  torchFX:
+    timeout-minutes: 40
+    defaults:
+      run:
+        shell: bash
+    runs-on: ubuntu-20.04-8-cores
+    env:
+      DEBIAN_FRONTEND: noninteractive
+    steps:
+      - name: Install dependencies
+        run : |
+          sudo apt-get update
+          sudo apt-get --assume-yes install gcc g++ build-essential ninja-build libgl1-mesa-dev libglib2.0-0
+      - uses: actions/checkout@a5ac7e51b41094c92402da3b24376905380afc29 # v4.1.6
+        with:
+          lfs: true
+      - uses: actions/setup-python@0a5c61591373683505ea898e09a3ea4f39ef2b9c # v5.0.0
+        with:
+          python-version: 3.8.18
+          cache: pip
+      - name: Runner info
+        continue-on-error: true
+        run: |
+          cat /etc/*release
+          cat /proc/cpuinfo
+      - name: Install NNCF and test requirements
+        run: make install-torch-fx-test
+      - name: Run PyTorch precommit test scope
+        run: |
+          make test-torch-fx
+        env:
+          NNCF_COVERAGE: 1
+          NUM_WORKERS: 4
+      - name: Upload coverage report as artifact
+        uses: actions/upload-artifact@65462800fd760344b1a7b4382951275a0abb4808 # v4.3.3
+        with:
+          name: coverage_pt_cpu
+          path: ./coverage.xml
+      - name: Upload coverage report to codecov
+        uses: codecov/codecov-action@125fc84a9a348dbcf27191600683ec096ec9021c # v4.4.1
+        with:
+          token: ${{ secrets.CODECOV_TOKEN }}
+          name: coverage_pt_cpu
+          flags: TORCH
   pytorch-cpu:
     timeout-minutes: 40
     defaults:
diff --git a/Makefile b/Makefile
index da65aa89e8..7842dd68ff 100644
--- a/Makefile
+++ b/Makefile
@@ -177,6 +177,17 @@ test-examples-torch:
 test-models-hub-torch:
 	pytest tests/torch/models_hub_test --junitxml ${JUNITXML_PATH}
 
+###############################################################################
+# TorchFX backend
+install-torch-fx-test:
+	pip install -U pip
+	pip install -e .
+	pip install -r tests/torch_fx/requirements.txt
+
+test-torch-fx:
+	pytest ${COVERAGE_ARGS} tests/torch_fx/  -ra \
+		--junitxml ${JUNITXML_PATH}
+
 ###############################################################################
 # Common part
 install-common-test:
diff --git a/tests/torch_fx/requirements.txt b/tests/torch_fx/requirements.txt
index 99ee43ce75..08617061e6 100644
--- a/tests/torch_fx/requirements.txt
+++ b/tests/torch_fx/requirements.txt
@@ -1 +1,7 @@
+-c ../../constraints.txt
+pytest
+pytest-cov
+openvino
+torch
+torchvision
 fastdownload==0.0.7
\ No newline at end of file
diff --git a/tests/torch_fx/test_sanity.py b/tests/torch_fx/test_sanity.py
index 197c2f9547..3f1ce72cb1 100644
--- a/tests/torch_fx/test_sanity.py
+++ b/tests/torch_fx/test_sanity.py
@@ -34,6 +34,11 @@
 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
@@ -47,7 +52,7 @@ class SanitySampleCase:
     SanitySampleCase(
         "resnet18",
         "https://storage.openvinotoolkit.org/repositories/nncf/openvino_notebook_ckpts/302_resnet18_fp32_v1.pth",
-        55.23,
+        55.2,
         51,
         58,
     ),
@@ -113,13 +118,12 @@ def count_q_dq(model: torch.fx.GraphModule):
 
 
 @pytest.mark.parametrize("test_case", MODELS)
-def test_sanity(test_case: SanitySampleCase):
+def test_sanity(test_case: SanitySampleCase, tiny_imagenet_dataset):
     with disable_patching():
+        torch.manual_seed(42)
         device = torch.device("cpu")
         model = get_model(test_case.model_id, test_case.checkpoint_url, device)
-        _, val_dataloader, calibration_dataset = TinyImagenetDatasetManager(
-            IMAGE_SIZE, BATCH_SIZE
-        ).create_data_loaders()
+        _, val_dataloader, calibration_dataset = tiny_imagenet_dataset
 
         def transform_fn(data_item):
             return data_item[0].to(device)
@@ -134,7 +138,7 @@ def transform_fn(data_item):
             quantized_model = torch.compile(quantized_model, backend="openvino")
 
         top1_int8 = validate(val_dataloader, quantized_model, device)
-        assert np.isclose(top1_int8, test_case.top1_int8_ref, atol=1e-2)
+        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