Remove hardcoded test directories (#3146)

* Remove hardcoded directories from (most if not all) tests, where they
remained.

---------

Signed-off-by: Michael Hofmann <quic_michof@quicinc.com>

NightlyTests/onnx/test_adaround.py

@@ -39,6 +39,7 @@
 import json
 import numpy as np
 import pytest
+import tempfile
 import torch
 from onnx import load_model
 from onnxruntime.quantization.onnx_quantizer import ONNXModel
@@ -52,8 +53,6 @@
 
 from aimet_onnx.adaround.adaround_weight import Adaround, AdaroundParameters
 
-WORKING_DIR = '/tmp/quantsim'
-
 image_size = 32
 batch_size = 64
 num_workers = 4
@@ -164,4 +163,4 @@ def build_session(model):
         sess_options=sess_options,
         providers=['CPUExecutionProvider'],
     )
-    return session
+    return session

NightlyTests/onnx/test_quantsim.py

@@ -38,6 +38,7 @@
 
 import numpy as np
 import pytest
+import tempfile
 import torch
 from onnx import load_model
 from torchvision import models
@@ -48,8 +49,6 @@
 from aimet_common.quantsim_config.utils import get_path_for_per_channel_config
 from torch_utils import get_cifar10_data_loaders, train_cifar10
 
-WORKING_DIR = '/tmp/quantsim'
-
 image_size = 32
 batch_size = 64
 num_workers = 4
@@ -79,42 +78,41 @@ class TestQuantizeAcceptance:
     @pytest.mark.parametrize("config_file", [None, get_path_for_per_channel_config()])
     @pytest.mark.cuda
     def test_quantized_accuracy(self, config_file):
-        if not os.path.exists(WORKING_DIR):
-            os.makedirs(WORKING_DIR)
-        np.random.seed(0)
-        torch.manual_seed(0)
-        model = models.resnet18(pretrained=False, num_classes=10)
-        if torch.cuda.is_available():
-            device = torch.device('cuda:0')
-            model.to(device)
-
-        train_cifar10(model, 2)
-        train_loader, val_loader = get_cifar10_data_loaders(drop_last=False)
-
-        torch.onnx.export(model, torch.rand(batch_size, 3, 32, 32).cuda(), os.path.join(WORKING_DIR, 'resnet18.onnx'),
-                          training=torch.onnx.TrainingMode.PRESERVE,
-                          input_names=['input'], output_names=['output'],
-                          dynamic_axes={
-                              'input': {0: 'batch_size'},
-                              'output': {0: 'batch_size'},
-                          }
-                          )
-
-        onnx_model = load_model(os.path.join(WORKING_DIR, 'resnet18.onnx'))
-        dummy_input = make_dummy_input(onnx_model)
-        sim = QuantizationSimModel(onnx_model, dummy_input, quant_scheme=QuantScheme.post_training_tf, default_param_bw=8,
-                                   default_activation_bw=8, use_cuda=True, config_file=config_file)
-
-        def onnx_callback(session, iters):
-            for i, batch in enumerate(train_loader):
-                x = batch[0].detach().cpu().numpy()
-                in_tensor = {'input': x}
-                session.run(None, in_tensor)
-                if i >= iters:
-                    break
-
-        sim.compute_encodings(onnx_callback, 10)
-
-        onnx_qs_acc = model_eval_onnx(sim.session, val_loader)
-
-        assert onnx_qs_acc > 0.5
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            np.random.seed(0)
+            torch.manual_seed(0)
+            model = models.resnet18(pretrained=False, num_classes=10)
+            if torch.cuda.is_available():
+                device = torch.device('cuda:0')
+                model.to(device)
+
+            train_cifar10(model, 2)
+            train_loader, val_loader = get_cifar10_data_loaders(drop_last=False)
+
+            torch.onnx.export(model, torch.rand(batch_size, 3, 32, 32).cuda(), os.path.join(tmp_dir, 'resnet18.onnx'),
+                            training=torch.onnx.TrainingMode.PRESERVE,
+                            input_names=['input'], output_names=['output'],
+                            dynamic_axes={
+                                'input': {0: 'batch_size'},
+                                'output': {0: 'batch_size'},
+                            }
+                            )
+
+            onnx_model = load_model(os.path.join(tmp_dir, 'resnet18.onnx'))
+            dummy_input = make_dummy_input(onnx_model)
+            sim = QuantizationSimModel(onnx_model, dummy_input, quant_scheme=QuantScheme.post_training_tf, default_param_bw=8,
+                                    default_activation_bw=8, use_cuda=True, config_file=config_file)
+
+            def onnx_callback(session, iters):
+                for i, batch in enumerate(train_loader):
+                    x = batch[0].detach().cpu().numpy()
+                    in_tensor = {'input': x}
+                    session.run(None, in_tensor)
+                    if i >= iters:
+                        break
+
+            sim.compute_encodings(onnx_callback, 10)
+
+            onnx_qs_acc = model_eval_onnx(sim.session, val_loader)
+
+            assert onnx_qs_acc > 0.5

