[BLACK PR] apply black on tests

tests/deci_core_integration_test_suite_runner.py

@@ -5,7 +5,6 @@
 
 
 class CoreIntegrationTestSuiteRunner:
-
     def __init__(self):
         self.test_loader = unittest.TestLoader()
         self.integration_tests_suite = unittest.TestSuite()
@@ -22,5 +21,5 @@ def _add_modules_to_integration_tests_suite(self):
         self.integration_tests_suite.addTest(self.test_loader.loadTestsFromModule(LRTest))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()

tests/end_to_end_tests/__init__.py

@@ -4,4 +4,4 @@
 
 from tests.end_to_end_tests.cifar_trainer_test import TestCifarTrainer
 
-__all__ = ['TestTrainer', 'TestCifarTrainer']
+__all__ = ["TestTrainer", "TestCifarTrainer"]

tests/end_to_end_tests/trainer_test.py

@@ -16,66 +16,68 @@ class TestTrainer(unittest.TestCase):
     def setUp(cls):
         super_gradients.init_trainer()
         # NAMES FOR THE EXPERIMENTS TO LATER DELETE
-        cls.folder_names = ['test_train', 'test_save_load', 'test_load_w', 'test_load_w2',
-                            'test_load_w3', 'test_checkpoint_content', 'analyze']
-        cls.training_params = {"max_epochs": 1,
-                               "silent_mode": True,
-                               "lr_decay_factor": 0.1,
-                               "initial_lr": 0.1,
-                               "lr_updates": [4],
-                               "lr_mode": "step",
-                               "loss": "cross_entropy", "train_metrics_list": [Accuracy(), Top5()],
-                               "valid_metrics_list": [Accuracy(), Top5()],
-                               "metric_to_watch": "Accuracy",
-                               "greater_metric_to_watch_is_better": True}
+        cls.folder_names = ["test_train", "test_save_load", "test_load_w", "test_load_w2", "test_load_w3", "test_checkpoint_content", "analyze"]
+        cls.training_params = {
+            "max_epochs": 1,
+            "silent_mode": True,
+            "lr_decay_factor": 0.1,
+            "initial_lr": 0.1,
+            "lr_updates": [4],
+            "lr_mode": "step",
+            "loss": "cross_entropy",
+            "train_metrics_list": [Accuracy(), Top5()],
+            "valid_metrics_list": [Accuracy(), Top5()],
+            "metric_to_watch": "Accuracy",
+            "greater_metric_to_watch_is_better": True,
+        }
 
     @classmethod
     def tearDownClass(cls) -> None:
         # ERASE ALL THE FOLDERS THAT WERE CREATED DURING THIS TEST
         for folder in cls.folder_names:
-            if os.path.isdir(os.path.join('checkpoints', folder)):
-                shutil.rmtree(os.path.join('checkpoints', folder))
+            if os.path.isdir(os.path.join("checkpoints", folder)):
+                shutil.rmtree(os.path.join("checkpoints", folder))
 
     @staticmethod
-    def get_classification_trainer(name=''):
+    def get_classification_trainer(name=""):
         trainer = Trainer(name)
         model = models.get("resnet18", num_classes=5)
         return trainer, model
 
     def test_train(self):
         trainer, model = self.get_classification_trainer(self.folder_names[0])
-        trainer.train(model=model, training_params=self.training_params, train_loader=classification_test_dataloader(),
-                      valid_loader=classification_test_dataloader())
+        trainer.train(
+            model=model, training_params=self.training_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader()
+        )
 
     def test_save_load(self):
         trainer, model = self.get_classification_trainer(self.folder_names[1])
-        trainer.train(model=model, training_params=self.training_params, train_loader=classification_test_dataloader(),
-                      valid_loader=classification_test_dataloader())
+        trainer.train(
+            model=model, training_params=self.training_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader()
+        )
         resume_training_params = self.training_params.copy()
         resume_training_params["resume"] = True
         resume_training_params["max_epochs"] = 2
         trainer, model = self.get_classification_trainer(self.folder_names[1])
-        trainer.train(model=model, training_params=resume_training_params,
-                      train_loader=classification_test_dataloader(),
-                      valid_loader=classification_test_dataloader())
+        trainer.train(
+            model=model, training_params=resume_training_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader()
+        )
 
     def test_checkpoint_content(self):
         """VERIFY THAT ALL CHECKPOINTS ARE SAVED AND CONTAIN ALL THE EXPECTED KEYS"""
         trainer, model = self.get_classification_trainer(self.folder_names[5])
         params = self.training_params.copy()
         params["save_ckpt_epoch_list"] = [1]
