From 9816729aa6f6e2f9708c59fd2aca1dfa3a4ebe34 Mon Sep 17 00:00:00 2001 From: Louis Dupont Date: Tue, 17 Jan 2023 15:45:00 +0200 Subject: [PATCH] apply black on tests --- ...deci_core_integration_test_suite_runner.py | 3 +- tests/end_to_end_tests/__init__.py | 2 +- tests/end_to_end_tests/trainer_test.py | 62 +++---- .../conversion_callback_test.py | 26 ++- tests/integration_tests/lr_test.py | 47 +++--- tests/unit_tests/all_architectures_test.py | 23 +-- tests/unit_tests/average_meter_test.py | 4 +- tests/unit_tests/cityscapes_dataset_test.py | 27 +-- .../coco_segmentation_dataset_test.py | 12 +- tests/unit_tests/conv_bn_relu_test.py | 43 ++--- tests/unit_tests/dataset_statistics_test.py | 31 ++-- tests/unit_tests/ddr_loss_test.py | 9 +- tests/unit_tests/detection_caching.py | 43 +++-- .../unit_tests/detection_sub_sampling_test.py | 9 +- tests/unit_tests/early_stop_test.py | 90 +++++----- tests/unit_tests/export_utils_test.py | 29 ++-- tests/unit_tests/forward_pass_prep_fn_test.py | 35 ++-- .../initialize_with_dataloaders_test.py | 37 +++-- tests/unit_tests/load_ema_ckpt_test.py | 37 +++-- tests/unit_tests/loss_loggings_test.py | 83 ++++++---- tests/unit_tests/lr_cooldown_test.py | 34 ++-- tests/unit_tests/mask_loss_test.py | 35 ++-- tests/unit_tests/multi_scaling_test.py | 3 +- tests/unit_tests/ohem_loss_test.py | 19 +-- tests/unit_tests/phase_context_test.py | 52 +++--- tests/unit_tests/phase_delegates_test.py | 71 +++++--- tests/unit_tests/random_erase_test.py | 10 +- tests/unit_tests/registry_test.py | 37 ++--- tests/unit_tests/regnet_unit_test.py | 131 +++++++++++---- tests/unit_tests/save_ckpt_test.py | 34 ++-- .../segmentation_transforms_test.py | 26 ++- tests/unit_tests/shelfnet_unit_test.py | 2 +- tests/unit_tests/test_auto_augment.py | 11 +- .../train_with_intialized_param_args_test.py | 154 ++++++++++++------ .../unit_tests/train_with_precise_bn_test.py | 73 ++++++--- .../training_params_factory_test.py | 2 +- .../update_param_groups_unit_test.py | 36 ++-- tests/unit_tests/yolox_unit_test.py | 2 +- .../zero_weight_decay_on_bias_bn_test.py | 109 +++++-------- 39 files changed, 859 insertions(+), 634 deletions(-) diff --git a/tests/deci_core_integration_test_suite_runner.py b/tests/deci_core_integration_test_suite_runner.py index 32a7daa463..b032e4a280 100644 --- a/tests/deci_core_integration_test_suite_runner.py +++ b/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() diff --git a/tests/end_to_end_tests/__init__.py b/tests/end_to_end_tests/__init__.py index d6cd72cddc..8ad5eb8169 100644 --- a/tests/end_to_end_tests/__init__.py +++ b/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"] diff --git a/tests/end_to_end_tests/trainer_test.py b/tests/end_to_end_tests/trainer_test.py index 21445bd3aa..7549c602c1 100644 --- a/tests/end_to_end_tests/trainer_test.py +++ b/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() diff --git a/tests/integration_tests/conversion_callback_test.py b/tests/integration_tests/conversion_callback_test.py index 9d888d1b1c..709287b4fc 100644 --- a/tests/integration_tests/conversion_callback_test.py +++ b/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,17 +82,14 @@ 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: @@ -102,8 +97,7 @@ def get_architecture_custom_config(architecture_name: str): 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: diff --git a/tests/integration_tests/lr_test.py b/tests/integration_tests/lr_test.py index cee92f0c10..0ef75df4ee 100644 --- a/tests/integration_tests/lr_test.py +++ b/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() diff --git a/tests/unit_tests/all_architectures_test.py b/tests/unit_tests/all_architectures_test.py index 45e51da6fb..491e9d0596 100644 --- a/tests/unit_tests/all_architectures_test.py +++ b/tests/unit_tests/all_architectures_test.py @@ -6,16 +6,19 @@ 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): """ @@ -23,10 +26,10 @@ def test_architecture_is_sg_module(self): """ 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() diff --git a/tests/unit_tests/average_meter_test.py b/tests/unit_tests/average_meter_test.py index 8f2da054be..e2ba41280e 100644 --- a/tests/unit_tests/average_meter_test.py +++ b/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() diff --git a/tests/unit_tests/cityscapes_dataset_test.py b/tests/unit_tests/cityscapes_dataset_test.py index 7b9c91c2ef..af82edcc80 100644 --- a/tests/unit_tests/cityscapes_dataset_test.py +++ b/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() diff --git a/tests/unit_tests/coco_segmentation_dataset_test.py b/tests/unit_tests/coco_segmentation_dataset_test.py index 1178b45f74..32b5b025bd 100644 --- a/tests/unit_tests/coco_segmentation_dataset_test.py +++ b/tests/unit_tests/coco_segmentation_dataset_test.py @@ -10,11 +10,9 @@ class CocoSegmentationDatasetTest(unittest.TestCase): - def setUp(self) -> None: - default_config_path = pkg_resources.resource_filename("super_gradients.recipes", - "dataset_params/coco_segmentation_dataset_params.yaml") - with open(default_config_path, 'r') as file: + default_config_path = pkg_resources.resource_filename("super_gradients.recipes", "dataset_params/coco_segmentation_dataset_params.yaml") + with open(default_config_path, "r") as file: self.recipe = yaml.safe_load(file) self.recipe = hydra.utils.instantiate(self.recipe) @@ -27,12 +25,12 @@ def dataloader_tester(self, dl: DataLoader): next(it) def test_train_dataset_creation(self): - train_dataset = CoCoSegmentationDataSet(**self.recipe['train_dataset_params']) + train_dataset = CoCoSegmentationDataSet(**self.recipe["train_dataset_params"]) for i in range(10): image, mask = train_dataset[i] def test_val_dataset_creation(self): - val_dataset = CoCoSegmentationDataSet(**self.recipe['val_dataset_params']) + val_dataset = CoCoSegmentationDataSet(**self.recipe["val_dataset_params"]) for i in range(10): image, mask = val_dataset[i] @@ -45,5 +43,5 @@ def test_coco_seg_val_dataloader(self): self.dataloader_tester(dl_val) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/conv_bn_relu_test.py b/tests/unit_tests/conv_bn_relu_test.py index 27853201cc..bbcabee60b 100644 --- a/tests/unit_tests/conv_bn_relu_test.py +++ b/tests/unit_tests/conv_bn_relu_test.py @@ -20,14 +20,20 @@ def test_conv_bn_relu(self): for stride in self.test_strides: for bias in self.biases: - conv_bn_relu = ConvBNReLU(32, 32, kernel_size=kernel, stride=stride, padding=kernel // 2, - bias=bias, use_activation=use_activation, - use_normalization=use_normalization) + conv_bn_relu = ConvBNReLU( + 32, + 32, + kernel_size=kernel, + stride=stride, + padding=kernel // 2, + bias=bias, + use_activation=use_activation, + use_normalization=use_normalization, + ) conv_bn_relu_seq = nn.Sequential( - nn.Conv2d(32, 32, kernel_size=kernel, stride=stride, padding=kernel // 2, - bias=bias), + nn.Conv2d(32, 32, kernel_size=kernel, stride=stride, padding=kernel // 2, bias=bias), nn.BatchNorm2d(32) if use_normalization else nn.Identity(), - nn.ReLU() if use_activation else nn.Identity() + nn.ReLU() if use_activation else nn.Identity(), ) # apply same conv weights and biases to compare output, # because conv weight and biases have random initialization. @@ -35,10 +41,12 @@ def test_conv_bn_relu(self): if bias: conv_bn_relu.seq[0].bias = conv_bn_relu_seq[0].bias - self.assertTrue(torch.equal(conv_bn_relu(self.sample), conv_bn_relu_seq(self.sample)), - msg=f"ConvBnRelu test failed for configuration: activation: " - f"{use_activation}, normalization: {use_normalization}, " - f"kernel: {kernel}, stride: {stride}") + self.assertTrue( + torch.equal(conv_bn_relu(self.sample), conv_bn_relu_seq(self.sample)), + msg=f"ConvBnRelu test failed for configuration: activation: " + f"{use_activation}, normalization: {use_normalization}, " + f"kernel: {kernel}, stride: {stride}", + ) def test_conv_bn_relu_with_default_torch_arguments(self): """ @@ -46,20 +54,17 @@ def test_conv_bn_relu_with_default_torch_arguments(self): Check that behavior of ConvBNRelu doesn't change with torch package upgrades. """ conv_bn_relu = ConvBNReLU(32, 32, kernel_size=1) - conv_bn_relu_defaults_torch = nn.Sequential( - nn.Conv2d(32, 32, kernel_size=1), - nn.BatchNorm2d(32), - nn.ReLU() - ) + conv_bn_relu_defaults_torch = nn.Sequential(nn.Conv2d(32, 32, kernel_size=1), nn.BatchNorm2d(32), nn.ReLU()) # apply same conv weights and biases to compare output, # because conv weight and biases have random initialization. conv_bn_relu.seq[0].weight = conv_bn_relu_defaults_torch[0].weight conv_bn_relu.seq[0].bias = conv_bn_relu_defaults_torch[0].bias - self.assertTrue(torch.equal(conv_bn_relu(self.sample), conv_bn_relu_defaults_torch(self.sample)), - msg="ConvBnRelu test failed for defaults arguments configuration: ConvBNRelu default" - "arguments are not aligned with torch defaults.") + self.assertTrue( + torch.equal(conv_bn_relu(self.sample), conv_bn_relu_defaults_torch(self.sample)), + msg="ConvBnRelu test failed for defaults arguments configuration: ConvBNRelu default" "arguments are not aligned with torch defaults.", + ) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/dataset_statistics_test.py b/tests/unit_tests/dataset_statistics_test.py index 99b6ef3338..ef63204664 100644 --- a/tests/unit_tests/dataset_statistics_test.py +++ b/tests/unit_tests/dataset_statistics_test.py @@ -8,7 +8,6 @@ class TestDatasetStatisticsTensorboardLogger(unittest.TestCase): - def test_dataset_statistics_tensorboard_logger(self): """ ** IMPORTANT NOTE ** @@ -18,27 +17,23 @@ def test_dataset_statistics_tensorboard_logger(self): """ # Create dataset - trainer = Trainer('dataset_statistics_visual_test') + trainer = Trainer("dataset_statistics_visual_test") model = models.get("yolox_s") - training_params = {"max_epochs": 1, # we dont really need the actual training to run - "lr_mode": "cosine", - "initial_lr": 0.01, - "loss": "yolox_loss", - "criterion_params": {"strides": [8, 16, 32], "num_classes": 80}, - "dataset_statistics": True, - "launch_tensorboard": True, - "valid_metrics_list": [ - DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), - normalize_targets=True, - num_cls=80)], - - - "metric_to_watch": "mAP@0.50:0.95", - } + training_params = { + "max_epochs": 1, # we dont really need the actual training to run + "lr_mode": "cosine", + "initial_lr": 0.01, + "loss": "yolox_loss", + "criterion_params": {"strides": [8, 16, 32], "num_classes": 80}, + "dataset_statistics": True, + "launch_tensorboard": True, + "valid_metrics_list": [DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), normalize_targets=True, num_cls=80)], + "metric_to_watch": "mAP@0.50:0.95", + } trainer.train(model=model, training_params=training_params, train_loader=coco2017_train(), valid_loader=coco2017_val()) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/ddr_loss_test.py b/tests/unit_tests/ddr_loss_test.py index 74dc8af316..50c0c03266 100644 --- a/tests/unit_tests/ddr_loss_test.py +++ b/tests/unit_tests/ddr_loss_test.py @@ -19,7 +19,7 @@ def test_without_auxiliary_loss(self): """ No Auxiliary loss, only one prediction map """ - weights = [1.] + weights = [1.0] criterion = DDRNetLoss((1 / self.num_classes - self.eps), ohem_percentage=0.1, weights=weights) bce_loss = -torch.log(torch.tensor(1 / self.num_classes)) expected_loss = bce_loss * weights[0] @@ -32,15 +32,14 @@ def test_with_auxiliary_loss(self): Auxiliary loss, 2 prediction maps, as DDRNet paper. """ predictions = [self.predictions, self.predictions] - weights = [1., 0.4] + weights = [1.0, 0.4] criterion = DDRNetLoss((1 / self.num_classes - self.eps), ohem_percentage=0.1, weights=weights) - expected_loss = -torch.log(torch.tensor(1 / self.num_classes)) * weights[0] + \ - -torch.log(torch.tensor(1 / self.num_classes)) * weights[1] + expected_loss = -torch.log(torch.tensor(1 / self.num_classes)) * weights[0] + -torch.log(torch.tensor(1 / self.num_classes)) * weights[1] loss, _ = criterion(predictions, self.targets) self.assertAlmostEqual(expected_loss, loss, delta=1e-5) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/detection_caching.py b/tests/unit_tests/detection_caching.py index 7bcee47ac9..0ccbbe5cd6 100644 --- a/tests/unit_tests/detection_caching.py +++ b/tests/unit_tests/detection_caching.py @@ -14,8 +14,8 @@ def __init__(self, input_dim, *args, **kwargs): self.image_size = input_dim self.n_samples = 321 - kwargs['all_classes_list'] = ["class_0", "class_1", "class_2"] - kwargs['original_target_format'] = DetectionTargetsFormat.XYXY_LABEL + kwargs["all_classes_list"] = ["class_0", "class_1", "class_2"] + kwargs["original_target_format"] = DetectionTargetsFormat.XYXY_LABEL super().__init__(input_dim=input_dim, *args, **kwargs) def _setup_data_source(self): @@ -26,10 +26,7 @@ def _load_annotation(self, sample_id: int) -> dict: a seed to allow the random image to the same for a given sample_id """ cls_id = sample_id % len(self.all_classes_list) - return {"img_path": str(sample_id), - "target": np.array([[0, 0, 10, 10, cls_id]]), - "resized_img_shape": self.image_size, - "seed": sample_id} + return {"img_path": str(sample_id), "target": np.array([[0, 0, 10, 10, cls_id]]), "resized_img_shape": self.image_size, "seed": sample_id} # We overwrite this to fake images def _load_image(self, index: int) -> np.ndarray: @@ -39,23 +36,29 @@ def _load_image(self, index: int) -> np.ndarray: class TestDetectionDatasetCaching(unittest.TestCase): def setUp(self) -> None: - self.temp_cache_dir = tempfile.TemporaryDirectory(prefix='cache').name + self.temp_cache_dir = tempfile.TemporaryDirectory(prefix="cache").name if not os.path.isdir(self.temp_cache_dir): os.mkdir(self.temp_cache_dir) def _count_cached_array(self): - return len(list(Path(self.temp_cache_dir).glob('*.array'))) + return len(list(Path(self.temp_cache_dir).glob("*.array"))) def _empty_cache(self): - for cache_file in Path(self.temp_cache_dir).glob('*.array'): + for cache_file in Path(self.temp_cache_dir).glob("*.array"): cache_file.unlink() def test_cache_keep_empty(self): self._empty_cache() datasets = [ - DummyDetectionDataset(input_dim=(640, 512), ignore_empty_annotations=False, class_inclusion_list=class_inclusion_list, - cache=True, cache_dir=self.temp_cache_dir, data_dir='/home/') + DummyDetectionDataset( + input_dim=(640, 512), + ignore_empty_annotations=False, + class_inclusion_list=class_inclusion_list, + cache=True, + cache_dir=self.temp_cache_dir, + data_dir="/home/", + ) for class_inclusion_list in [["class_0", "class_1", "class_2"], ["class_0"], ["class_1"], ["class_2"], ["class_1", "class_2"]] ] @@ -69,8 +72,14 @@ def test_cache_ignore_empty(self): self._empty_cache() datasets = [ - DummyDetectionDataset(input_dim=(640, 512), ignore_empty_annotations=True, class_inclusion_list=class_inclusion_list, - cache=True, cache_dir=self.temp_cache_dir, data_dir='/home/') + DummyDetectionDataset( + input_dim=(640, 512), + ignore_empty_annotations=True, + class_inclusion_list=class_inclusion_list, + cache=True, + cache_dir=self.temp_cache_dir, + data_dir="/home/", + ) for class_inclusion_list in [["class_0", "class_1", "class_2"], ["class_0"], ["class_1"], ["class_2"], ["class_1", "class_2"]] ] @@ -86,17 +95,15 @@ def test_cache_saved(self): self._empty_cache() self.assertEqual(0, self._count_cached_array()) - _ = DummyDetectionDataset(input_dim=(640, 512), ignore_empty_annotations=True, - cache=True, cache_dir=self.temp_cache_dir, data_dir='/home/') + _ = DummyDetectionDataset(input_dim=(640, 512), ignore_empty_annotations=True, cache=True, cache_dir=self.temp_cache_dir, data_dir="/home/") self.assertEqual(1, self._count_cached_array()) for _ in range(5): - _ = DummyDetectionDataset(input_dim=(640, 512), ignore_empty_annotations=True, - cache=True, cache_dir=self.temp_cache_dir, data_dir='/home/') + _ = DummyDetectionDataset(input_dim=(640, 512), ignore_empty_annotations=True, cache=True, cache_dir=self.temp_cache_dir, data_dir="/home/") self.assertEqual(1, self._count_cached_array()) self._empty_cache() -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/detection_sub_sampling_test.py b/tests/unit_tests/detection_sub_sampling_test.py index 2264c6057a..e99cb6676d 100644 --- a/tests/unit_tests/detection_sub_sampling_test.py +++ b/tests/unit_tests/detection_sub_sampling_test.py @@ -12,9 +12,9 @@ def __init__(self, dataset_size, input_dim, *args, **kwargs): self.dataset_size = dataset_size self.image_size = input_dim - kwargs['all_classes_list'] = ["class_0", "class_1", "class_2"] - kwargs['original_target_format'] = DetectionTargetsFormat.XYXY_LABEL - super().__init__(data_dir='', input_dim=input_dim, *args, **kwargs) + kwargs["all_classes_list"] = ["class_0", "class_1", "class_2"] + kwargs["original_target_format"] = DetectionTargetsFormat.XYXY_LABEL + super().__init__(data_dir="", input_dim=input_dim, *args, **kwargs) def _setup_data_source(self): return self.dataset_size @@ -30,7 +30,6 @@ def _load_image(self, index: int) -> np.ndarray: class TestDetectionDatasetSubsampling(unittest.TestCase): - def test_subsampling(self): """Check that subsampling works""" for max_num_samples in [1, 1_000, 1_000_000]: @@ -38,5 +37,5 @@ def test_subsampling(self): self.assertEqual(len(test_dataset), min(max_num_samples, 100_000)) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/early_stop_test.py b/tests/unit_tests/early_stop_test.py index e882c70647..1feeb6a9df 100644 --- a/tests/unit_tests/early_stop_test.py +++ b/tests/unit_tests/early_stop_test.py @@ -45,13 +45,23 @@ def setUp(self) -> None: # batch_size is equal to length of dataset, to have only one step per epoch, to ease the test. self.net = ResNet18(num_classes=5, arch_params={}) self.max_epochs = 10 - self.train_params = {"max_epochs": self.max_epochs, "lr_updates": [1], "lr_decay_factor": 0.1, - "lr_mode": "step", - "lr_warmup_epochs": 0, "initial_lr": 0.1, "loss": "cross_entropy", "optimizer": "SGD", - "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, - "train_metrics_list": [Accuracy()], "valid_metrics_list": [Top5()], - "metric_to_watch": "Top5", - "greater_metric_to_watch_is_better": True, "average_best_models": False} + self.train_params = { + "max_epochs": self.max_epochs, + "lr_updates": [1], + "lr_decay_factor": 0.1, + "lr_mode": "step", + "lr_warmup_epochs": 0, + "initial_lr": 0.1, + "loss": "cross_entropy", + "optimizer": "SGD", + "criterion_params": {}, + "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, + "train_metrics_list": [Accuracy()], + "valid_metrics_list": [Top5()], + "metric_to_watch": "Top5", + "greater_metric_to_watch_is_better": True, + "average_best_models": False, + } def test_min_mode_patience_metric(self): """ @@ -63,13 +73,14 @@ def test_min_mode_patience_metric(self): early_stop_loss = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="LossTest", mode="min", patience=3, verbose=True) phase_callbacks = [early_stop_loss] - loss_values = torch.tensor([1., 0.8, 0.81, 0.8, 0.9, 0.2, 0.1, 0.3, 0.05, 0.9]) + loss_values = torch.tensor([1.0, 0.8, 0.81, 0.8, 0.9, 0.2, 0.1, 0.3, 0.05, 0.9]) fake_loss = LossTest(loss_values) train_params = self.train_params.copy() train_params.update({"loss": fake_loss, "phase_callbacks": phase_callbacks}) - trainer.train(model=self.net, training_params=train_params, train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + trainer.train( + model=self.net, training_params=train_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader() + ) excepted_end_epoch = 5 # count divided by 2, because loss counter used for both train and eval. @@ -81,18 +92,17 @@ def test_max_mode_patience_metric(self): epochs. """ trainer = Trainer("early_stop_test") - early_stop_acc = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="MetricTest", mode="max", patience=3, - verbose=True) + early_stop_acc = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="MetricTest", mode="max", patience=3, verbose=True) phase_callbacks = [early_stop_acc] metric_values = torch.tensor([0.2, 0.1, 0.3, 0.28, 0.2, 0.1, 0.33, 0.05, 0.9, 0.99]) fake_metric = MetricTest(metric_values) train_params = self.train_params.copy() - train_params.update( - {"valid_metrics_list": [fake_metric], "metric_to_watch": "MetricTest", "phase_callbacks": phase_callbacks}) + train_params.update({"valid_metrics_list": [fake_metric], "metric_to_watch": "MetricTest", "phase_callbacks": phase_callbacks}) - trainer.train(model=self.net, training_params=train_params, train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + trainer.train( + model=self.net, training_params=train_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader() + ) excepted_end_epoch = 6 self.assertEqual(excepted_end_epoch, fake_metric.count) @@ -106,13 +116,14 @@ def test_min_mode_threshold_metric(self): early_stop_loss = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="LossTest", mode="min", threshold=0.1, verbose=True) phase_callbacks = [early_stop_loss] - loss_values = torch.tensor([1., 0.8, 0.4, 0.2, 0.09, 0.11, 0.105, 0.3, 0.05, 0.02]) + loss_values = torch.tensor([1.0, 0.8, 0.4, 0.2, 0.09, 0.11, 0.105, 0.3, 0.05, 0.02]) fake_loss = LossTest(loss_values) train_params = self.train_params.copy() train_params.update({"loss": fake_loss, "phase_callbacks": phase_callbacks}) - trainer.train(model=self.net, training_params=train_params, train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + trainer.train( + model=self.net, training_params=train_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader() + ) excepted_end_epoch = 5 # count divided by 2, because loss counter used for both train and eval. self.assertEqual(excepted_end_epoch, fake_loss.count // 2) @@ -123,18 +134,17 @@ def test_max_mode_threshold_metric(self): """ trainer = Trainer("early_stop_test") - early_stop_acc = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="MetricTest", mode="max", threshold=0.94, - verbose=True) + early_stop_acc = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="MetricTest", mode="max", threshold=0.94, verbose=True) phase_callbacks = [early_stop_acc] metric_values = torch.tensor([0.2, 0.1, 0.6, 0.8, 0.9, 0.92, 0.95, 0.94, 0.948, 0.99]) fake_metric = MetricTest(metric_values) train_params = self.train_params.copy() - train_params.update( - {"valid_metrics_list": [fake_metric], "metric_to_watch": "MetricTest", "phase_callbacks": phase_callbacks}) + train_params.update({"valid_metrics_list": [fake_metric], "metric_to_watch": "MetricTest", "phase_callbacks": phase_callbacks}) - trainer.train(model=self.net, training_params=train_params, train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + trainer.train( + model=self.net, training_params=train_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader() + ) excepted_end_epoch = 7 self.assertEqual(excepted_end_epoch, fake_metric.count) @@ -146,17 +156,17 @@ def test_no_finite_stoppage(self): # test Nan value trainer = Trainer("early_stop_test") - early_stop_loss = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="LossTest", mode="min", check_finite=True, - verbose=True) + early_stop_loss = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="LossTest", mode="min", check_finite=True, verbose=True) phase_callbacks = [early_stop_loss] - loss_values = torch.tensor([1., float('nan'), 0.81, 0.8, 0.9, 0.2, 0.1, 0.3, 0.05, 0.9]) + loss_values = torch.tensor([1.0, float("nan"), 0.81, 0.8, 0.9, 0.2, 0.1, 0.3, 0.05, 0.9]) fake_loss = LossTest(loss_values) train_params = self.train_params.copy() train_params.update({"loss": fake_loss, "phase_callbacks": phase_callbacks}) - trainer.train(model=self.net, training_params=train_params, train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + trainer.train( + model=self.net, training_params=train_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader() + ) excepted_end_epoch = 2 self.assertEqual(excepted_end_epoch, fake_loss.count // 2) @@ -167,13 +177,14 @@ def test_no_finite_stoppage(self): early_stop_loss = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="LossTest", mode="min", patience=3, verbose=True) phase_callbacks = [early_stop_loss] - loss_values = torch.tensor([1., 0.8, float('inf'), 0.8, 0.9, 0.2, 0.1, 0.3, 0.05, 0.9]) + loss_values = torch.tensor([1.0, 0.8, float("inf"), 0.8, 0.9, 0.2, 0.1, 0.3, 0.05, 0.9]) fake_loss = LossTest(loss_values) train_params = self.train_params.copy() train_params.update({"loss": fake_loss, "phase_callbacks": phase_callbacks}) - trainer.train(model=self.net, training_params=train_params, train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + trainer.train( + model=self.net, training_params=train_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader() + ) excepted_end_epoch = 3 # count divided by 2, because loss counter used for both train and eval. self.assertEqual(excepted_end_epoch, fake_loss.count // 2) @@ -185,23 +196,22 @@ def test_min_delta(self): """ trainer = Trainer("early_stop_test") - early_stop_acc = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="MetricTest", mode="max", patience=2, - min_delta=0.1, verbose=True) + early_stop_acc = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="MetricTest", mode="max", patience=2, min_delta=0.1, verbose=True) phase_callbacks = [early_stop_acc] metric_values = torch.tensor([0.1, 0.2, 0.305, 0.31, 0.34, 0.42, 0.6, 0.8, 0.9, 0.99]) fake_metric = MetricTest(metric_values) train_params = self.train_params.copy() - train_params.update( - {"valid_metrics_list": [fake_metric], "metric_to_watch": "MetricTest", "phase_callbacks": phase_callbacks}) + train_params.update({"valid_metrics_list": [fake_metric], "metric_to_watch": "MetricTest", "phase_callbacks": phase_callbacks}) - trainer.train(model=self.net, training_params=train_params, train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + trainer.train( + model=self.net, training_params=train_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader() + ) excepted_end_epoch = 5 self.assertEqual(excepted_end_epoch, fake_metric.count) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/export_utils_test.py b/tests/unit_tests/export_utils_test.py index b99cd92145..575ad63326 100644 --- a/tests/unit_tests/export_utils_test.py +++ b/tests/unit_tests/export_utils_test.py @@ -14,18 +14,17 @@ def count_parameters(model): class TestUtil(unittest.TestCase): - def test_fuse_conv_bn_real_archs(self): """ test the fuse_conv_bn function. run the function on some Sg architectures and assert the result of the original net are the same as the results of the fused net """ - archs = ['resnet18', 'mobilenet_v2', 'densenet121', 'regnetY200'] + archs = ["resnet18", "mobilenet_v2", "densenet121", "regnetY200"] for arch_name in archs: - model1 = ARCHITECTURES[arch_name](HpmStruct(**{'num_classes': 10, 'dropout': 0.1})) + model1 = ARCHITECTURES[arch_name](HpmStruct(**{"num_classes": 10, "dropout": 0.1})) model2 = copy.deepcopy(model1) model1.eval() @@ -41,7 +40,7 @@ def test_fuse_conv_bn_real_archs(self): param_count2 = count_parameters(model2) self.assertTrue(torch.allclose(output1, output2, atol=1e-6)) - print(f'Tested fuse Conv BN on {arch_name}: OK ({param_count1 - param_count2} less params)') + print(f"Tested fuse Conv BN on {arch_name}: OK ({param_count1 - param_count2} less params)") def test_fuse_conv_bn_on_sequential_models(self): @@ -50,35 +49,33 @@ def test_fuse_conv_bn_on_sequential_models(self): model.eval() fuse_conv_bn(model, replace_bn_with_identity=True) self.assertEqual(len(model._modules), 2) - self.assertIsInstance(model._modules['0'], nn.Conv2d) - self.assertIsInstance(model._modules['1'], nn.Identity) + self.assertIsInstance(model._modules["0"], nn.Conv2d) + self.assertIsInstance(model._modules["1"], nn.Identity) # assert the bn module was removed model = nn.Sequential(nn.Conv2d(3, 3, 3), nn.BatchNorm2d(3)) model.eval() fuse_conv_bn(model, replace_bn_with_identity=False) self.assertEqual(len(model._modules), 1) - self.assertIsInstance(model._modules['0'], nn.Conv2d) + self.assertIsInstance(model._modules["0"], nn.Conv2d) # assert all bn module were removed model = nn.Sequential(nn.Conv2d(3, 3, 3), nn.BatchNorm2d(3), nn.Conv2d(3, 3, 3), nn.BatchNorm2d(3)) model.eval() fuse_conv_bn(model, replace_bn_with_identity=False) self.assertEqual(len(model._modules), 2) - self.assertIsInstance(model._modules['0'], nn.Conv2d) + self.assertIsInstance(model._modules["0"], nn.Conv2d) # assert only merged bn module were removed model = nn.Sequential(nn.Conv2d(3, 3, 3), nn.Conv2d(3, 3, 3), nn.BatchNorm2d(3)) model.eval() fuse_conv_bn(model, replace_bn_with_identity=False) self.assertEqual(len(model._modules), 2) - self.assertIsInstance(model._modules['0'], nn.Conv2d) - self.assertIsInstance(model._modules['1'], nn.Conv2d) + self.assertIsInstance(model._modules["0"], nn.Conv2d) + self.assertIsInstance(model._modules["1"], nn.Conv2d) def test_fuse_conv_bn_on_toy_models(self): - class Toy(nn.Module): - def __init__(self): super().__init__() self.conv1 = nn.Conv2d(3, 3, 3) @@ -105,16 +102,16 @@ def forward(self, x): model = Toy() model.eval() fuse_conv_bn(model, replace_bn_with_identity=False) - self.assertFalse(hasattr(model, 'bn1')) + self.assertFalse(hasattr(model, "bn1")) self.assertIsInstance(model.conv1, nn.Conv2d) # assert all bn module were removed model = Toy() model.eval() fuse_conv_bn(model, replace_bn_with_identity=False) - self.assertFalse(hasattr(model, 'bn1')) + self.assertFalse(hasattr(model, "bn1")) self.assertIsInstance(model.conv1, nn.Conv2d) - self.assertFalse(hasattr(model, 'bn2')) + self.assertFalse(hasattr(model, "bn2")) self.assertIsInstance(model.conv2, nn.Conv2d) # assert correct number of parameters removed @@ -126,5 +123,5 @@ def forward(self, x): self.assertEqual(before - after, 12) # each bn of 3 channels has 6 parameters (12 together) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/forward_pass_prep_fn_test.py b/tests/unit_tests/forward_pass_prep_fn_test.py index d7701f1277..868d0f2ee8 100644 --- a/tests/unit_tests/forward_pass_prep_fn_test.py +++ b/tests/unit_tests/forward_pass_prep_fn_test.py @@ -20,14 +20,11 @@ def __call__(self, context: PhaseContext): def test_forward_pass_prep_fn(inputs, targets, *args, **kwargs): - inputs = torch.nn.functional.interpolate( - inputs, size=(50, 50), mode="bilinear", align_corners=False - ) + inputs = torch.nn.functional.interpolate(inputs, size=(50, 50), mode="bilinear", align_corners=False) return inputs, targets class ForwardpassPrepFNTest(unittest.TestCase): - def test_resizing_with_forward_pass_prep_fn(self): # Define Model trainer = Trainer("ForwardpassPrepFNTest") @@ -36,16 +33,26 @@ def test_resizing_with_forward_pass_prep_fn(self): sizes = [] phase_callbacks = [TestInputSizesCallback(sizes)] - train_params = {"max_epochs": 2, "cosine_final_lr_ratio": 0.2, "lr_mode": "cosine", - "lr_cooldown_epochs": 2, - "lr_warmup_epochs": 3, "initial_lr": 1, "loss": "cross_entropy", "optimizer": 'SGD', - "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, - "train_metrics_list": [Accuracy()], "valid_metrics_list": [Accuracy()], - "metric_to_watch": "Accuracy", - "greater_metric_to_watch_is_better": True, "ema": False, "phase_callbacks": phase_callbacks, - "pre_prediction_callback": test_forward_pass_prep_fn} - trainer.train(model=model, training_params=train_params, train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + train_params = { + "max_epochs": 2, + "cosine_final_lr_ratio": 0.2, + "lr_mode": "cosine", + "lr_cooldown_epochs": 2, + "lr_warmup_epochs": 3, + "initial_lr": 1, + "loss": "cross_entropy", + "optimizer": "SGD", + "criterion_params": {}, + "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, + "train_metrics_list": [Accuracy()], + "valid_metrics_list": [Accuracy()], + "metric_to_watch": "Accuracy", + "greater_metric_to_watch_is_better": True, + "ema": False, + "phase_callbacks": phase_callbacks, + "pre_prediction_callback": test_forward_pass_prep_fn, + } + trainer.train(model=model, training_params=train_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader()) # ALTHOUGH NOT SEEN IN HERE, THE 4TH EPOCH USES LR=1, SO THIS IS THE EXPECTED LIST AS WE COLLECT # THE LRS AFTER THE UPDATE diff --git a/tests/unit_tests/initialize_with_dataloaders_test.py b/tests/unit_tests/initialize_with_dataloaders_test.py index c1efce8294..365f1efa31 100644 --- a/tests/unit_tests/initialize_with_dataloaders_test.py +++ b/tests/unit_tests/initialize_with_dataloaders_test.py @@ -28,24 +28,27 @@ def setUp(self): def test_train_with_dataloaders(self): trainer = Trainer(experiment_name="test_name") model = models.get("resnet18", num_classes=5) - trainer.train(model=model, - training_params={"max_epochs": 2, - "lr_updates": [5, 6, 12], - "lr_decay_factor": 0.01, - "lr_mode": "step", - "initial_lr": 0.01, - "loss": "cross_entropy", - "optimizer": "SGD", - "optimizer_params": {"weight_decay": 1e-5, "momentum": 0.9}, - "train_metrics_list": [Accuracy()], - "valid_metrics_list": [Accuracy()], - "metric_to_watch": "Accuracy", - "greater_metric_to_watch_is_better": True}, - train_loader=self.testcase_trainloader, - valid_loader=self.testcase_validloader, - ) + trainer.train( + model=model, + training_params={ + "max_epochs": 2, + "lr_updates": [5, 6, 12], + "lr_decay_factor": 0.01, + "lr_mode": "step", + "initial_lr": 0.01, + "loss": "cross_entropy", + "optimizer": "SGD", + "optimizer_params": {"weight_decay": 1e-5, "momentum": 0.9}, + "train_metrics_list": [Accuracy()], + "valid_metrics_list": [Accuracy()], + "metric_to_watch": "Accuracy", + "greater_metric_to_watch_is_better": True, + }, + train_loader=self.testcase_trainloader, + valid_loader=self.testcase_validloader, + ) self.assertTrue(0 < trainer.best_metric.item() < 1) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/load_ema_ckpt_test.py b/tests/unit_tests/load_ema_ckpt_test.py index 4d09f4fa3b..b070c8d862 100644 --- a/tests/unit_tests/load_ema_ckpt_test.py +++ b/tests/unit_tests/load_ema_ckpt_test.py @@ -19,20 +19,31 @@ def __call__(self, context: PhaseContext): class LoadCheckpointWithEmaTest(unittest.TestCase): def setUp(self) -> None: - self.train_params = {"max_epochs": 2, "lr_updates": [1], "lr_decay_factor": 0.1, "lr_mode": "step", - "lr_warmup_epochs": 0, "initial_lr": 0.1, "loss": "cross_entropy", "optimizer": 'SGD', - "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, - "train_metrics_list": [Accuracy(), Top5()], "valid_metrics_list": [Accuracy(), Top5()], - "metric_to_watch": "Accuracy", - "greater_metric_to_watch_is_better": True, "ema": True} + self.train_params = { + "max_epochs": 2, + "lr_updates": [1], + "lr_decay_factor": 0.1, + "lr_mode": "step", + "lr_warmup_epochs": 0, + "initial_lr": 0.1, + "loss": "cross_entropy", + "optimizer": "SGD", + "criterion_params": {}, + "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, + "train_metrics_list": [Accuracy(), Top5()], + "valid_metrics_list": [Accuracy(), Top5()], + "metric_to_watch": "Accuracy", + "greater_metric_to_watch_is_better": True, + "ema": True, + } def test_ema_ckpt_reload(self): # Define Model net = LeNet() trainer = Trainer("ema_ckpt_test") - trainer.train(model=net, training_params=self.train_params, - train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + trainer.train( + model=net, training_params=self.train_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader() + ) ema_model = trainer.ema_model.ema @@ -44,9 +55,9 @@ def test_ema_ckpt_reload(self): self.train_params["resume"] = True self.train_params["max_epochs"] = 3 self.train_params["phase_callbacks"] = [net_collector] - trainer.train(model=net, training_params=self.train_params, - train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + trainer.train( + model=net, training_params=self.train_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader() + ) reloaded_ema_model = net_collector.net.ema @@ -54,5 +65,5 @@ def test_ema_ckpt_reload(self): assert check_models_have_same_weights(ema_model, reloaded_ema_model) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/loss_loggings_test.py b/tests/unit_tests/loss_loggings_test.py index 56eaa0d790..882b97d391 100644 --- a/tests/unit_tests/loss_loggings_test.py +++ b/tests/unit_tests/loss_loggings_test.py @@ -26,52 +26,77 @@ def __init__(self): class LossLoggingsTest(unittest.TestCase): def test_single_item_logging(self): - trainer = Trainer("test_single_item_logging", model_checkpoints_location='local') + trainer = Trainer("test_single_item_logging", model_checkpoints_location="local") dataloader = classification_test_dataloader(batch_size=10) model = models.get("resnet18", arch_params={"num_classes": 5}) - train_params = {"max_epochs": 1, "lr_updates": [1], "lr_decay_factor": 0.1, "lr_mode": "step", - "lr_warmup_epochs": 0, "initial_lr": 0.1, "loss": torch.nn.CrossEntropyLoss(), - "optimizer": "SGD", - "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, - "train_metrics_list": [Accuracy()], "valid_metrics_list": [Accuracy()], - "metric_to_watch": "Accuracy", - "greater_metric_to_watch_is_better": True} + train_params = { + "max_epochs": 1, + "lr_updates": [1], + "lr_decay_factor": 0.1, + "lr_mode": "step", + "lr_warmup_epochs": 0, + "initial_lr": 0.1, + "loss": torch.nn.CrossEntropyLoss(), + "optimizer": "SGD", + "criterion_params": {}, + "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, + "train_metrics_list": [Accuracy()], + "valid_metrics_list": [Accuracy()], + "metric_to_watch": "Accuracy", + "greater_metric_to_watch_is_better": True, + } trainer.train(model=model, training_params=train_params, train_loader=dataloader, valid_loader=dataloader) self.assertListEqual(trainer.loss_logging_items_names, ["CrossEntropyLoss"]) def test_multiple_unnamed_components_loss_logging(self): - trainer = Trainer("test_multiple_unnamed_components_loss_logging", model_checkpoints_location='local') + trainer = Trainer("test_multiple_unnamed_components_loss_logging", model_checkpoints_location="local") dataloader = classification_test_dataloader(batch_size=10) model = models.get("resnet18", arch_params={"num_classes": 5}) - train_params = {"max_epochs": 1, "lr_updates": [1], "lr_decay_factor": 0.1, "lr_mode": "step", - "lr_warmup_epochs": 0, "initial_lr": 0.1, "loss": CriterionWithUnnamedComponents(), - "optimizer": "SGD", - "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, - "train_metrics_list": [Accuracy()], "valid_metrics_list": [Accuracy()], - "metric_to_watch": "Accuracy", - "greater_metric_to_watch_is_better": True} + train_params = { + "max_epochs": 1, + "lr_updates": [1], + "lr_decay_factor": 0.1, + "lr_mode": "step", + "lr_warmup_epochs": 0, + "initial_lr": 0.1, + "loss": CriterionWithUnnamedComponents(), + "optimizer": "SGD", + "criterion_params": {}, + "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, + "train_metrics_list": [Accuracy()], + "valid_metrics_list": [Accuracy()], + "metric_to_watch": "Accuracy", + "greater_metric_to_watch_is_better": True, + } trainer.train(model=model, training_params=train_params, train_loader=dataloader, valid_loader=dataloader) - self.assertListEqual(trainer.loss_logging_items_names, ["CriterionWithUnnamedComponents/loss_0", - "CriterionWithUnnamedComponents/loss_1"]) + self.assertListEqual(trainer.loss_logging_items_names, ["CriterionWithUnnamedComponents/loss_0", "CriterionWithUnnamedComponents/loss_1"]) def test_multiple_named_components_loss_logging(self): - trainer = Trainer("test_multiple_named_components_loss_logging", model_checkpoints_location='local') + trainer = Trainer("test_multiple_named_components_loss_logging", model_checkpoints_location="local") dataloader = classification_test_dataloader(batch_size=10) model = models.get("resnet18", arch_params={"num_classes": 5}) - train_params = {"max_epochs": 1, "lr_updates": [1], "lr_decay_factor": 0.1, "lr_mode": "step", - "lr_warmup_epochs": 0, "initial_lr": 0.1, "loss": CriterionWithNamedComponents(), - "optimizer": "SGD", - "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, - "train_metrics_list": [Accuracy()], "valid_metrics_list": [Accuracy()], - "metric_to_watch": "Accuracy", - "greater_metric_to_watch_is_better": True} + train_params = { + "max_epochs": 1, + "lr_updates": [1], + "lr_decay_factor": 0.1, + "lr_mode": "step", + "lr_warmup_epochs": 0, + "initial_lr": 0.1, + "loss": CriterionWithNamedComponents(), + "optimizer": "SGD", + "criterion_params": {}, + "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, + "train_metrics_list": [Accuracy()], + "valid_metrics_list": [Accuracy()], + "metric_to_watch": "Accuracy", + "greater_metric_to_watch_is_better": True, + } trainer.train(model=model, training_params=train_params, train_loader=dataloader, valid_loader=dataloader) - self.assertListEqual(trainer.loss_logging_items_names, ["CriterionWithNamedComponents/loss_A", - "CriterionWithNamedComponents/loss_B"]) + self.assertListEqual(trainer.loss_logging_items_names, ["CriterionWithNamedComponents/loss_A", "CriterionWithNamedComponents/loss_B"]) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/lr_cooldown_test.py b/tests/unit_tests/lr_cooldown_test.py index 053073b7db..746c56a31c 100644 --- a/tests/unit_tests/lr_cooldown_test.py +++ b/tests/unit_tests/lr_cooldown_test.py @@ -15,18 +15,32 @@ def test_lr_cooldown_with_lr_scheduling(self): lrs = [] phase_callbacks = [TestLRCallback(lr_placeholder=lrs)] - train_params = {"max_epochs": 7, "cosine_final_lr_ratio": 0.2, "lr_mode": "cosine", - "lr_cooldown_epochs": 2, - "lr_warmup_epochs": 3, "initial_lr": 1, "loss": "cross_entropy", "optimizer": 'SGD', - "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, - "train_metrics_list": [Accuracy()], "valid_metrics_list": [Accuracy()], - "metric_to_watch": "Accuracy", - "greater_metric_to_watch_is_better": True, "ema": False, "phase_callbacks": phase_callbacks} + train_params = { + "max_epochs": 7, + "cosine_final_lr_ratio": 0.2, + "lr_mode": "cosine", + "lr_cooldown_epochs": 2, + "lr_warmup_epochs": 3, + "initial_lr": 1, + "loss": "cross_entropy", + "optimizer": "SGD", + "criterion_params": {}, + "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, + "train_metrics_list": [Accuracy()], + "valid_metrics_list": [Accuracy()], + "metric_to_watch": "Accuracy", + "greater_metric_to_watch_is_better": True, + "ema": False, + "phase_callbacks": phase_callbacks, + } expected_lrs = [0.25, 0.5, 0.75, 0.9236067977499791, 0.4763932022500211, 0.4763932022500211, 0.4763932022500211] - trainer.train(model=net, training_params=train_params, - train_loader=classification_test_dataloader(dataset_size=5, batch_size=4), - valid_loader=classification_test_dataloader(dataset_size=5, batch_size=4)) + trainer.train( + model=net, + training_params=train_params, + train_loader=classification_test_dataloader(dataset_size=5, batch_size=4), + valid_loader=classification_test_dataloader(dataset_size=5, batch_size=4), + ) # ALTHOUGH NOT SEEN IN HERE, THE 4TH EPOCH USES LR=1, SO THIS IS THE EXPECTED LIST AS WE COLLECT # THE LRS AFTER THE UPDATE diff --git a/tests/unit_tests/mask_loss_test.py b/tests/unit_tests/mask_loss_test.py index 69d8685df2..7febdb3383 100644 --- a/tests/unit_tests/mask_loss_test.py +++ b/tests/unit_tests/mask_loss_test.py @@ -6,7 +6,6 @@ class MaskAttentionLossTest(unittest.TestCase): - def setUp(self) -> None: self.img_size = 32 self.num_classes = 4 @@ -22,14 +21,11 @@ def _get_default_target_tensor(self): def _get_default_mask_tensor(self): mask = torch.zeros(self.batch, 1, self.img_size, self.img_size) # half tensor rows as 1 - mask[:, :, self.img_size // 2:] = 1 + mask[:, :, self.img_size // 2 :] = 1 return mask.float() def _assertion_torch_values(self, expected_value: torch.Tensor, found_value: torch.Tensor, rtol: float = 1e-5): - self.assertTrue( - torch.allclose(found_value, expected_value, rtol=rtol), - msg=f"Unequal torch tensors: excepted: {expected_value}, found: {found_value}" - ) + self.assertTrue(torch.allclose(found_value, expected_value, rtol=rtol), msg=f"Unequal torch tensors: excepted: {expected_value}, found: {found_value}") def test_with_cross_entropy_loss(self): """ @@ -40,14 +36,14 @@ def test_with_cross_entropy_loss(self): target = self._get_default_target_tensor() mask = self._get_default_mask_tensor() - loss_weigths = [1., 0.5] + loss_weigths = [1.0, 0.5] ce_crit = nn.CrossEntropyLoss(reduction="none") mask_ce_crit = MaskAttentionLoss(criterion=ce_crit, loss_weights=loss_weigths) # expected result ce_loss = ce_crit(predict, target) _mask = mask.view_as(ce_loss) - mask_loss = (ce_loss * _mask) + mask_loss = ce_loss * _mask mask_loss = mask_loss[_mask == 1] # consider only mask samples for mask loss computing expected_loss = ce_loss.mean() * loss_weigths[0] + mask_loss.mean() * loss_weigths[1] @@ -65,14 +61,14 @@ def test_with_binary_cross_entropy_loss(self): target = torch.randn(self.batch, self.num_classes, self.img_size, self.img_size) mask = self._get_default_mask_tensor() - loss_weigths = [1., 0.5] + loss_weigths = [1.0, 0.5] ce_crit = nn.BCEWithLogitsLoss(reduction="none") mask_ce_crit = MaskAttentionLoss(criterion=ce_crit, loss_weights=loss_weigths) # expected result ce_loss = ce_crit(predict, target) _mask = mask.expand_as(ce_loss) - mask_loss = (ce_loss * _mask) + mask_loss = ce_loss * _mask mask_loss = mask_loss[_mask == 1] # consider only mask samples for mask loss computing expected_loss = ce_loss.mean() * loss_weigths[0] + mask_loss.mean() * loss_weigths[1] @@ -90,14 +86,14 @@ def test_reduction_none(self): target = self._get_default_target_tensor() mask = self._get_default_mask_tensor() - loss_weigths = [1., 0.5] + loss_weigths = [1.0, 0.5] ce_crit = nn.CrossEntropyLoss(reduction="none") mask_ce_crit = MaskAttentionLoss(criterion=ce_crit, loss_weights=loss_weigths, reduction="none") # expected result ce_loss = ce_crit(predict, target) _mask = mask.view_as(ce_loss) - mask_loss = (ce_loss * _mask) + mask_loss = ce_loss * _mask expected_loss = ce_loss * loss_weigths[0] + mask_loss * loss_weigths[1] # mask ce loss result @@ -111,10 +107,10 @@ def test_assert_valid_arguments(self): kwargs = {"criterion": nn.CrossEntropyLoss(reduction="mean")} self.failUnlessRaises(ValueError, MaskAttentionLoss, **kwargs) # loss_weights must have only 2 values - kwargs = {"criterion": nn.CrossEntropyLoss(reduction="none"), "loss_weights": [1., 1., 1.]} + kwargs = {"criterion": nn.CrossEntropyLoss(reduction="none"), "loss_weights": [1.0, 1.0, 1.0]} self.failUnlessRaises(ValueError, MaskAttentionLoss, **kwargs) # mask loss_weight must be a positive value - kwargs = {"criterion": nn.CrossEntropyLoss(reduction="none"), "loss_weights": [1., 0.]} + kwargs = {"criterion": nn.CrossEntropyLoss(reduction="none"), "loss_weights": [1.0, 0.0]} self.failUnlessRaises(ValueError, MaskAttentionLoss, **kwargs) def test_multi_class_mask(self): @@ -128,13 +124,13 @@ def test_multi_class_mask(self): target = to_one_hot(target, self.num_classes).float() mask = torch.randint(0, 2, size=(self.batch, self.num_classes, self.img_size, self.img_size)).float() - loss_weigths = [1., 0.5] + loss_weigths = [1.0, 0.5] ce_crit = nn.MSELoss(reduction="none") mask_ce_crit = MaskAttentionLoss(criterion=ce_crit, loss_weights=loss_weigths) # expected result mse_loss = ce_crit(predict, target) - mask_loss = (mse_loss * mask) + mask_loss = mse_loss * mask mask_loss = mask_loss[mask == 1] # consider only mask samples for mask loss computing expected_loss = mse_loss.mean() * loss_weigths[0] + mask_loss.mean() * loss_weigths[1] @@ -148,10 +144,9 @@ def test_broadcast_exceptions(self): Test assertion in mask broadcasting """ predict = torch.randn(self.batch, self.num_classes, self.img_size, self.img_size) - target = torch.randint(0, self.num_classes, - size=(self.batch, self.num_classes, self.img_size, self.img_size)).float() + target = torch.randint(0, self.num_classes, size=(self.batch, self.num_classes, self.img_size, self.img_size)).float() - loss_weigths = [1., 0.5] + loss_weigths = [1.0, 0.5] ce_crit = nn.BCEWithLogitsLoss(reduction="none") mask_ce_crit = MaskAttentionLoss(criterion=ce_crit, loss_weights=loss_weigths) @@ -168,5 +163,5 @@ def test_broadcast_exceptions(self): self.failUnlessRaises(AssertionError, mask_ce_crit, *(predict, target, mask)) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/multi_scaling_test.py b/tests/unit_tests/multi_scaling_test.py index 317d0b407f..6f9b38e32d 100644 --- a/tests/unit_tests/multi_scaling_test.py +++ b/tests/unit_tests/multi_scaling_test.py @@ -5,7 +5,6 @@ class MultiScaleTest(unittest.TestCase): - def setUp(self) -> None: self.size = (1024, 512) self.batch_size = 12 @@ -31,5 +30,5 @@ def test_multiscale_keep_state(self): self.assertListEqual(post_multiscale_input_shapes[0], post_multiscale_input_shapes[-1]) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/ohem_loss_test.py b/tests/unit_tests/ohem_loss_test.py index 1eb4969f94..482c00c257 100644 --- a/tests/unit_tests/ohem_loss_test.py +++ b/tests/unit_tests/ohem_loss_test.py @@ -6,7 +6,6 @@ class OhemLossTest(unittest.TestCase): - def setUp(self) -> None: self.img_size = 64 self.eps = 0.01 @@ -18,7 +17,7 @@ def test_all_hard_no_mining(self): targets = torch.randint(0, num_classes, (1, self.img_size, self.img_size)) predictions = torch.ones((1, num_classes, self.img_size, self.img_size)) - probability = (1 / num_classes) + probability = 1 / num_classes # All samples are hard, No Hard-mining criterion = OhemCELoss(threshold=probability + self.eps, mining_percent=0.1) @@ -35,11 +34,11 @@ def test_hard_mining(self): # create hard samples hard_class = 0 mask = targets == hard_class - predictions[:, hard_class, mask.squeeze()] = 0. + predictions[:, hard_class, mask.squeeze()] = 0.0 hard_percent = mask.sum() / targets.numel() - predicted_prob = F.softmax(torch.tensor([0., 1.]), dim=0)[0].item() + predicted_prob = F.softmax(torch.tensor([0.0, 1.0]), dim=0)[0].item() criterion = OhemCELoss(threshold=predicted_prob + self.eps, mining_percent=hard_percent) expected_loss = -torch.log(torch.tensor(predicted_prob)) @@ -55,12 +54,12 @@ def test_ignore_label(self): # create hard samples, to be ignored later hard_class = 0 mask = targets == hard_class - predictions[:, hard_class, mask.squeeze()] = 0. + predictions[:, hard_class, mask.squeeze()] = 0.0 # except loss to be an equal distribution, w.r.t ignoring the hard label - predicted_prob = F.softmax(torch.tensor([1., 1.]), dim=0)[0].item() + predicted_prob = F.softmax(torch.tensor([1.0, 1.0]), dim=0)[0].item() - criterion = OhemCELoss(threshold=predicted_prob + self.eps, mining_percent=1., ignore_lb=hard_class) + criterion = OhemCELoss(threshold=predicted_prob + self.eps, mining_percent=1.0, ignore_lb=hard_class) expected_loss = -torch.log(torch.tensor(predicted_prob)) loss = criterion(predictions, targets) @@ -72,12 +71,12 @@ def test_all_are_ignore_label(self): targets = torch.zeros(1, self.img_size, self.img_size).long() # all targets are 0 class ignore_class = 0 - criterion = OhemCELoss(threshold=0.5, mining_percent=1., ignore_lb=ignore_class) - expected_loss = 0.0 # except empty zero tensor, because all are ignore labels + criterion = OhemCELoss(threshold=0.5, mining_percent=1.0, ignore_lb=ignore_class) + expected_loss = 0.0 # except empty zero tensor, because all are ignore labels loss = criterion(predictions, targets) self.assertAlmostEqual(expected_loss, loss, delta=1e-5) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/phase_context_test.py b/tests/unit_tests/phase_context_test.py index 15d227d35c..a9d37f7f6f 100644 --- a/tests/unit_tests/phase_context_test.py +++ b/tests/unit_tests/phase_context_test.py @@ -16,22 +16,33 @@ def context_information_in_train_test(self): net = ResNet18(num_classes=5, arch_params={}) - phase_callbacks = [PhaseContextTestCallback(Phase.TRAIN_BATCH_END), - PhaseContextTestCallback(Phase.TRAIN_BATCH_STEP), - PhaseContextTestCallback(Phase.TRAIN_EPOCH_END), - PhaseContextTestCallback(Phase.VALIDATION_BATCH_END), - PhaseContextTestCallback(Phase.VALIDATION_EPOCH_END)] + phase_callbacks = [ + PhaseContextTestCallback(Phase.TRAIN_BATCH_END), + PhaseContextTestCallback(Phase.TRAIN_BATCH_STEP), + PhaseContextTestCallback(Phase.TRAIN_EPOCH_END), + PhaseContextTestCallback(Phase.VALIDATION_BATCH_END), + PhaseContextTestCallback(Phase.VALIDATION_EPOCH_END), + ] - train_params = {"max_epochs": 2, "lr_updates": [1], "lr_decay_factor": 0.1, "lr_mode": "step", - "lr_warmup_epochs": 0, "initial_lr": 0.1, "loss": "cross_entropy", "optimizer": "SGD", - "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, - "train_metrics_list": [Accuracy()], "valid_metrics_list": [Top5()], - "metric_to_watch": "Top5", - "greater_metric_to_watch_is_better": True, "phase_callbacks": phase_callbacks} + train_params = { + "max_epochs": 2, + "lr_updates": [1], + "lr_decay_factor": 0.1, + "lr_mode": "step", + "lr_warmup_epochs": 0, + "initial_lr": 0.1, + "loss": "cross_entropy", + "optimizer": "SGD", + "criterion_params": {}, + "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, + "train_metrics_list": [Accuracy()], + "valid_metrics_list": [Top5()], + "metric_to_watch": "Top5", + "greater_metric_to_watch_is_better": True, + "phase_callbacks": phase_callbacks, + } - trainer.train(model=net, training_params=train_params, - train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + trainer.train(model=net, training_params=train_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader()) context_callbacks = list(filter(lambda cb: isinstance(cb, PhaseContextTestCallback), trainer.phase_callbacks)) # CHECK THAT PHASE CONTEXES HAVE THE EXACT INFORMATION THERY'RE SUPPOSE TO HOLD @@ -48,13 +59,16 @@ def context_information_in_train_test(self): if phase_callback.phase == Phase.VALIDATION_BATCH_END: self.assertTrue(phase_callback.context.epoch == 2) - self.assertTrue(isinstance(phase_callback.context.metrics_compute_fn, MetricCollection) and - hasattr(phase_callback.context.metrics_compute_fn, "Top5")) + self.assertTrue( + isinstance(phase_callback.context.metrics_compute_fn, MetricCollection) and hasattr(phase_callback.context.metrics_compute_fn, "Top5") + ) else: self.assertTrue(phase_callback.context.epoch == 1) - self.assertTrue(isinstance(phase_callback.context.metrics_compute_fn, MetricCollection) and - hasattr(phase_callback.context.metrics_compute_fn, "Accuracy")) + self.assertTrue( + isinstance(phase_callback.context.metrics_compute_fn, MetricCollection) + and hasattr(phase_callback.context.metrics_compute_fn, "Accuracy") + ) if phase_callback.phase in [Phase.TRAIN_EPOCH_END, Phase.VALIDATION_EPOCH_END]: self.assertTrue(phase_callback.context.batch_idx is None) @@ -73,5 +87,5 @@ def context_information_in_train_test(self): self.assertTrue("Top5" in phase_callback.context.metrics_dict.keys()) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/phase_delegates_test.py b/tests/unit_tests/phase_delegates_test.py index 6e607a64bb..53bdbea809 100644 --- a/tests/unit_tests/phase_delegates_test.py +++ b/tests/unit_tests/phase_delegates_test.py @@ -35,38 +35,59 @@ def test_access_to_methods_by_phase(self): phase_callbacks = [] for phase in Phase: - if phase in [Phase.PRE_TRAINING, Phase.TRAIN_EPOCH_START, Phase.TRAIN_EPOCH_END, Phase.VALIDATION_EPOCH_END, - Phase.VALIDATION_END_BEST_EPOCH, Phase.POST_TRAINING]: - phase_callbacks.append(ContextMethodsCheckerCallback(phase=phase, accessible_method_names=["get_net", - "set_net", - "set_ckpt_best_name", - "reset_best_metric", - "validate_epoch"], - non_accessible_method_names=[])) + if phase in [ + Phase.PRE_TRAINING, + Phase.TRAIN_EPOCH_START, + Phase.TRAIN_EPOCH_END, + Phase.VALIDATION_EPOCH_END, + Phase.VALIDATION_END_BEST_EPOCH, + Phase.POST_TRAINING, + ]: + phase_callbacks.append( + ContextMethodsCheckerCallback( + phase=phase, + accessible_method_names=["get_net", "set_net", "set_ckpt_best_name", "reset_best_metric", "validate_epoch"], + non_accessible_method_names=[], + ) + ) else: phase_callbacks.append( - ContextMethodsCheckerCallback(phase=phase, non_accessible_method_names=["get_net", - "set_net", - "set_ckpt_best_name", - "reset_best_metric", - "validate_epoch", - "set_ema"], - accessible_method_names=[])) + ContextMethodsCheckerCallback( + phase=phase, + non_accessible_method_names=["get_net", "set_net", "set_ckpt_best_name", "reset_best_metric", "validate_epoch", "set_ema"], + accessible_method_names=[], + ) + ) - train_params = {"max_epochs": 1, "lr_updates": [], "lr_decay_factor": 0.1, "lr_mode": "step", - "lr_warmup_epochs": 0, "initial_lr": 1, "loss": "cross_entropy", "optimizer": 'SGD', - "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, - "train_metrics_list": [Accuracy()], "valid_metrics_list": [Accuracy()], - "metric_to_watch": "Accuracy", - "greater_metric_to_watch_is_better": True, "ema": False, "phase_callbacks": phase_callbacks} + train_params = { + "max_epochs": 1, + "lr_updates": [], + "lr_decay_factor": 0.1, + "lr_mode": "step", + "lr_warmup_epochs": 0, + "initial_lr": 1, + "loss": "cross_entropy", + "optimizer": "SGD", + "criterion_params": {}, + "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, + "train_metrics_list": [Accuracy()], + "valid_metrics_list": [Accuracy()], + "metric_to_watch": "Accuracy", + "greater_metric_to_watch_is_better": True, + "ema": False, + "phase_callbacks": phase_callbacks, + } - trainer.train(model=net, training_params=train_params, - train_loader=classification_test_dataloader(batch_size=4), - valid_loader=classification_test_dataloader(batch_size=4)) + trainer.train( + model=net, + training_params=train_params, + train_loader=classification_test_dataloader(batch_size=4), + valid_loader=classification_test_dataloader(batch_size=4), + ) for phase_callback in phase_callbacks: if isinstance(phase_callback, ContextMethodsCheckerCallback): self.assertTrue(phase_callback.result) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/random_erase_test.py b/tests/unit_tests/random_erase_test.py index 1a36c7e3af..06fde6d722 100644 --- a/tests/unit_tests/random_erase_test.py +++ b/tests/unit_tests/random_erase_test.py @@ -6,15 +6,15 @@ class RandomEraseTest(unittest.TestCase): def test_random_erase(self): dummy_input = torch.randn(1, 3, 32, 32) - one_erase = RandomErase(probability=0, value='1.') + one_erase = RandomErase(probability=0, value="1.") self.assertEqual(one_erase.p, 0) - self.assertEqual(one_erase.value, 1.) + self.assertEqual(one_erase.value, 1.0) one_erase(dummy_input) - rndm_erase = RandomErase(probability=0, value='random') - self.assertEqual(rndm_erase.value, 'random') + rndm_erase = RandomErase(probability=0, value="random") + self.assertEqual(rndm_erase.value, "random") rndm_erase(dummy_input) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/registry_test.py b/tests/unit_tests/registry_test.py index 84ba2b5bcb..f68490d609 100644 --- a/tests/unit_tests/registry_test.py +++ b/tests/unit_tests/registry_test.py @@ -13,10 +13,8 @@ class RegistryTest(unittest.TestCase): - def setUp(self): - - @register_model('myconvnet') + @register_model("myconvnet") class MyConvNet(nn.Module): def __init__(self, num_classes): super().__init__() @@ -40,7 +38,7 @@ def forward(self, x): def myconvnet_for_cifar10(): return MyConvNet(num_classes=10) - @register_metric('custom_accuracy') # Will be registered as "custom_accuracy" + @register_metric("custom_accuracy") # Will be registered as "custom_accuracy" class CustomAccuracy(torchmetrics.Accuracy): def update(self, preds: torch.Tensor, target: torch.Tensor): if target.shape == preds.shape: @@ -54,30 +52,30 @@ def forward(self, output, target): return 1 - criterion_mse(output, target).item() / torch.var(target).item() def tearDown(self): - ARCHITECTURES.pop('myconvnet', None) - ARCHITECTURES.pop('myconvnet_for_cifar10', None) - METRICS.pop('custom_accuracy', None) - LOSSES.pop('custom_rsquared_loss', None) + ARCHITECTURES.pop("myconvnet", None) + ARCHITECTURES.pop("myconvnet_for_cifar10", None) + METRICS.pop("custom_accuracy", None) + LOSSES.pop("custom_rsquared_loss", None) def test_cls_is_registered(self): - assert ARCHITECTURES['myconvnet'] - assert METRICS['custom_accuracy'] - assert LOSSES['custom_rsquared_loss'] + assert ARCHITECTURES["myconvnet"] + assert METRICS["custom_accuracy"] + assert LOSSES["custom_rsquared_loss"] def test_fn_is_registered(self): - assert ARCHITECTURES['myconvnet_for_cifar10'] + assert ARCHITECTURES["myconvnet_for_cifar10"] def test_is_instantiable(self): - assert ARCHITECTURES['myconvnet_for_cifar10']() - assert ARCHITECTURES['myconvnet'](num_classes=10) - assert METRICS['custom_accuracy']() - assert LOSSES['custom_rsquared_loss']() + assert ARCHITECTURES["myconvnet_for_cifar10"]() + assert ARCHITECTURES["myconvnet"](num_classes=10) + assert METRICS["custom_accuracy"]() + assert LOSSES["custom_rsquared_loss"]() def test_model_outputs(self): torch.manual_seed(0) - model_1 = ARCHITECTURES['myconvnet_for_cifar10']() + model_1 = ARCHITECTURES["myconvnet_for_cifar10"]() torch.manual_seed(0) - model_2 = ARCHITECTURES['myconvnet'](num_classes=10) + model_2 = ARCHITECTURES["myconvnet"](num_classes=10) dummy_input = torch.randn(1, 3, 32, 32, requires_grad=False) x = model_1(dummy_input) y = model_2(dummy_input) @@ -85,10 +83,11 @@ def test_model_outputs(self): def test_existing_key(self): with self.assertRaises(Exception): + @register_model() def myconvnet_for_cifar10(): return -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/regnet_unit_test.py b/tests/unit_tests/regnet_unit_test.py index 3b8961f1bc..49c37669a7 100644 --- a/tests/unit_tests/regnet_unit_test.py +++ b/tests/unit_tests/regnet_unit_test.py @@ -2,15 +2,24 @@ import torch.nn as nn import unittest -from super_gradients.training.models.classification_models.regnet import CustomRegNet, NASRegNet, RegNetY200, RegNetY400, RegNetY600, RegNetY800, \ - Stem, Stage, XBlock +from super_gradients.training.models.classification_models.regnet import ( + CustomRegNet, + NASRegNet, + RegNetY200, + RegNetY400, + RegNetY600, + RegNetY800, + Stem, + Stage, + XBlock, +) from super_gradients.training.utils.utils import HpmStruct class TestRegnet(unittest.TestCase): @classmethod def setUp(cls): - cls.arch_params = HpmStruct(**{'num_classes': 1000}) + cls.arch_params = HpmStruct(**{"num_classes": 1000}) @staticmethod def verify_2_archs_are_identical(model_1: nn.Module, model_2: nn.Module): @@ -21,10 +30,17 @@ def test_custom_and_nas_regnet_can_build_regnetY200(self): """Test that when build Nas Regnet and Custom Regnet with the correct params - they build RegnetY200""" regnet_y_200 = RegNetY200(arch_params=self.arch_params) # Parameters identical to regnet_y_200 - nas_regnet = NASRegNet(arch_params=HpmStruct(**{'structure': [24, 36, 2.5, 13, 1, 8, 2, 4], - 'num_classes': 1000})) - regnet_y_200_arch_params = {'initial_width': 24, 'slope': 36, 'quantized_param': 2.5, 'network_depth': 13, - 'bottleneck_ratio': 1, 'group_width': 8, 'stride': 2, 'num_classes': 1000} + nas_regnet = NASRegNet(arch_params=HpmStruct(**{"structure": [24, 36, 2.5, 13, 1, 8, 2, 4], "num_classes": 1000})) + regnet_y_200_arch_params = { + "initial_width": 24, + "slope": 36, + "quantized_param": 2.5, + "network_depth": 13, + "bottleneck_ratio": 1, + "group_width": 8, + "stride": 2, + "num_classes": 1000, + } custom_regnet = CustomRegNet(arch_params=HpmStruct(**regnet_y_200_arch_params)) self.verify_2_archs_are_identical(regnet_y_200, nas_regnet) @@ -50,7 +66,7 @@ def test_dropout_forward_backward(self): """ Test that output is stochastic in training and is fixed in eval with Dropout. """ - arch_params = HpmStruct(**{'num_classes': 1000, 'dropout_prob': 0.3}) + arch_params = HpmStruct(**{"num_classes": 1000, "dropout_prob": 0.3}) model = RegNetY200(arch_params=arch_params) dummy_input = torch.randn(1, 3, 224, 224) @@ -64,7 +80,7 @@ def test_droppath_forward_backward(self): """ Test that output is stochastic in training and is fixed in eval with DropPath. """ - arch_params = HpmStruct(**{'num_classes': 1000, 'droppath_prob': 0.2}) + arch_params = HpmStruct(**{"num_classes": 1000, "droppath_prob": 0.2}) model = RegNetY200(arch_params=arch_params) dummy_input = torch.randn(1, 3, 224, 224) @@ -81,34 +97,79 @@ def test_nas_regnet_logic_is_backward_competible(self): """ # THE LIST CONSISTS SEVERAL CUSTOM REGNET "ENCODINGS" AND THE CORRESPONDING XBLOCK STRUCTURE OF THE MODEL selected_arch_and_corresponding_configs = [ - {"struct": [56, 10, 2.2, 8, 2, 8, 2, 0], - "expected_config": [3, 32, 32, - 32, 16, 16, 16, 2, (2, 2), None, - 32, 32, 32, 32, 4, (2, 2), None]}, - {"struct": [56, 10, 2.3, 11, 1, 1, 3, 0], - "expected_config": [3, 32, 32, - 32, 56, 56, 56, 56, (3, 3), None, - 56, 128, 128, 128, 128, (3, 3), None]}, - {"struct": [70, 20, 2.6, 13, 0.5, 16, 2, 4], - "expected_config": [3, 32, 32, - 32, 288, 288, 288, 18, (2, 2), nn.Module, - 144, 736, 736, 736, 46, (2, 2), nn.Module, - 368, 1888, 1888, 1888, 118, (2, 2), nn.Module]}, - {"struct": [8, 20, 2.3, 13, 0.16666666666666666, 1, 2, 2], - "expected_config": [3, 32, 32, - 32, 288, 288, 288, 288, (2, 2), nn.Module, - 48, 1440, 1440, 1440, 1440, (2, 2), nn.Module, - 240, 3456, 3456, 3456, 3456, (2, 2), nn.Module, - 576, 8064, 8064, 8064, 8064, (2, 2), nn.Module]}, - {"struct": [56, 10, 2.4, 13, 2, 8, 1, 0], - "expected_config": [3, 32, 32, - 32, 16, 16, 16, 2, (1, 1), None, - 32, 32, 32, 32, 4, (1, 1), None]}, + {"struct": [56, 10, 2.2, 8, 2, 8, 2, 0], "expected_config": [3, 32, 32, 32, 16, 16, 16, 2, (2, 2), None, 32, 32, 32, 32, 4, (2, 2), None]}, + {"struct": [56, 10, 2.3, 11, 1, 1, 3, 0], "expected_config": [3, 32, 32, 32, 56, 56, 56, 56, (3, 3), None, 56, 128, 128, 128, 128, (3, 3), None]}, + { + "struct": [70, 20, 2.6, 13, 0.5, 16, 2, 4], + "expected_config": [ + 3, + 32, + 32, + 32, + 288, + 288, + 288, + 18, + (2, 2), + nn.Module, + 144, + 736, + 736, + 736, + 46, + (2, 2), + nn.Module, + 368, + 1888, + 1888, + 1888, + 118, + (2, 2), + nn.Module, + ], + }, + { + "struct": [8, 20, 2.3, 13, 0.16666666666666666, 1, 2, 2], + "expected_config": [ + 3, + 32, + 32, + 32, + 288, + 288, + 288, + 288, + (2, 2), + nn.Module, + 48, + 1440, + 1440, + 1440, + 1440, + (2, 2), + nn.Module, + 240, + 3456, + 3456, + 3456, + 3456, + (2, 2), + nn.Module, + 576, + 8064, + 8064, + 8064, + 8064, + (2, 2), + nn.Module, + ], + }, + {"struct": [56, 10, 2.4, 13, 2, 8, 1, 0], "expected_config": [3, 32, 32, 32, 16, 16, 16, 2, (1, 1), None, 32, 32, 32, 32, 4, (1, 1), None]}, ] for arch_conf_pair in selected_arch_and_corresponding_configs: - expected_config = iter(arch_conf_pair['expected_config']) - model = NASRegNet(HpmStruct(**{'structure': arch_conf_pair['struct'], 'num_classes': 1000})) + expected_config = iter(arch_conf_pair["expected_config"]) + model = NASRegNet(HpmStruct(**{"structure": arch_conf_pair["struct"], "num_classes": 1000})) for stage in model.net.children(): # CHECK CORRECTNESS OF THE STEM @@ -132,5 +193,5 @@ def test_nas_regnet_logic_is_backward_competible(self): break -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/save_ckpt_test.py b/tests/unit_tests/save_ckpt_test.py index 8b64dfa1b0..1605bab2a1 100644 --- a/tests/unit_tests/save_ckpt_test.py +++ b/tests/unit_tests/save_ckpt_test.py @@ -8,14 +8,24 @@ class SaveCkptListUnitTest(unittest.TestCase): def setUp(self): # Define Parameters - train_params = {"max_epochs": 4, "lr_decay_factor": 0.1, "lr_updates": [4], "lr_mode": "step", - "lr_warmup_epochs": 0, "initial_lr": 0.1, "loss": "cross_entropy", "optimizer": "SGD", - "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, - "save_ckpt_epoch_list": [1, 3], - "loss": "cross_entropy", "train_metrics_list": [Accuracy(), Top5()], - "valid_metrics_list": [Accuracy(), Top5()], - "metric_to_watch": "Accuracy", - "greater_metric_to_watch_is_better": True} + train_params = { + "max_epochs": 4, + "lr_decay_factor": 0.1, + "lr_updates": [4], + "lr_mode": "step", + "lr_warmup_epochs": 0, + "initial_lr": 0.1, + "loss": "cross_entropy", + "optimizer": "SGD", + "criterion_params": {}, + "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, + "save_ckpt_epoch_list": [1, 3], + "loss": "cross_entropy", + "train_metrics_list": [Accuracy(), Top5()], + "valid_metrics_list": [Accuracy(), Top5()], + "metric_to_watch": "Accuracy", + "greater_metric_to_watch_is_better": True, + } # Define Model trainer = Trainer("save_ckpt_test") @@ -24,17 +34,15 @@ def setUp(self): model = models.get("resnet18_cifar", arch_params={"num_classes": 10}) # Train Model (and save ckpt_epoch_list) - 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()) dir_path = trainer.checkpoints_dir_path - self.file_names_list = [dir_path + f'/ckpt_epoch_{epoch}.pth' for epoch in train_params["save_ckpt_epoch_list"]] + self.file_names_list = [dir_path + f"/ckpt_epoch_{epoch}.pth" for epoch in train_params["save_ckpt_epoch_list"]] def test_save_ckpt_epoch_list(self): self.assertTrue(os.path.exists(self.file_names_list[0])) self.assertTrue(os.path.exists(self.file_names_list[1])) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/segmentation_transforms_test.py b/tests/unit_tests/segmentation_transforms_test.py index 5bc6d91ad9..2cb0406d67 100644 --- a/tests/unit_tests/segmentation_transforms_test.py +++ b/tests/unit_tests/segmentation_transforms_test.py @@ -8,7 +8,6 @@ class SegmentationTransformsTest(unittest.TestCase): - def setUp(self) -> None: self.default_image_value = 0 self.default_mask_value = 0 @@ -16,7 +15,7 @@ def setUp(self) -> None: def create_sample(self, size): sample = { "image": Image.new(mode="RGB", size=size, color=self.default_image_value), - "mask": Image.new(mode="L", size=size, color=self.default_mask_value) + "mask": Image.new(mode="L", size=size, color=self.default_mask_value), } return sample @@ -151,8 +150,8 @@ def test_padding_fill_values(self): out_image = image_to_tensor(out["image"]) # test transformed mask values - original_values = out_mask[128 // 2:-128 // 2].unique().tolist() - pad_values = torch.cat([out_mask[:128 // 2], out_mask[-128 // 2:]], dim=0).unique().tolist() + original_values = out_mask[128 // 2 : -128 // 2].unique().tolist() + pad_values = torch.cat([out_mask[: 128 // 2], out_mask[-128 // 2 :]], dim=0).unique().tolist() self.assertEqual(len(original_values), 1) self.assertEqual(original_values[0], self.default_mask_value) @@ -161,8 +160,8 @@ def test_padding_fill_values(self): self.assertEqual(pad_values[0], fill_mask_value) # test transformed image values - original_values = out_image[:, 128 // 2:-128 // 2].unique().tolist() - pad_values = torch.cat([out_image[:, :128 // 2], out_image[:, -128 // 2:]], dim=1).unique().tolist() + original_values = out_image[:, 128 // 2 : -128 // 2].unique().tolist() + pad_values = torch.cat([out_image[:, : 128 // 2], out_image[:, -128 // 2 :]], dim=1).unique().tolist() self.assertEqual(len(original_values), 1) self.assertEqual(original_values[0], self.default_image_value) @@ -201,10 +200,7 @@ def test_rescale_padding(self): out_size = (512, 512) sample = self.create_sample(in_size) - transform = Compose([ - SegRescale(long_size=out_size[0]), # rescale to (512, 256) - SegPadShortToCropSize(crop_size=out_size) # pad to (512, 512) - ]) + transform = Compose([SegRescale(long_size=out_size[0]), SegPadShortToCropSize(crop_size=out_size)]) # rescale to (512, 256) # pad to (512, 512) out = transform(sample) self.assertEqual(out_size, out["image"].size) @@ -213,15 +209,13 @@ def test_random_rescale_padding_random_crop(self): crop_size = (256, 128) sample = self.create_sample(img_size) - transform = Compose([ - SegRandomRescale(scales=(0.1, 2.0)), - SegPadShortToCropSize(crop_size=crop_size), - SegCropImageAndMask(crop_size=crop_size, mode="random") - ]) + transform = Compose( + [SegRandomRescale(scales=(0.1, 2.0)), SegPadShortToCropSize(crop_size=crop_size), SegCropImageAndMask(crop_size=crop_size, mode="random")] + ) out = transform(sample) self.assertEqual(crop_size, out["image"].size) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/shelfnet_unit_test.py b/tests/unit_tests/shelfnet_unit_test.py index fde64558d3..2396183014 100644 --- a/tests/unit_tests/shelfnet_unit_test.py +++ b/tests/unit_tests/shelfnet_unit_test.py @@ -35,5 +35,5 @@ def test_shelfnet_creation(self): self.assertIsNotNone(output) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/test_auto_augment.py b/tests/unit_tests/test_auto_augment.py index 325f2a5927..d6cfa03362 100644 --- a/tests/unit_tests/test_auto_augment.py +++ b/tests/unit_tests/test_auto_augment.py @@ -7,11 +7,8 @@ class TestAutoAugment(unittest.TestCase): - def setUp(self): - self.dataset_params = {"batch_size": 1, - "color_jitter": 0.1, - 'rand_augment_config_string': "m9-mstd0.5"} + self.dataset_params = {"batch_size": 1, "color_jitter": 0.1, "rand_augment_config_string": "m9-mstd0.5"} def test_autoaugment_call(self): """ @@ -20,16 +17,16 @@ def test_autoaugment_call(self): image_size = 224 color_augmentation = get_color_augmentation("m9-mstd0.5", color_jitter=None, crop_size=image_size) self.assertTrue(isinstance(color_augmentation, RandAugment)) - img = Image.fromarray(np.ones((image_size, image_size, 3)).astype('uint8')) + img = Image.fromarray(np.ones((image_size, image_size, 3)).astype("uint8")) augmented_image = color_augmentation(img) self.assertTrue(augmented_image.size == (image_size, image_size)) color_augmentation = get_color_augmentation(None, color_jitter=(0.7, 0.7, 0.7), crop_size=image_size) self.assertTrue(isinstance(color_augmentation, transforms.ColorJitter)) - img = Image.fromarray(np.random.randn(image_size, image_size, 3).astype('uint8')) + img = Image.fromarray(np.random.randn(image_size, image_size, 3).astype("uint8")) augmented_image = color_augmentation(img) self.assertTrue(augmented_image.size == (image_size, image_size)) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/train_with_intialized_param_args_test.py b/tests/unit_tests/train_with_intialized_param_args_test.py index 2527e20a92..0a21404cf1 100644 --- a/tests/unit_tests/train_with_intialized_param_args_test.py +++ b/tests/unit_tests/train_with_intialized_param_args_test.py @@ -23,13 +23,22 @@ def test_train_with_external_criterion(self): dataloader = classification_test_dataloader(batch_size=10) model = models.get("resnet18", arch_params={"num_classes": 5}) - train_params = {"max_epochs": 2, "lr_updates": [1], "lr_decay_factor": 0.1, "lr_mode": "step", - "lr_warmup_epochs": 0, "initial_lr": 0.1, "loss": torch.nn.CrossEntropyLoss(), - "optimizer": "SGD", - "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, - "train_metrics_list": [Accuracy()], "valid_metrics_list": [Accuracy()], - "metric_to_watch": "Accuracy", - "greater_metric_to_watch_is_better": True} + train_params = { + "max_epochs": 2, + "lr_updates": [1], + "lr_decay_factor": 0.1, + "lr_mode": "step", + "lr_warmup_epochs": 0, + "initial_lr": 0.1, + "loss": torch.nn.CrossEntropyLoss(), + "optimizer": "SGD", + "criterion_params": {}, + "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, + "train_metrics_list": [Accuracy()], + "valid_metrics_list": [Accuracy()], + "metric_to_watch": "Accuracy", + "greater_metric_to_watch_is_better": True, + } trainer.train(model=model, training_params=train_params, train_loader=dataloader, valid_loader=dataloader) def test_train_with_external_optimizer(self): @@ -38,12 +47,22 @@ def test_train_with_external_optimizer(self): model = models.get("resnet18", arch_params={"num_classes": 5}) optimizer = SGD(params=model.parameters(), lr=0.1) - train_params = {"max_epochs": 2, "lr_updates": [1], "lr_decay_factor": 0.1, "lr_mode": "step", - "lr_warmup_epochs": 0, "initial_lr": 0.1, "loss": "cross_entropy", "optimizer": optimizer, - "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, - "train_metrics_list": [Accuracy(), Top5()], "valid_metrics_list": [Accuracy(), Top5()], - "metric_to_watch": "Accuracy", - "greater_metric_to_watch_is_better": True} + train_params = { + "max_epochs": 2, + "lr_updates": [1], + "lr_decay_factor": 0.1, + "lr_mode": "step", + "lr_warmup_epochs": 0, + "initial_lr": 0.1, + "loss": "cross_entropy", + "optimizer": optimizer, + "criterion_params": {}, + "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, + "train_metrics_list": [Accuracy(), Top5()], + "valid_metrics_list": [Accuracy(), Top5()], + "metric_to_watch": "Accuracy", + "greater_metric_to_watch_is_better": True, + } trainer.train(model=model, training_params=train_params, train_loader=dataloader, valid_loader=dataloader) def test_train_with_external_scheduler(self): @@ -56,12 +75,19 @@ def test_train_with_external_scheduler(self): lr_scheduler = MultiStepLR(optimizer=optimizer, milestones=[1, 2], gamma=0.1) phase_callbacks = [LRSchedulerCallback(lr_scheduler, Phase.TRAIN_EPOCH_END)] - train_params = {"max_epochs": 2, "phase_callbacks": phase_callbacks, - "lr_warmup_epochs": 0, "initial_lr": lr, "loss": "cross_entropy", "optimizer": optimizer, - "criterion_params": {}, - "train_metrics_list": [Accuracy(), Top5()], "valid_metrics_list": [Accuracy(), Top5()], - "metric_to_watch": "Accuracy", - "greater_metric_to_watch_is_better": True} + train_params = { + "max_epochs": 2, + "phase_callbacks": phase_callbacks, + "lr_warmup_epochs": 0, + "initial_lr": lr, + "loss": "cross_entropy", + "optimizer": optimizer, + "criterion_params": {}, + "train_metrics_list": [Accuracy(), Top5()], + "valid_metrics_list": [Accuracy(), Top5()], + "metric_to_watch": "Accuracy", + "greater_metric_to_watch_is_better": True, + } trainer.train(model=model, training_params=train_params, train_loader=dataloader, valid_loader=dataloader) self.assertTrue(lr_scheduler.get_last_lr()[0] == lr * 0.1 * 0.1) @@ -72,12 +98,18 @@ def test_train_with_external_scheduler_class(self): model = models.get("resnet18", arch_params={"num_classes": 5}) optimizer = SGD # a class - not an instance - train_params = {"max_epochs": 2, - "lr_warmup_epochs": 0, "initial_lr": 0.3, "loss": "cross_entropy", "optimizer": optimizer, - "criterion_params": {}, - "train_metrics_list": [Accuracy(), Top5()], "valid_metrics_list": [Accuracy(), Top5()], - "metric_to_watch": "Accuracy", - "greater_metric_to_watch_is_better": True} + train_params = { + "max_epochs": 2, + "lr_warmup_epochs": 0, + "initial_lr": 0.3, + "loss": "cross_entropy", + "optimizer": optimizer, + "criterion_params": {}, + "train_metrics_list": [Accuracy(), Top5()], + "valid_metrics_list": [Accuracy(), Top5()], + "metric_to_watch": "Accuracy", + "greater_metric_to_watch_is_better": True, + } trainer.train(model=model, training_params=train_params, train_loader=dataloader, valid_loader=dataloader) def test_train_with_reduce_on_plateau(self): @@ -90,13 +122,19 @@ def test_train_with_reduce_on_plateau(self): lr_scheduler = ReduceLROnPlateau(optimizer=optimizer, patience=0) phase_callbacks = [LRSchedulerCallback(lr_scheduler, Phase.VALIDATION_EPOCH_END, "ToyTestClassificationMetric")] - train_params = {"max_epochs": 2, "phase_callbacks": phase_callbacks, - "lr_warmup_epochs": 0, "initial_lr": lr, "loss": "cross_entropy", "optimizer": optimizer, - "criterion_params": {}, - "train_metrics_list": [Accuracy(), Top5()], - "valid_metrics_list": [Accuracy(), Top5(), ToyTestClassificationMetric()], - "metric_to_watch": "Accuracy", - "greater_metric_to_watch_is_better": True} + train_params = { + "max_epochs": 2, + "phase_callbacks": phase_callbacks, + "lr_warmup_epochs": 0, + "initial_lr": lr, + "loss": "cross_entropy", + "optimizer": optimizer, + "criterion_params": {}, + "train_metrics_list": [Accuracy(), Top5()], + "valid_metrics_list": [Accuracy(), Top5(), ToyTestClassificationMetric()], + "metric_to_watch": "Accuracy", + "greater_metric_to_watch_is_better": True, + } trainer.train(model=model, training_params=train_params, train_loader=dataloader, valid_loader=dataloader) self.assertTrue(lr_scheduler._last_lr[0] == lr * 0.1) @@ -105,36 +143,54 @@ def test_train_with_external_metric(self): dataloader = classification_test_dataloader(batch_size=10) model = models.get("resnet18", arch_params={"num_classes": 5}) - train_params = {"max_epochs": 2, "lr_updates": [1], "lr_decay_factor": 0.1, "lr_mode": "step", - "lr_warmup_epochs": 0, "initial_lr": 0.1, "loss": "cross_entropy", "optimizer": "SGD", - "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, - "train_metrics_list": [F1Score()], "valid_metrics_list": [F1Score()], - "metric_to_watch": "F1Score", - "greater_metric_to_watch_is_better": True} + train_params = { + "max_epochs": 2, + "lr_updates": [1], + "lr_decay_factor": 0.1, + "lr_mode": "step", + "lr_warmup_epochs": 0, + "initial_lr": 0.1, + "loss": "cross_entropy", + "optimizer": "SGD", + "criterion_params": {}, + "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, + "train_metrics_list": [F1Score()], + "valid_metrics_list": [F1Score()], + "metric_to_watch": "F1Score", + "greater_metric_to_watch_is_better": True, + } trainer.train(model=model, training_params=train_params, train_loader=dataloader, valid_loader=dataloader) def test_train_with_external_dataloaders(self): trainer = Trainer("external_data_loader_test") batch_size = 5 - trainset = torch.utils.data.TensorDataset(torch.Tensor(np.random.random((10, 3, 32, 32))), - torch.LongTensor(np.zeros((10)))) + trainset = torch.utils.data.TensorDataset(torch.Tensor(np.random.random((10, 3, 32, 32))), torch.LongTensor(np.zeros((10)))) - valset = torch.utils.data.TensorDataset(torch.Tensor(np.random.random((10, 3, 32, 32))), - torch.LongTensor(np.zeros((10)))) + valset = torch.utils.data.TensorDataset(torch.Tensor(np.random.random((10, 3, 32, 32))), torch.LongTensor(np.zeros((10)))) train_loader = torch.utils.data.DataLoader(trainset, batch_size=batch_size) val_loader = torch.utils.data.DataLoader(valset, batch_size=batch_size) model = models.get("resnet18", num_classes=5) - train_params = {"max_epochs": 2, "lr_updates": [1], "lr_decay_factor": 0.1, "lr_mode": "step", - "lr_warmup_epochs": 0, "initial_lr": 0.1, "loss": "cross_entropy", "optimizer": "SGD", - "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, - "train_metrics_list": [F1Score()], "valid_metrics_list": [F1Score()], - "metric_to_watch": "F1Score", - "greater_metric_to_watch_is_better": True} + train_params = { + "max_epochs": 2, + "lr_updates": [1], + "lr_decay_factor": 0.1, + "lr_mode": "step", + "lr_warmup_epochs": 0, + "initial_lr": 0.1, + "loss": "cross_entropy", + "optimizer": "SGD", + "criterion_params": {}, + "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, + "train_metrics_list": [F1Score()], + "valid_metrics_list": [F1Score()], + "metric_to_watch": "F1Score", + "greater_metric_to_watch_is_better": True, + } trainer.train(model=model, training_params=train_params, train_loader=train_loader, valid_loader=val_loader) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/train_with_precise_bn_test.py b/tests/unit_tests/train_with_precise_bn_test.py index 98408f1775..7a2eff2a99 100644 --- a/tests/unit_tests/train_with_precise_bn_test.py +++ b/tests/unit_tests/train_with_precise_bn_test.py @@ -14,32 +14,59 @@ class TrainWithPreciseBNTest(unittest.TestCase): def test_train_with_precise_bn_explicit_size(self): trainer = Trainer("test_train_with_precise_bn_explicit_size") net = ResNet18(num_classes=5, arch_params={}) - train_params = {"max_epochs": 2, "lr_updates": [1], "lr_decay_factor": 0.1, "lr_mode": "step", - "lr_warmup_epochs": 0, "initial_lr": 0.1, "loss": "cross_entropy", "optimizer": "SGD", - "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, - "train_metrics_list": [Accuracy(), Top5()], "valid_metrics_list": [Accuracy(), Top5()], - "metric_to_watch": "Accuracy", - "greater_metric_to_watch_is_better": True, - "precise_bn": True, "precise_bn_batch_size": 100} - trainer.train(model=net, training_params=train_params, - train_loader=classification_test_dataloader(batch_size=10), - valid_loader=classification_test_dataloader(batch_size=10)) + train_params = { + "max_epochs": 2, + "lr_updates": [1], + "lr_decay_factor": 0.1, + "lr_mode": "step", + "lr_warmup_epochs": 0, + "initial_lr": 0.1, + "loss": "cross_entropy", + "optimizer": "SGD", + "criterion_params": {}, + "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, + "train_metrics_list": [Accuracy(), Top5()], + "valid_metrics_list": [Accuracy(), Top5()], + "metric_to_watch": "Accuracy", + "greater_metric_to_watch_is_better": True, + "precise_bn": True, + "precise_bn_batch_size": 100, + } + trainer.train( + model=net, + training_params=train_params, + train_loader=classification_test_dataloader(batch_size=10), + valid_loader=classification_test_dataloader(batch_size=10), + ) def test_train_with_precise_bn_implicit_size(self): trainer = Trainer("test_train_with_precise_bn_implicit_size") net = ResNet18(num_classes=5, arch_params={}) - train_params = {"max_epochs": 2, "lr_updates": [1], "lr_decay_factor": 0.1, "lr_mode": "step", - "lr_warmup_epochs": 0, "initial_lr": 0.1, "loss": "cross_entropy", "optimizer": "SGD", - "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, - "train_metrics_list": [Accuracy(), Top5()], "valid_metrics_list": [Accuracy(), Top5()], - "metric_to_watch": "Accuracy", - "greater_metric_to_watch_is_better": True, - "precise_bn": True} - trainer.train(model=net, training_params=train_params, - train_loader=classification_test_dataloader(batch_size=10), - valid_loader=classification_test_dataloader(batch_size=10)) - - -if __name__ == '__main__': + train_params = { + "max_epochs": 2, + "lr_updates": [1], + "lr_decay_factor": 0.1, + "lr_mode": "step", + "lr_warmup_epochs": 0, + "initial_lr": 0.1, + "loss": "cross_entropy", + "optimizer": "SGD", + "criterion_params": {}, + "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, + "train_metrics_list": [Accuracy(), Top5()], + "valid_metrics_list": [Accuracy(), Top5()], + "metric_to_watch": "Accuracy", + "greater_metric_to_watch_is_better": True, + "precise_bn": True, + } + trainer.train( + model=net, + training_params=train_params, + train_loader=classification_test_dataloader(batch_size=10), + valid_loader=classification_test_dataloader(batch_size=10), + ) + + +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/training_params_factory_test.py b/tests/unit_tests/training_params_factory_test.py index b4cb34acbf..b574cce8a2 100644 --- a/tests/unit_tests/training_params_factory_test.py +++ b/tests/unit_tests/training_params_factory_test.py @@ -14,5 +14,5 @@ def test_get_train_params_with_overrides(self): self.assertTrue(train_params["max_epochs"] == 5) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/update_param_groups_unit_test.py b/tests/unit_tests/update_param_groups_unit_test.py index ba9a9737b7..f0c85c71a2 100644 --- a/tests/unit_tests/update_param_groups_unit_test.py +++ b/tests/unit_tests/update_param_groups_unit_test.py @@ -16,9 +16,7 @@ class TestNet(LeNet): def __init__(self): super(TestNet, self).__init__() - def update_param_groups( - self, param_groups: list, lr: float, epoch: int, iter: int, training_params: HpmStruct, total_batch: int - ) -> list: + def update_param_groups(self, param_groups: list, lr: float, epoch: int, iter: int, training_params: HpmStruct, total_batch: int) -> list: initial_lr = get_param(training_params, "initial_lr") for param_group in param_groups: param_group["lr"] = initial_lr * (epoch + 1) @@ -34,21 +32,25 @@ def test_lr_scheduling_with_update_param_groups(self): lrs = [] phase_callbacks = [TestLRCallback(lr_placeholder=lrs)] - train_params = {"max_epochs": 3, - "lr_mode": "step", - "lr_updates": [0, 1, 2], - "initial_lr": 0.1, - "lr_decay_factor": 1, - "loss": "cross_entropy", "optimizer": 'SGD', - "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, - "train_metrics_list": [Accuracy()], "valid_metrics_list": [Accuracy()], - "metric_to_watch": "Accuracy", - "greater_metric_to_watch_is_better": True, "ema": False, "phase_callbacks": phase_callbacks, - } + train_params = { + "max_epochs": 3, + "lr_mode": "step", + "lr_updates": [0, 1, 2], + "initial_lr": 0.1, + "lr_decay_factor": 1, + "loss": "cross_entropy", + "optimizer": "SGD", + "criterion_params": {}, + "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9}, + "train_metrics_list": [Accuracy()], + "valid_metrics_list": [Accuracy()], + "metric_to_watch": "Accuracy", + "greater_metric_to_watch_is_better": True, + "ema": False, + "phase_callbacks": phase_callbacks, + } expected_lrs = np.array([0.1, 0.2, 0.3]) - trainer.train(model=net, training_params=train_params, - train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + trainer.train(model=net, training_params=train_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader()) self.assertTrue(np.allclose(np.array(lrs), expected_lrs, rtol=0.0000001)) diff --git a/tests/unit_tests/yolox_unit_test.py b/tests/unit_tests/yolox_unit_test.py index 4cddf260ea..215d4fd90e 100644 --- a/tests/unit_tests/yolox_unit_test.py +++ b/tests/unit_tests/yolox_unit_test.py @@ -40,5 +40,5 @@ def test_yolox_creation(self): self.assertIsNotNone(output_augment) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/zero_weight_decay_on_bias_bn_test.py b/tests/unit_tests/zero_weight_decay_on_bias_bn_test.py index 8869219480..75878e5f44 100644 --- a/tests/unit_tests/zero_weight_decay_on_bias_bn_test.py +++ b/tests/unit_tests/zero_weight_decay_on_bias_bn_test.py @@ -13,13 +13,14 @@ class ToyLinearKernel(nn.Module): Custom Toy linear module to test custom modules with bias parameter, that are not instances of primitive torch modules. """ + def __init__(self, in_features: int, out_features: int, bias: bool = True): super().__init__() self.weight = nn.Parameter(torch.Tensor(out_features, in_features)) if bias: self.bias = nn.Parameter(torch.Tensor(out_features)) else: - self.register_parameter('bias', None) + self.register_parameter("bias", None) def forward(self, input: torch.Tensor): return F.linear(input, self.weight, self.bias) @@ -29,12 +30,12 @@ class ToySgModule(SgModule): """ Toy Module to test zero of weight decay, support multiple group of parameters. """ + CONV_CLASSES = {1: nn.Conv1d, 2: nn.Conv2d, 3: nn.Conv3d} CONV_TRANSPOSE_CLASSES = {1: nn.ConvTranspose1d, 2: nn.ConvTranspose2d, 3: nn.ConvTranspose3d} BN_CLASSES = {1: nn.BatchNorm1d, 2: nn.BatchNorm2d, 3: nn.BatchNorm3d} - def __init__(self, input_dimension=2, multiple_param_groups=False, module_groups=False, - linear_cls=nn.Linear): + def __init__(self, input_dimension=2, multiple_param_groups=False, module_groups=False, linear_cls=nn.Linear): """ :param input_dimension: input dimension, 1 for 1D, 2 for 2D ... :param multiple_param_groups: if True create multiple param groups with different optimizer args. @@ -70,16 +71,9 @@ def __init__(self, input_dimension=2, multiple_param_groups=False, module_groups self.conv_transpose(128, 128), ) self.avg_pool = nn.AdaptiveAvgPool2d(1) - self.classifier = nn.Sequential( - self.linear(128, 2 * num_classes, bias=False), - self.linear(2 * num_classes, num_classes, bias=True) - ) + self.classifier = nn.Sequential(self.linear(128, 2 * num_classes, bias=False), self.linear(2 * num_classes, num_classes, bias=True)) - self.head = nn.Sequential( - self.more_convs, - self.avg_pool, - self.classifier - ) + self.head = nn.Sequential(self.more_convs, self.avg_pool, self.classifier) self.head_params = (self.num_no_decay_params() - self.base_params[0], self.num_decay_params() - self.base_params[1]) @@ -89,11 +83,7 @@ def conv1(self, ch_in: int, ch_out: int, stride: int, bias=False): conv = self.conv_cls(ch_in, ch_out, 1, stride=stride, bias=bias) self.num_biases += 1 else: - conv = nn.Sequential( - self.conv_cls(ch_in, ch_out, 1, stride=stride, bias=bias), - self.bn_cls(ch_out), - nn.ReLU() - ) + conv = nn.Sequential(self.conv_cls(ch_in, ch_out, 1, stride=stride, bias=bias), self.bn_cls(ch_out), nn.ReLU()) self.num_bn += 1 return conv @@ -103,11 +93,7 @@ def conv_transpose(self, ch_in: int, ch_out: int, bias=False): conv = self.conv_transpose_cls(ch_in, ch_out, 2, stride=2, bias=bias) self.num_biases += 1 else: - conv = nn.Sequential( - self.conv_transpose_cls(ch_in, ch_out, 2, stride=2, bias=bias), - self.bn_cls(ch_out), - nn.ReLU() - ) + conv = nn.Sequential(self.conv_transpose_cls(ch_in, ch_out, 2, stride=2, bias=bias), self.bn_cls(ch_out), nn.ReLU()) self.num_bn += 1 return conv @@ -126,8 +112,7 @@ def num_no_decay_params(self): def initialize_param_groups(self, lr: float, training_params: HpmStruct) -> list: # Example to different learning rates, similar to ShelfNet, in order to create multiple groups. if self.multiple_param_groups: - params_list = [{'named_params': self.base.named_parameters(), 'lr': lr}, - {'named_params': self.head.named_parameters(), 'lr': lr * 10}] + params_list = [{"named_params": self.base.named_parameters(), "lr": lr}, {"named_params": self.head.named_parameters(), "lr": lr * 10}] return params_list @@ -145,30 +130,29 @@ def setUp(self): self.lr = 0.1 # input dimension beside batch and channels, i.e 2 for vision, 1 for audio, 3 for point cloud. self.input_dimensions = (1, 2, 3) - self.train_params_zero_wd = {"initial_lr": self.lr, - "optimizer": "SGD", - "optimizer_params": {"weight_decay": self.weight_decay, "momentum": 0.9}} - - def _assert_optimizer_param_groups(self, - param_groups: list, - excpected_num_groups: int, - excpected_num_params_per_group: list, - excpected_weight_decay_per_group: list): + self.train_params_zero_wd = {"initial_lr": self.lr, "optimizer": "SGD", "optimizer_params": {"weight_decay": self.weight_decay, "momentum": 0.9}} + + def _assert_optimizer_param_groups( + self, param_groups: list, excpected_num_groups: int, excpected_num_params_per_group: list, excpected_weight_decay_per_group: list + ): """ Helper method to assert, num of param_groups, num of parameters in each param group and weight decay value in each param group. """ - self.assertEqual(len(param_groups), excpected_num_groups, - msg=f"Optimizer should have {excpected_num_groups} groups") - for (param_group, excpected_num_params, excpected_weight_decay) in zip(param_groups, - excpected_num_params_per_group, - excpected_weight_decay_per_group): - self.assertEqual(len(param_group["params"]), excpected_num_params, - msg="Wrong number of params for optimizer param group, excpected: {}, found: {}".format( - excpected_num_params, len(param_group["params"]))) - self.assertEqual(param_group['weight_decay'], excpected_weight_decay, - msg="Wrong weight decay value found for optimizer param group, excpected: {}, found: {}".format( - excpected_weight_decay, param_group["weight_decay"])) + self.assertEqual(len(param_groups), excpected_num_groups, msg=f"Optimizer should have {excpected_num_groups} groups") + for (param_group, excpected_num_params, excpected_weight_decay) in zip(param_groups, excpected_num_params_per_group, excpected_weight_decay_per_group): + self.assertEqual( + len(param_group["params"]), + excpected_num_params, + msg="Wrong number of params for optimizer param group, excpected: {}, found: {}".format(excpected_num_params, len(param_group["params"])), + ) + self.assertEqual( + param_group["weight_decay"], + excpected_weight_decay, + msg="Wrong weight decay value found for optimizer param group, excpected: {}, found: {}".format( + excpected_weight_decay, param_group["weight_decay"] + ), + ) def test_zero_wd_one_group(self): """ @@ -178,17 +162,13 @@ def test_zero_wd_one_group(self): net = ToySgModule(input_dimension=input_dim) train_params = HpmStruct(**self.train_params_zero_wd) - optimizer_params_groups = separate_zero_wd_params_groups_for_optimizer( - net, - net.initialize_param_groups(self.lr, train_params), - self.weight_decay - ) + optimizer_params_groups = separate_zero_wd_params_groups_for_optimizer(net, net.initialize_param_groups(self.lr, train_params), self.weight_decay) self._assert_optimizer_param_groups( optimizer_params_groups, excpected_num_groups=2, excpected_num_params_per_group=[net.num_no_decay_params(), net.num_decay_params()], - excpected_weight_decay_per_group=[0, self.weight_decay] + excpected_weight_decay_per_group=[0, self.weight_decay], ) def test_zero_wd_multiple_group(self): @@ -199,18 +179,13 @@ def test_zero_wd_multiple_group(self): net = ToySgModule(input_dimension=input_dim, multiple_param_groups=True) train_params = HpmStruct(**self.train_params_zero_wd) - optimizer_params_groups = separate_zero_wd_params_groups_for_optimizer( - net, - net.initialize_param_groups(self.lr, train_params), - self.weight_decay - ) + optimizer_params_groups = separate_zero_wd_params_groups_for_optimizer(net, net.initialize_param_groups(self.lr, train_params), self.weight_decay) self._assert_optimizer_param_groups( optimizer_params_groups, excpected_num_groups=4, - excpected_num_params_per_group=[net.base_params[0], net.base_params[1], net.head_params[0], - net.head_params[1]], - excpected_weight_decay_per_group=[0, self.weight_decay, 0, self.weight_decay] + excpected_num_params_per_group=[net.base_params[0], net.base_params[1], net.head_params[0], net.head_params[1]], + excpected_weight_decay_per_group=[0, self.weight_decay, 0, self.weight_decay], ) def test_zero_wd_sync_bn(self): @@ -224,17 +199,13 @@ def test_zero_wd_sync_bn(self): train_params = HpmStruct(**self.train_params_zero_wd) - optimizer_params_groups = separate_zero_wd_params_groups_for_optimizer( - net, - net.initialize_param_groups(self.lr, train_params), - self.weight_decay - ) + optimizer_params_groups = separate_zero_wd_params_groups_for_optimizer(net, net.initialize_param_groups(self.lr, train_params), self.weight_decay) self._assert_optimizer_param_groups( optimizer_params_groups, excpected_num_groups=2, excpected_num_params_per_group=[net.num_no_decay_params(), net.num_decay_params()], - excpected_weight_decay_per_group=[0, self.weight_decay] + excpected_weight_decay_per_group=[0, self.weight_decay], ) def test_zero_wd_custom_module_with_bias(self): @@ -246,19 +217,15 @@ def test_zero_wd_custom_module_with_bias(self): train_params = HpmStruct(**self.train_params_zero_wd) - optimizer_params_groups = separate_zero_wd_params_groups_for_optimizer( - net, - net.initialize_param_groups(self.lr, train_params), - self.weight_decay - ) + optimizer_params_groups = separate_zero_wd_params_groups_for_optimizer(net, net.initialize_param_groups(self.lr, train_params), self.weight_decay) self._assert_optimizer_param_groups( optimizer_params_groups, excpected_num_groups=2, excpected_num_params_per_group=[net.num_no_decay_params(), net.num_decay_params()], - excpected_weight_decay_per_group=[0, self.weight_decay] + excpected_weight_decay_per_group=[0, self.weight_decay], ) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main()