NightlyTests/onnx/test_rnn_quantsim.py

@@ -38,6 +38,7 @@
 
 import numpy as np
 import pytest
+import tempfile
 import torch
 from onnx import load_model
 from torchaudio import models
@@ -48,8 +49,6 @@
 from aimet_common.quantsim_config.utils import get_path_for_per_channel_config
 from torch_utils import get_librispeech_data_loaders, train_librispeech
 
-WORKING_DIR = '/tmp/quantsim'
-
 batch_size = 64
 n_feature = 128
 n_class = 29
@@ -87,44 +86,43 @@ class TestQuantizeAcceptance:
     @pytest.mark.parametrize("config_file", [None, get_path_for_per_channel_config()])
     @pytest.mark.cuda
     def test_quantized_accuracy(self, config_file):
-        if not os.path.exists(WORKING_DIR):
-            os.makedirs(WORKING_DIR)
-        np.random.seed(0)
-        torch.manual_seed(0)
-        model = models.DeepSpeech(n_feature=n_feature, n_class=n_class)
-        if torch.cuda.is_available():
-            device = torch.device('cuda:0')
-            model.to(device)
-
-        train_librispeech(model, 1, max_batches=30)
-
-        train_loader, val_loader = get_librispeech_data_loaders(batch_size=batch_size, drop_last=False)
-
-        torch.onnx.export(model, torch.rand(1, 1, 1, 128).cuda(), os.path.join(WORKING_DIR, 'deepspeech.onnx'),
-                          training=torch.onnx.TrainingMode.PRESERVE,
-                          input_names=['input'], output_names=['output'],
-                          dynamic_axes={
-                              'input': {0: 'batch_size', 2: 'time'},
-                              'output': {0: 'batch_size', 1: 'time'},
-                          }
-                          )
-
-        onnx_model = load_model(os.path.join(WORKING_DIR, 'deepspeech.onnx'))
-        dummy_input = make_dummy_input(onnx_model)
-        sim = QuantizationSimModel(onnx_model, dummy_input, quant_scheme=QuantScheme.post_training_tf, default_param_bw=8,
-                                   default_activation_bw=8, use_cuda=True, config_file=config_file)
-
-        def onnx_callback(session, iters):
-            for i, batch in enumerate(train_loader):
-                x = batch[0].detach().cpu().numpy()
-                in_tensor = {'input': x}
-                session.run(None, in_tensor)
-                print(i, '/', iters)
-                if i+1 >= iters:
-                    break
-
-        sim.compute_encodings(onnx_callback, 1)
-
-        onnx_qs_test_loss = model_eval_onnx(sim.session, val_loader, max_batches=1)
-
-        assert onnx_qs_test_loss < 0.1
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            np.random.seed(0)
+            torch.manual_seed(0)
+            model = models.DeepSpeech(n_feature=n_feature, n_class=n_class)
+            if torch.cuda.is_available():
+                device = torch.device('cuda:0')
+                model.to(device)
+
+            train_librispeech(model, 1, max_batches=30)
+
+            train_loader, val_loader = get_librispeech_data_loaders(batch_size=batch_size, drop_last=False)
+
+            torch.onnx.export(model, torch.rand(1, 1, 1, 128).cuda(), os.path.join(tmp_dir, 'deepspeech.onnx'),
+                            training=torch.onnx.TrainingMode.PRESERVE,
+                            input_names=['input'], output_names=['output'],
+                            dynamic_axes={
+                                'input': {0: 'batch_size', 2: 'time'},
+                                'output': {0: 'batch_size', 1: 'time'},
+                            }
+                            )
+
+            onnx_model = load_model(os.path.join(tmp_dir, 'deepspeech.onnx'))
+            dummy_input = make_dummy_input(onnx_model)
+            sim = QuantizationSimModel(onnx_model, dummy_input, quant_scheme=QuantScheme.post_training_tf, default_param_bw=8,
+                                    default_activation_bw=8, use_cuda=True, config_file=config_file)
+
+            def onnx_callback(session, iters):
+                for i, batch in enumerate(train_loader):
+                    x = batch[0].detach().cpu().numpy()
+                    in_tensor = {'input': x}
+                    session.run(None, in_tensor)
+                    print(i, '/', iters)
+                    if i+1 >= iters:
+                        break
+
+            sim.compute_encodings(onnx_callback, 1)
+
+            onnx_qs_test_loss = model_eval_onnx(sim.session, val_loader, max_batches=1)
+
+            assert onnx_qs_test_loss < 0.1