-        trainer.train(model=model, training_params=params, train_loader=classification_test_dataloader(),
-                      valid_loader=classification_test_dataloader())
-        ckpt_filename = ['ckpt_best.pth', 'ckpt_latest.pth', 'ckpt_epoch_1.pth']
+        trainer.train(model=model, training_params=params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader())
+        ckpt_filename = ["ckpt_best.pth", "ckpt_latest.pth", "ckpt_epoch_1.pth"]
         ckpt_paths = [os.path.join(trainer.checkpoints_dir_path, suf) for suf in ckpt_filename]
         for ckpt_path in ckpt_paths:
             ckpt = torch.load(ckpt_path)
-            self.assertListEqual(['net', 'acc', 'epoch', 'optimizer_state_dict', 'scaler_state_dict'],
-                                 list(ckpt.keys()))
+            self.assertListEqual(["net", "acc", "epoch", "optimizer_state_dict", "scaler_state_dict"], list(ckpt.keys()))
         trainer._save_checkpoint()
-        weights_only = torch.load(os.path.join(trainer.checkpoints_dir_path, 'ckpt_latest_weights_only.pth'))
-        self.assertListEqual(['net'], list(weights_only.keys()))
+        weights_only = torch.load(os.path.join(trainer.checkpoints_dir_path, "ckpt_latest_weights_only.pth"))
+        self.assertListEqual(["net"], list(weights_only.keys()))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()

tests/integration_tests/conversion_callback_test.py

@@ -5,8 +5,7 @@
 from super_gradients.training import models
 
 from super_gradients import Trainer
-from super_gradients.training.dataloaders.dataloaders import segmentation_test_dataloader, \
-    classification_test_dataloader
+from super_gradients.training.dataloaders.dataloaders import segmentation_test_dataloader, classification_test_dataloader
 from super_gradients.training.utils.callbacks import ModelConversionCheckCallback
 from super_gradients.training.metrics import Accuracy, Top5, IoU
 from super_gradients.training.losses.stdc_loss import STDCLoss
@@ -63,18 +62,17 @@ def test_classification_architectures(self):
                 "criterion_params": {},
                 "train_metrics_list": [Accuracy(), Top5()],
                 "valid_metrics_list": [Accuracy(), Top5()],
-
                 "metric_to_watch": "Accuracy",
                 "greater_metric_to_watch_is_better": True,
                 "phase_callbacks": phase_callbacks,
             }
 
-            trainer = Trainer(f"{architecture}_example",
-                              ckpt_root_dir=checkpoint_dir)
+            trainer = Trainer(f"{architecture}_example", ckpt_root_dir=checkpoint_dir)
             model = models.get(architecture=architecture, arch_params={"use_aux_heads": True, "aux_head": True})
             try:
-                trainer.train(model=model, training_params=train_params, train_loader=classification_test_dataloader(),
-                              valid_loader=classification_test_dataloader())
+                trainer.train(
+                    model=model, training_params=train_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader()
+                )
             except Exception as e:
                 self.fail(f"Model training didn't succeed due to {e}")
             else:
@@ -84,26 +82,22 @@ def test_segmentation_architectures(self):
         def get_architecture_custom_config(architecture_name: str):
             if re.search(r"ddrnet", architecture_name):
                 return {
-
                     "loss": DDRNetLoss(num_pixels_exclude_ignored=False),
                 }
             elif re.search(r"stdc", architecture_name):
                 return {
-
                     "loss": STDCLoss(num_classes=5),
                 }
             elif re.search(r"regseg", architecture_name):
                 return {
-
                     "loss": "cross_entropy",
                 }
             else:
                 raise Exception("You tried to run a conversion test on an unknown architecture")
 
         for architecture in SEMANTIC_SEGMENTATION:
             model_meta_data = generate_model_metadata(architecture=architecture, task=Task.SEMANTIC_SEGMENTATION)
