Merge pull request #727 from mv1388/add_new_unittests

Add new unittests

aitoolbox/experiment/local_save/local_model_save.py

@@ -125,7 +125,6 @@ def check_model_dict_contents(model):
         Returns:
             None
         """
-        # TODO: maybe add some check about the actual values/content of the dict as well
         for required_element in ['model_state_dict', 'optimizer_state_dict', 'epoch', 'hyperparams']:
             if required_element not in model:
                 raise ValueError(f'Required element of the model dict {required_element} is missing. Given model'

tests/test_experiment/test_local_experiment_saver.py

@@ -49,6 +49,39 @@ def test_init(self):
                 local_model_result_folder_path=THIS_DIR
             )
 
+    def test_save_experiment_experiment_timestamp_not_provided(self):
+        model = Net()
+        project_dir_name = 'projectPyTorchLocalModelSaver'
+        exp_dir_name = 'experimentSubDirPT'
+        current_time = datetime.datetime.fromtimestamp(time.time()).strftime('%Y-%m-%d_%H-%M-%S')
+
+        project_path = os.path.join(THIS_DIR, project_dir_name)
+        exp_path = os.path.join(project_path, f'{exp_dir_name}_{current_time}')
+        model_path = os.path.join(exp_path, 'model')
+        model_file_path = os.path.join(model_path, f'model_{exp_dir_name}_{current_time}.pth')
+        results_path = os.path.join(exp_path, 'results')
+        results_file_path = os.path.join(results_path, f'results_hyperParams_hist_{exp_dir_name}_{current_time}.p')
+
+        saver = BaseFullExperimentLocalSaver(
+            PyTorchLocalModelSaver(local_model_result_folder_path=THIS_DIR),
+            project_name=project_dir_name, experiment_name=exp_dir_name,
+            local_model_result_folder_path=THIS_DIR
+        )
+
+        result_pkg = DummyFullResultPackage({'metric1': 33434, 'acc1': 223.43, 'loss': 4455.6},
+                                            {'epoch': 20, 'lr': 0.334})
+        training_history = TrainingHistory().wrap_pre_prepared_history({})
+
+        model_checkpoint = {'model_state_dict': model.state_dict(), 'optimizer_state_dict': None,
+                            'epoch': 10, 'hyperparams': {}}
+
+        saved_paths = saver.save_experiment(model_checkpoint, result_pkg, training_history)
+
+        self.assertEqual(saved_paths, [model_file_path, results_file_path])
+
+        if os.path.exists(project_path):
+            shutil.rmtree(project_path)
+
 
 class TestFullPyTorchExperimentLocalSaver(unittest.TestCase):
     def test_init(self):

tests/test_experiment/test_local_save/test_local_model_save.py

@@ -239,6 +239,43 @@ def save_model_local(self, checkpoint_model, epoch):
         if os.path.exists(project_path):
             shutil.rmtree(project_path)
 
+    def test_fail_check_model_dict_contents(self):
+        project_dir_name = 'projectPyTorchLocalModelSaver'
+        exp_dir_name = 'experimentSubDirPT'
+        current_time = datetime.datetime.fromtimestamp(time.time()).strftime('%Y-%m-%d_%H-%M-%S')
+
+        model = Net()
+        saver = PyTorchLocalModelSaver(local_model_result_folder_path=THIS_DIR)
+
+        model_checkpoint = {'model_state_dict': model.state_dict(), 'optimizer_state_dict': None,
+                            'epoch': 10, 'hyperparams': {}}
+
+        for key_excluded in model_checkpoint.keys():
+            model_check = {k: v for k, v in model_checkpoint.items() if k != key_excluded}
+
+            with self.assertRaises(ValueError):
+                saver.check_model_dict_contents(model_check)
+
+        for key_excluded in model_checkpoint.keys():
+            model_check = {k: v for k, v in model_checkpoint.items() if k != key_excluded}
+            model_check['additional'] = 'aaaa'
+
+            with self.assertRaises(ValueError):
+                saver.check_model_dict_contents(model_check)
+
+        for key_excluded in model_checkpoint.keys():
+            model_check = {k: v for k, v in model_checkpoint.items() if k != key_excluded}
+
+            with self.assertRaises(ValueError):
+                saver.save_model(model_check, project_dir_name, exp_dir_name, current_time, 1)
+
+        for key_excluded in model_checkpoint.keys():
+            model_check = {k: v for k, v in model_checkpoint.items() if k != key_excluded}
+            model_check['additional'] = 'aaaa'
+
+            with self.assertRaises(ValueError):
+                saver.save_model(model_check, project_dir_name, exp_dir_name, current_time, 1)
+
 
 class TestKerasLocalModelSaver(unittest.TestCase):
     def test_init(self):

tests/test_experiment/test_result_package/test_abstract_result_packages.py

@@ -15,7 +15,29 @@ def test_basic(self):
 
         self.assertEqual(str(result_pkg), 'dummy: 111.0\nextended_dummy: 1323123.44')
         self.assertEqual(len(result_pkg), 2)
-
+
+    def test_get_results_error_trigger(self):
+        result_pkg = DummyResultPackage()
+        self.assertIsNone(result_pkg.get_results())
+
+        result_pkg = DummyResultPackage(strict_content_check=True)
+        with self.assertRaises(ValueError):
+            result_pkg.get_results()
+
+    def test_qa_check_hyperparameters_dict_error_trigger(self):
+        result_pkg = DummyResultPackage(strict_content_check=True)
+
+        with self.assertRaises(ValueError):
+            result_pkg.qa_check_hyperparameters_dict()
+
+    def test_warn_about_result_data_problem(self):
+        result_pkg = DummyResultPackage(strict_content_check=False)
+        result_pkg.warn_about_result_data_problem('aaaaa')
+
+        result_pkg = DummyResultPackage(strict_content_check=True)
+        with self.assertRaises(ValueError):
+            result_pkg.warn_about_result_data_problem('aaaaa')
+
     def test_get_additional_results_dump_paths(self):
         paths_1 = [['filename', 'file/path/filename']]
         result_pkg_1 = DummyResultPackageExtendV2(paths_1)

tests/test_torchtrain/test_tl_components/test_pred_collate_fns.py

@@ -1,5 +1,6 @@
 import unittest
 import numpy as np
+import torch
 
 from aitoolbox.torchtrain.train_loop.components.pred_collate_fns import *
 
@@ -17,15 +18,8 @@ def test_append_predictions(self):
         self.assertEqual([preds_ep_1, preds_ep_2], preds)
 
     def test_append_concat_predictions(self):
-        preds = []
-        preds_ep_1 = np.random.rand(100, 1)
-        preds = append_predictions(preds_ep_1, preds)
-        self.assertEqual([preds_ep_1], preds)
-        preds_ep_2 = np.random.rand(45, 1)
-        preds = append_predictions(preds_ep_2, preds)
-        self.assertEqual([preds_ep_1, preds_ep_2], preds)
-
         preds_list = []
+
         preds_list_ep_1 = np.random.rand(100).tolist()
         preds_list = append_concat_predictions(preds_list_ep_1, preds_list)
         self.assertEqual(preds_list_ep_1, preds_list)
@@ -34,6 +28,48 @@ def test_append_concat_predictions(self):
         preds_list = append_concat_predictions(preds_list_ep_2, preds_list)
         self.assertEqual(preds_list_ep_1 + preds_list_ep_2, preds_list)
 
+    def test_append_concat_predictions_non_list_np_array(self):
+        preds_list = []
+
+        preds_list_ep_1 = np.random.rand(100)
+        preds_list = append_concat_predictions(preds_list_ep_1, preds_list)
+        self.assertEqual([preds_list_ep_1], preds_list)
+
+        preds_list_ep_2 = np.random.rand(100)
+        preds_list = append_concat_predictions(preds_list_ep_2, preds_list)
+        self.assertEqual([preds_list_ep_1, preds_list_ep_2], preds_list)
+
+        preds_list = []
+
+        preds_list_ep_1 = np.random.rand(100, 5)
+        preds_list = append_concat_predictions(preds_list_ep_1, preds_list)
+        self.assertEqual([preds_list_ep_1], preds_list)
+
+        preds_list_ep_2 = np.random.rand(100, 5)
+        preds_list = append_concat_predictions(preds_list_ep_2, preds_list)
+        self.assertEqual([preds_list_ep_1, preds_list_ep_2], preds_list)
+
+    def test_append_concat_predictions_non_list_tensor(self):
+        preds_list = []
+
+        preds_list_ep_1 = torch.rand(100)
+        preds_list = append_concat_predictions(preds_list_ep_1, preds_list)
+        self.assertEqual([preds_list_ep_1], preds_list)
+
+        preds_list_ep_2 = torch.rand(100)
+        preds_list = append_concat_predictions(preds_list_ep_2, preds_list)
+        self.assertEqual([preds_list_ep_1, preds_list_ep_2], preds_list)
+
+        preds_list = []
+
+        preds_list_ep_1 = torch.rand(100, 5)
+        preds_list = append_concat_predictions(preds_list_ep_1, preds_list)
+        self.assertEqual([preds_list_ep_1], preds_list)
+
+        preds_list_ep_2 = torch.rand(100, 5)
+        preds_list = append_concat_predictions(preds_list_ep_2, preds_list)
+        self.assertEqual([preds_list_ep_1, preds_list_ep_2], preds_list)
+
 
 class TestAllPredTransformFns(unittest.TestCase):
     def test_torch_cat_transf(self):
@@ -47,3 +83,7 @@ def test_torch_cat_transf(self):
                               torch.Tensor([5, 6, 7]), torch.Tensor([100, 200])]).numpy().tolist(),
             np.array([1., 2., 3., 4., 5., 6., 7., 100., 200.]).tolist()
         )
+
+    def test_keep_list_transf(self):
+        data = list(range(100))
+        self.assertEqual(keep_list_transf(data), data)

tests/test_torchtrain/test_train_loop/test_train_loop.py

@@ -645,6 +645,74 @@ def test_get_schedulers(self):
         tl_filtered_schedulers = train_loop.get_schedulers()
         self.assertEqual(tl_filtered_schedulers, schedulers_list)
 
+    def test_get_num_training_steps(self):
+        model = NetUnifiedBatchFeed()
+        dummy_optimizer = DummyOptimizer()
+        dummy_loss = DummyLoss()
+        dummy_train_loader = list(range(100))
+
+        train_loop = TrainLoop(model, dummy_train_loader, None, None, dummy_optimizer, dummy_loss)
+        train_loop.fit(num_epochs=3)
+        self.assertEqual(train_loop.get_num_training_steps(), 300)
+
+        train_loop = TrainLoop(model, dummy_train_loader, None, None, dummy_optimizer, dummy_loss)
+        train_loop.fit(num_epochs=10)
+        self.assertEqual(train_loop.get_num_training_steps(), 1000)
+
+        train_loop = TrainLoop(model, dummy_train_loader, None, None, dummy_optimizer, dummy_loss)
+        train_loop.fit(num_epochs=3, grad_accumulation=5)
+        self.assertEqual(train_loop.get_num_training_steps(), 60)
+
+        train_loop = TrainLoop(model, dummy_train_loader, None, None, dummy_optimizer, dummy_loss)
+        train_loop.fit(num_epochs=3, grad_accumulation=10)
+        self.assertEqual(train_loop.get_num_training_steps(), 30)
+
+        train_loop = TrainLoop(model, dummy_train_loader, None, None, dummy_optimizer, dummy_loss)
+        train_loop.fit(num_epochs=3, grad_accumulation=20)
+        self.assertEqual(train_loop.get_num_training_steps(), 15)
+
+        train_loop = TrainLoop(model, dummy_train_loader, None, None, dummy_optimizer, dummy_loss)
+        train_loop.fit(num_epochs=10, grad_accumulation=5)
+        self.assertEqual(train_loop.get_num_training_steps(), 200)
+
+    def test_get_num_training_steps_num_iterations(self):
+        model = NetUnifiedBatchFeed()
+        dummy_optimizer = DummyOptimizer()
+        dummy_loss = DummyLoss()
+        dummy_train_loader = list(range(4))
+
+        train_loop = TrainLoop(model, dummy_train_loader, None, None, dummy_optimizer, dummy_loss)
+        train_loop.fit(num_iterations=120)
+        self.assertEqual(train_loop.get_num_training_steps(), 120)
+
+        train_loop = TrainLoop(model, dummy_train_loader, None, None, dummy_optimizer, dummy_loss)
+        train_loop.fit(num_iterations=100)
+        self.assertEqual(train_loop.get_num_training_steps(), 100)
+
+        train_loop = TrainLoop(model, dummy_train_loader, None, None, dummy_optimizer, dummy_loss)
+        train_loop.fit(num_iterations=100, grad_accumulation=5)
+        self.assertEqual(train_loop.get_num_training_steps(), 100 / 5)
+
+        train_loop = TrainLoop(model, dummy_train_loader, None, None, dummy_optimizer, dummy_loss)
+        train_loop.fit(num_iterations=100, grad_accumulation=20)
+        self.assertEqual(train_loop.get_num_training_steps(), 100 / 20)
+
+        train_loop = TrainLoop(model, dummy_train_loader, None, None, dummy_optimizer, dummy_loss)
+        train_loop.fit(num_iterations=100, grad_accumulation=21)
+        self.assertEqual(train_loop.get_num_training_steps(), 100 // 21)
+
+    def test_call(self):
+        model = NetUnifiedBatchFeed()
+        dummy_optimizer = DummyOptimizer()
+        dummy_loss = DummyLoss()
+        dummy_train_loader = list(range(100))
+
+        train_loop = TrainLoop(model, dummy_train_loader, None, None, dummy_optimizer, dummy_loss)
+        self.assertEqual(train_loop.epoch, 0)
+        train_loop(num_epochs=3)
+        self.assertEqual(train_loop.epoch, 2)
+        self.assertEqual(train_loop.total_iteration_idx, 100 * 3 - 1)
+
 
 class TestMultiLossOptiTrainLoop(unittest.TestCase):
     def test_multi_loss_in_train_loop(self):

tests/test_utils/test_file_system.py

@@ -57,9 +57,20 @@ def test_create_folder_hierarchy_two_prong(self):
             shutil.rmtree(os.path.join(THIS_DIR, 'folder_1'))
 
     def test_zip_folder(self):
+        self.zip_folder()
+
+    def test_zip_folder_with_zip_extension(self):
+        self.zip_folder(with_zip_extension=True)
+
+    def zip_folder(self, with_zip_extension=False):
         dummy_dir_path, dummy_files_content = self.prepare_dummy_folder()
 
-        file_system.zip_folder(dummy_dir_path, dummy_dir_path)
+        if with_zip_extension:
+            dummy_zip_path = f'{dummy_dir_path}.zip'
+        else:
+            dummy_zip_path = dummy_dir_path
+
+        file_system.zip_folder(dummy_dir_path, dummy_zip_path)
         self.assertTrue(os.path.exists(dummy_dir_path + '.zip'))
 
         with zipfile.ZipFile(dummy_dir_path + '.zip', 'r') as zip_ref:

tests/utils.py

@@ -353,8 +353,8 @@ class DummyNonAbstractResultPackage:
 
 
 class DummyResultPackage(AbstractResultPackage):
-    def __init__(self):
-        AbstractResultPackage.__init__(self, 'DummyPackage', False)
+    def __init__(self, strict_content_check=False):
+        AbstractResultPackage.__init__(self, 'DummyPackage', strict_content_check)
         self.experiment_path = None
 
     def prepare_results_dict(self):