NightlyTests/tensorflow/dependencies.py

@@ -46,7 +46,6 @@
 import wget
 import zipfile
 import tarfile
-import shutil
 from transformers import BertTokenizer, DistilBertTokenizer
 from aimet_common.utils import AimetLogger
 

NightlyTests/tensorflow/eager/test_bn_reestimation_keras.py

@@ -36,6 +36,8 @@
 # =============================================================================
 """ Keras bn_reestimation Nightly Tests """
 import json
+import tempfile
+from Pathlib import Path
 import tensorflow as tf
 import numpy as np
 from aimet_common.defs import QuantScheme
@@ -114,11 +116,14 @@ def _qsim_setup_for_fold_scale(model, dummy_inputs):
             {}
     }
 
-    with open("/tmp/default_config_per_channel.json", "w") as f:
-        json.dump(default_config_per_channel, f)
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        config_filename = Path(tmp_dir, "default_config_per_channel.json")
+        with open(config_filename, "w") as f:
+            json.dump(default_config_per_channel, f)
 
-    qsim = QuantizationSimModel(model, quant_scheme=QuantScheme.training_range_learning_with_tf_init,
-                                config_file="/tmp/default_config_per_channel.json")
+        qsim = QuantizationSimModel(model,
+                                    quant_scheme=QuantScheme.training_range_learning_with_tf_init,
+                                    config_file=config_filename)
 
 
     qsim.compute_encodings(lambda m, _: m.predict(dummy_inputs), None)

NightlyTests/tensorflow/non_eager/test_bn_reestimation.py

@@ -36,6 +36,8 @@
 # =============================================================================
 
 import pytest
+import tempfile
+from pathlib import Path
 import os
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'
 import logging
@@ -246,39 +248,40 @@ def test_model_rewriter_ptq_reestimation_fold(self, bn_re_estimation_dataset, bn
                 {}
         }
 
-        config_file_path = "/tmp/default_config_per_channel.json"
-        with open(config_file_path, "w") as f:
-            json.dump(default_config_per_channel, f)
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            config_file_path = Path(tmp_dir, "default_config_per_channel.json")
+            with open(config_file_path, "w") as f:
+                json.dump(default_config_per_channel, f)
 
