From 9721cfbd0d5185737427b78f35e679e8074c5444 Mon Sep 17 00:00:00 2001
From: Daniil Lyakhov <daniil.lyakhov@intel.com>
Date: Wed, 27 Mar 2024 14:49:24 +0100
Subject: [PATCH] [Test][Examples] Make resnet18 qat test deterministic (#2594)

### Changes

Deterministic mode is enabled for resnet18 qat example

### Reason for changes

To make test_examples stable

### Related tickets

136387

### Tests
test_examples/314/ PASSED
test_examples/317/ PASSED
---
 tests/cross_fw/examples/example_scope.json |  9 +++++---
 tests/cross_fw/examples/run_example.py     | 24 ++++++++++++++++++++++
 tests/cross_fw/examples/test_examples.py   |  7 +++++++
 3 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/tests/cross_fw/examples/example_scope.json b/tests/cross_fw/examples/example_scope.json
index a890e2c8495..edebf9a839e 100644
--- a/tests/cross_fw/examples/example_scope.json
+++ b/tests/cross_fw/examples/example_scope.json
@@ -172,12 +172,15 @@
         "backend": "torch",
         "requirements": "examples/quantization_aware_training/torch/resnet18/requirements.txt",
         "cpu": "Intel(R) Core(TM) i9-10980XE CPU @ 3.00GHz",
+        "accuracy_tolerance_after_training": 1.0,
         "accuracy_tolerance": 0.2,
         "accuracy_metrics": {
             "fp32_top1": 55.52000045776367,
-            "int8_init_top1": 55.279998779296875,
-            "int8_top1": 56.7721,
-            "accuracy_drop": -1.3499984741210938
+            "int8_init_top1": 55.279998779296875
+        },
+        "accuracy_metrics_after_training":{
+            "int8_top1": 56.74446202531646,
+            "accuracy_drop": -1.2244606018066406
         },
         "performance_metrics": {
             "fp32_fps": 3646.13,
diff --git a/tests/cross_fw/examples/run_example.py b/tests/cross_fw/examples/run_example.py
index 2edbd061cb3..dc44a00e650 100644
--- a/tests/cross_fw/examples/run_example.py
+++ b/tests/cross_fw/examples/run_example.py
@@ -10,6 +10,7 @@
 # limitations under the License.
 
 import json
+import os
 import sys
 from argparse import ArgumentParser
 from typing import Dict, Tuple
@@ -172,6 +173,8 @@ def llm_tune_params() -> Dict[str, float]:
 def quantization_aware_training_torch_resnet18():
     from examples.quantization_aware_training.torch.resnet18.main import main as resnet18_main
 
+    # Set manual seed and determenistic cuda mode to make the test determenistic
+    set_torch_cuda_seed()
     results = resnet18_main()
 
     return {
@@ -188,6 +191,27 @@ def quantization_aware_training_torch_resnet18():
     }
 
 
+def set_torch_cuda_seed(seed: int = 42):
+    """
+    Sets torch, cuda and python random module to determenistic mode with
+    given seed.
+    :param seed: Seed to use for determenistic run.
+    """
+    import random
+
+    import numpy as np
+    import torch
+    from torch.backends import cudnn
+
+    np.random.seed(seed)
+    random.seed(seed)
+    torch.manual_seed(seed)
+    cudnn.deterministic = True
+    cudnn.benchmark = False
+    torch.use_deterministic_algorithms(True)
+    os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+
+
 def main(argv):
     parser = ArgumentParser()
     parser.add_argument("--name", help="Example name", required=True)
diff --git a/tests/cross_fw/examples/test_examples.py b/tests/cross_fw/examples/test_examples.py
index 8127315d94d..9dc38ea2d24 100644
--- a/tests/cross_fw/examples/test_examples.py
+++ b/tests/cross_fw/examples/test_examples.py
@@ -33,6 +33,7 @@
 MODEL_SIZE_RELATIVE_TOLERANCE = 0.05
 
 ACCURACY_METRICS = "accuracy_metrics"
+ACCURACY_METRICS_AFTER_TRAINING = "accuracy_metrics_after_training"
 MODEL_SIZE_METRICS = "model_size_metrics"
 PERFORMANCE_METRICS = "performance_metrics"
 
@@ -83,6 +84,12 @@ def test_examples(
             value, abs=example_params.get("accuracy_tolerance", ACCURACY_TOLERANCE)
         )
 
+    if ACCURACY_METRICS_AFTER_TRAINING in example_params:
+        for name, value in example_params[ACCURACY_METRICS_AFTER_TRAINING].items():
+            assert measured_metrics[name] == pytest.approx(
+                value, abs=example_params.get("accuracy_tolerance_after_training", ACCURACY_TOLERANCE)
+            )
+
     if MODEL_SIZE_METRICS in example_params:
         for name, value in example_params[MODEL_SIZE_METRICS].items():
             assert measured_metrics[name] == pytest.approx(value, rel=MODEL_SIZE_RELATIVE_TOLERANCE)