-            trainer = Trainer(f"{architecture}_example",
-                              ckpt_root_dir=checkpoint_dir)
+            trainer = Trainer(f"{architecture}_example", ckpt_root_dir=checkpoint_dir)
             model = models.get(model_name=architecture, arch_params={"use_aux_heads": True, "aux_head": True})
 
             phase_callbacks = [
@@ -128,8 +122,12 @@ def get_architecture_custom_config(architecture_name: str):
             train_params.update(custom_config)
 
             try:
-                trainer.train(model=model, training_params=train_params, train_loader=segmentation_test_dataloader(image_size=512),
-                              valid_loader=segmentation_test_dataloader(image_size=512))
+                trainer.train(
+                    model=model,
+                    training_params=train_params,
+                    train_loader=segmentation_test_dataloader(image_size=512),
+                    valid_loader=segmentation_test_dataloader(image_size=512),
+                )
             except Exception as e:
                 self.fail(f"Model training didn't succeed for {architecture} due to {e}")
             else:

tests/integration_tests/lr_test.py

@@ -13,23 +13,26 @@ class LRTest(unittest.TestCase):
     @classmethod
     def setUp(cls):
         # NAMES FOR THE EXPERIMENTS TO LATER DELETE
-        cls.folder_name = 'lr_test'
-        cls.training_params = {"max_epochs": 1,
-                               "silent_mode": True,
-                               "initial_lr": 0.1,
-                               "loss": "cross_entropy", "train_metrics_list": [Accuracy(), Top5()],
-                               "valid_metrics_list": [Accuracy(), Top5()],
-                               "metric_to_watch": "Accuracy",
-                               "greater_metric_to_watch_is_better": True}
+        cls.folder_name = "lr_test"
+        cls.training_params = {
+            "max_epochs": 1,
+            "silent_mode": True,
+            "initial_lr": 0.1,
+            "loss": "cross_entropy",
+            "train_metrics_list": [Accuracy(), Top5()],
+            "valid_metrics_list": [Accuracy(), Top5()],
+            "metric_to_watch": "Accuracy",
+            "greater_metric_to_watch_is_better": True,
+        }
 
     @classmethod
     def tearDownClass(cls) -> None:
         # ERASE THE FOLDER THAT WAS CREATED DURING THIS TEST
-        if os.path.isdir(os.path.join('checkpoints', cls.folder_name)):
-            shutil.rmtree(os.path.join('checkpoints', cls.folder_name))
+        if os.path.isdir(os.path.join("checkpoints", cls.folder_name)):
+            shutil.rmtree(os.path.join("checkpoints", cls.folder_name))
 
     @staticmethod
-    def get_trainer(name=''):
+    def get_trainer(name=""):
         trainer = Trainer(name)
         model = models.get("resnet18_cifar", num_classes=5)
         return trainer, model
@@ -42,26 +45,30 @@ def test_lr_function(initial_lr, epoch, iter, max_epoch, iters_per_epoch, **kwar
 
         # test if we are able that lr_function supports functions with this structure
         training_params = {**self.training_params, "lr_mode": "function", "lr_schedule_function": test_lr_function}
-        trainer.train(model=model, training_params=training_params, train_loader=classification_test_dataloader(),
-                      valid_loader=classification_test_dataloader())
+        trainer.train(
+            model=model, training_params=training_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader()
+        )
         # test that we assert lr_function is callable
         training_params = {**self.training_params, "lr_mode": "function"}
         with self.assertRaises(AssertionError):
-            trainer.train(model=model, training_params=training_params, train_loader=classification_test_dataloader(),
-                          valid_loader=classification_test_dataloader())
+            trainer.train(
+                model=model, training_params=training_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader()
+            )
 
     def test_cosine_lr(self):
         trainer, model = self.get_trainer(self.folder_name)
         training_params = {**self.training_params, "lr_mode": "cosine", "cosine_final_lr_ratio": 0.01}
-        trainer.train(model=model, training_params=training_params, train_loader=classification_test_dataloader(),
-                      valid_loader=classification_test_dataloader())
+        trainer.train(
+            model=model, training_params=training_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader()
+        )
 
     def test_step_lr(self):
         trainer, model = self.get_trainer(self.folder_name)
         training_params = {**self.training_params, "lr_mode": "step", "lr_decay_factor": 0.1, "lr_updates": [4]}
-        trainer.train(model=model, training_params=training_params, train_loader=classification_test_dataloader(),
-                      valid_loader=classification_test_dataloader())
+        trainer.train(
+            model=model, training_params=training_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader()
+        )
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()

tests/unit_tests/all_architectures_test.py

@@ -6,27 +6,30 @@
 
 
 class AllArchitecturesTest(unittest.TestCase):
-
     def setUp(self):
         # contains all arch_params needed for initialization of all architectures
-        self.all_arch_params = HpmStruct(**{'num_classes': 10,
-                                            'width_mult': 1,
-                                            'threshold': 1,
-                                            'sml_net': torch.nn.Identity(),
-                                            'big_net': torch.nn.Identity(),
-                                            'dropout': 0,
-                                            'build_residual_branches': True})
+        self.all_arch_params = HpmStruct(
+            **{
+                "num_classes": 10,
+                "width_mult": 1,
+                "threshold": 1,
+                "sml_net": torch.nn.Identity(),
+                "big_net": torch.nn.Identity(),
+                "dropout": 0,
+                "build_residual_branches": True,
+            }
+        )
 
     def test_architecture_is_sg_module(self):
         """
         Validate all models from all_architectures.py are SgModule
         """
         for arch_name in ARCHITECTURES:
             # skip custom constructors to keep all_arch_params as general as a possible
-            if 'custom' in arch_name.lower() or 'nas' in arch_name.lower() or 'kd' in arch_name.lower():
+            if "custom" in arch_name.lower() or "nas" in arch_name.lower() or "kd" in arch_name.lower():
                 continue
             self.assertTrue(isinstance(ARCHITECTURES[arch_name](arch_params=self.all_arch_params), SgModule))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()

tests/unit_tests/average_meter_test.py

@@ -14,7 +14,7 @@ def setUp(cls):
         cls.avg_tensor = AverageMeter()
         cls.left_empty = AverageMeter()
         cls.list_of_avg_meter = [cls.avg_float, cls.avg_tuple, cls.avg_list, cls.avg_tensor]
-        cls.score_types = [1.2, (3., 4.), [5., 6., 7.], torch.FloatTensor([8., 9., 10.])]
+        cls.score_types = [1.2, (3.0, 4.0), [5.0, 6.0, 7.0], torch.FloatTensor([8.0, 9.0, 10.0])]
         cls.batch_size = 3
 
     def test_empty_return_0(self):
@@ -42,5 +42,5 @@ def test_correctness_and_typing(self):
                     self.assertListEqual(list(avg_meter.average), list(score))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()

tests/unit_tests/cityscapes_dataset_test.py

@@ -4,18 +4,25 @@
 import yaml
 from torch.utils.data import DataLoader
 
-from super_gradients.training.dataloaders.dataloaders import cityscapes_train, cityscapes_val, \
-    cityscapes_stdc_seg50_train, cityscapes_stdc_seg50_val, cityscapes_stdc_seg75_val, cityscapes_ddrnet_train, \
-    cityscapes_regseg48_val, cityscapes_regseg48_train, cityscapes_ddrnet_val, cityscapes_stdc_seg75_train
+from super_gradients.training.dataloaders.dataloaders import (
+    cityscapes_train,
+    cityscapes_val,
+    cityscapes_stdc_seg50_train,
+    cityscapes_stdc_seg50_val,
+    cityscapes_stdc_seg75_val,
+    cityscapes_ddrnet_train,
+    cityscapes_regseg48_val,
+    cityscapes_regseg48_train,
+    cityscapes_ddrnet_val,
+    cityscapes_stdc_seg75_train,
+)
 from super_gradients.training.datasets.segmentation_datasets.cityscape_segmentation import CityscapesDataset
 
 
 class CityscapesDatasetTest(unittest.TestCase):
-
     def setUp(self) -> None:
-        default_config_path = pkg_resources.resource_filename("super_gradients.recipes",
-                                                              "dataset_params/cityscapes_dataset_params.yaml")
-        with open(default_config_path, 'r') as file:
+        default_config_path = pkg_resources.resource_filename("super_gradients.recipes", "dataset_params/cityscapes_dataset_params.yaml")
+        with open(default_config_path, "r") as file:
             self.recipe = yaml.safe_load(file)
 
     def dataloader_tester(self, dl: DataLoader):
@@ -26,12 +33,12 @@ def dataloader_tester(self, dl: DataLoader):
             next(it)
 
     def test_train_dataset_creation(self):
-        train_dataset = CityscapesDataset(**self.recipe['train_dataset_params'])
+        train_dataset = CityscapesDataset(**self.recipe["train_dataset_params"])
         for i in range(10):
             image, mask = train_dataset[i]
 
     def test_val_dataset_creation(self):
-        val_dataset = CityscapesDataset(**self.recipe['val_dataset_params'])
+        val_dataset = CityscapesDataset(**self.recipe["val_dataset_params"])
         for i in range(10):
             image, mask = val_dataset[i]
 
@@ -76,5 +83,5 @@ def test_cityscapes_ddrnet_val_dataloader(self):
         self.dataloader_tester(dl_val)
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     unittest.main()