-        sim = QuantizationSimModel(sess, start_op_names, end_op_names, use_cuda=True,
-                                   quant_scheme=QuantScheme.training_range_learning_with_tf_init,
-                                   config_file=config_file_path)
-        def dummy_forward_pass(sess, args):
-            model_input = sess.graph.get_tensor_by_name("input_1:0")
-            model_output = sess.graph.get_tensor_by_name('predictions/Softmax:0')
-            dummy_val = np.random.randn(1, *model_input.shape[1:])
-            sess.run(model_output, feed_dict={model_input: dummy_val})
-        sim.compute_encodings(dummy_forward_pass, None)
+            sim = QuantizationSimModel(sess, start_op_names, end_op_names, use_cuda=True,
+                                    quant_scheme=QuantScheme.training_range_learning_with_tf_init,
+                                    config_file=config_file_path)
+            def dummy_forward_pass(sess, args):
+                model_input = sess.graph.get_tensor_by_name("input_1:0")
+                model_output = sess.graph.get_tensor_by_name('predictions/Softmax:0')
+                dummy_val = np.random.randn(1, *model_input.shape[1:])
+                sess.run(model_output, feed_dict={model_input: dummy_val})
+            sim.compute_encodings(dummy_forward_pass, None)
 
-        # check bn_re_estimation
-        self._reestimate_and_compare_results(sim, bn_re_estimation_dataset, bn_num_batches, start_op_names, end_op_names)
+            # check bn_re_estimation
+            self._reestimate_and_compare_results(sim, bn_re_estimation_dataset, bn_num_batches, start_op_names, end_op_names)
 
-        # check bn_fold
-        model_input = sim.session.graph.get_tensor_by_name("input_1:0")
-        model_output = sim.session.graph.get_tensor_by_name('predictions/Softmax:0')
-        dummy_val = np.random.randn(128, *model_input.shape[1:])
+            # check bn_fold
+            model_input = sim.session.graph.get_tensor_by_name("input_1:0")
+            model_output = sim.session.graph.get_tensor_by_name('predictions/Softmax:0')
+            dummy_val = np.random.randn(128, *model_input.shape[1:])
 
-        output_baseline = sim.session.run(model_output, feed_dict={model_input: dummy_val})
+            output_baseline = sim.session.run(model_output, feed_dict={model_input: dummy_val})
 
-        fold_all_batch_norms_to_scale(sim, start_op_names, end_op_names)
+            fold_all_batch_norms_to_scale(sim, start_op_names, end_op_names)
 
-        model_input_after_fold = sim.session.graph.get_tensor_by_name("input_1:0")
-        model_output_after_fold = sim.session.graph.get_tensor_by_name('predictions/Softmax:0')
+            model_input_after_fold = sim.session.graph.get_tensor_by_name("input_1:0")
+            model_output_after_fold = sim.session.graph.get_tensor_by_name('predictions/Softmax:0')
 
-        output_fold_after_fold = sim.session.run(model_output_after_fold, feed_dict={model_input_after_fold: dummy_val})
+            output_fold_after_fold = sim.session.run(model_output_after_fold, feed_dict={model_input_after_fold: dummy_val})
 
-        assert np.allclose(output_baseline, output_fold_after_fold, atol=1e-2)
-        sim.session.close()
+            assert np.allclose(output_baseline, output_fold_after_fold, atol=1e-2)
+            sim.session.close()
 
     def test_remove_bn_update_ops_with_training_ops(self):
         """ verify that the BNs UPDATE_OPS are removed correctly after training ops are added (QAT) """

NightlyTests/tensorflow/non_eager/test_compress_example_tf.py

@@ -36,7 +36,6 @@
 # =============================================================================
 """ Acceptance tests for various compression techniques """
 
-import pytest
 import math
 import os
 os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'