From e72dad116979c06f200667fea5795dba46866e5b Mon Sep 17 00:00:00 2001 From: Louis Dupont Date: Mon, 12 Dec 2022 16:09:15 +0200 Subject: [PATCH] black --- .../cifar10_training_torch_objects_example.py | 31 +-- .../ema_train_integration_test.py | 46 ++--- tests/unit_tests/kd_ema_test.py | 91 +++++---- tests/unit_tests/kd_trainer_test.py | 178 ++++++++++-------- .../unit_tests/pretrained_models_unit_test.py | 24 +-- tests/unit_tests/test_without_train_test.py | 37 ++-- tests/unit_tests/vit_unit_test.py | 32 ++-- 7 files changed, 239 insertions(+), 200 deletions(-) diff --git a/src/super_gradients/examples/cifar10_training_torch_objects/cifar10_training_torch_objects_example.py b/src/super_gradients/examples/cifar10_training_torch_objects/cifar10_training_torch_objects_example.py index 7244b6958d..297736740b 100644 --- a/src/super_gradients/examples/cifar10_training_torch_objects/cifar10_training_torch_objects_example.py +++ b/src/super_gradients/examples/cifar10_training_torch_objects/cifar10_training_torch_objects_example.py @@ -23,8 +23,8 @@ from torch.utils.data import DataLoader # Define any torch DataLoaders, need at least train & valid loaders -train_dataset = CIFAR10(root='data/', download=True, train=True, transform=ToTensor()) -valid_dataset = CIFAR10(root='data/', download=True, train=False, transform=ToTensor()) +train_dataset = CIFAR10(root="data/", download=True, train=True, transform=ToTensor()) +valid_dataset = CIFAR10(root="data/", download=True, train=False, transform=ToTensor()) train_loader = DataLoader(train_dataset, shuffle=True, batch_size=16) valid_loader = DataLoader(valid_dataset, batch_size=32) @@ -45,21 +45,24 @@ # Define phase callbacks phase_callbacks = [ LRSchedulerCallback(scheduler=rop_lr_scheduler, phase=Phase.VALIDATION_EPOCH_END, metric_name="Accuracy"), - LRSchedulerCallback(scheduler=step_lr_scheduler, phase=Phase.TRAIN_EPOCH_END)] + LRSchedulerCallback(scheduler=step_lr_scheduler, phase=Phase.TRAIN_EPOCH_END), +] # Bring everything together with Trainer and start training trainer = Trainer("Cifar10_external_objects_example", multi_gpu=MultiGPUMode.OFF) -train_params = {"max_epochs": 300, - "phase_callbacks": phase_callbacks, - "initial_lr": lr, - "loss": loss_fn, - "criterion_params": {}, - 'optimizer': optimizer, - "train_metrics_list": [Accuracy(), Top5()], - "valid_metrics_list": [Accuracy(), Top5()], - "metric_to_watch": "Accuracy", - "greater_metric_to_watch_is_better": True, - "lr_scheduler_step_type": "epoch"} +train_params = { + "max_epochs": 300, + "phase_callbacks": phase_callbacks, + "initial_lr": lr, + "loss": loss_fn, + "criterion_params": {}, + "optimizer": optimizer, + "train_metrics_list": [Accuracy(), Top5()], + "valid_metrics_list": [Accuracy(), Top5()], + "metric_to_watch": "Accuracy", + "greater_metric_to_watch_is_better": True, + "lr_scheduler_step_type": "epoch", +} trainer.train(model=net, training_params=train_params, train_loader=train_loader, valid_loader=valid_loader) diff --git a/tests/integration_tests/ema_train_integration_test.py b/tests/integration_tests/ema_train_integration_test.py index 646bbd22aa..233b9b5221 100644 --- a/tests/integration_tests/ema_train_integration_test.py +++ b/tests/integration_tests/ema_train_integration_test.py @@ -10,7 +10,6 @@ def do_nothing(): class CallWrapper: - def __init__(self, f, check_before=do_nothing): self.f = f self.check_before = check_before @@ -21,10 +20,8 @@ def __call__(self, *args, **kwargs): class EMAIntegrationTest(unittest.TestCase): - def _init_model(self) -> None: - self.trainer = Trainer("resnet18_cifar_ema_test", - device='cpu', multi_gpu=MultiGPUMode.OFF) + self.trainer = Trainer("resnet18_cifar_ema_test", device="cpu", multi_gpu=MultiGPUMode.OFF) self.model = models.get("resnet18_cifar", arch_params={"num_classes": 5}) @classmethod @@ -38,21 +35,24 @@ def test_train(self): self._train({"exp_activation": False}) def _train(self, ema_params): - training_params = {"max_epochs": 4, - "lr_updates": [4], - "lr_mode": "step", - "lr_decay_factor": 0.1, - "lr_warmup_epochs": 0, - "initial_lr": 0.1, - "loss": "cross_entropy", - "optimizer": "SGD", - "criterion_params": {}, - "ema": True, - "ema_params": ema_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} + training_params = { + "max_epochs": 4, + "lr_updates": [4], + "lr_mode": "step", + "lr_decay_factor": 0.1, + "lr_warmup_epochs": 0, + "initial_lr": 0.1, + "loss": "cross_entropy", + "optimizer": "SGD", + "criterion_params": {}, + "ema": True, + "ema_params": ema_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, + } def before_test(): self.assertEqual(self.trainer.net, self.trainer.ema_model.ema) @@ -63,12 +63,12 @@ def before_train_epoch(): self.trainer.test = CallWrapper(self.trainer.test, check_before=before_test) self.trainer._train_epoch = CallWrapper(self.trainer._train_epoch, check_before=before_train_epoch) - self.trainer.train(model=self.model, training_params=training_params, - train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + self.trainer.train( + model=self.model, training_params=training_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader() + ) self.assertIsNotNone(self.trainer.ema_model) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/kd_ema_test.py b/tests/unit_tests/kd_ema_test.py index c707f0721c..dce5dab97c 100644 --- a/tests/unit_tests/kd_ema_test.py +++ b/tests/unit_tests/kd_ema_test.py @@ -13,30 +13,42 @@ class KDEMATest(unittest.TestCase): @classmethod def setUp(cls): - cls.sg_trained_teacher = Trainer("sg_trained_teacher", device='cpu') - - cls.kd_train_params = {"max_epochs": 3, "lr_updates": [1], "lr_decay_factor": 0.1, "lr_mode": "step", - "lr_warmup_epochs": 0, "initial_lr": 0.1, - "loss": KDLogitsLoss(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", - 'loss_logging_items_names': ["Loss", "Task Loss", "Distillation Loss"], - "greater_metric_to_watch_is_better": True, "average_best_models": False, - "ema": True} + cls.sg_trained_teacher = Trainer("sg_trained_teacher", device="cpu") + + cls.kd_train_params = { + "max_epochs": 3, + "lr_updates": [1], + "lr_decay_factor": 0.1, + "lr_mode": "step", + "lr_warmup_epochs": 0, + "initial_lr": 0.1, + "loss": KDLogitsLoss(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", + "loss_logging_items_names": ["Loss", "Task Loss", "Distillation Loss"], + "greater_metric_to_watch_is_better": True, + "average_best_models": False, + "ema": True, + } def test_teacher_ema_not_duplicated(self): """Check that the teacher EMA is a reference to the teacher net (not a copy).""" - kd_model = KDTrainer("test_teacher_ema_not_duplicated", device='cpu') - student = models.get('resnet18', arch_params={'num_classes': 1000}) - teacher = models.get('resnet50', arch_params={'num_classes': 1000}, - pretrained_weights="imagenet") + kd_model = KDTrainer("test_teacher_ema_not_duplicated", device="cpu") + student = models.get("resnet18", arch_params={"num_classes": 1000}) + teacher = models.get("resnet50", arch_params={"num_classes": 1000}, pretrained_weights="imagenet") - kd_model.train(training_params=self.kd_train_params, student=student, teacher=teacher, - train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + kd_model.train( + training_params=self.kd_train_params, + student=student, + teacher=teacher, + train_loader=classification_test_dataloader(), + valid_loader=classification_test_dataloader(), + ) self.assertTrue(kd_model.ema_model.ema.module.teacher is kd_model.net.module.teacher) self.assertTrue(kd_model.ema_model.ema.module.student is not kd_model.net.module.student) @@ -46,27 +58,33 @@ def test_kd_ckpt_reload_net(self): # Create a KD trainer and train it train_params = self.kd_train_params.copy() - kd_model = KDTrainer("test_kd_ema_ckpt_reload", device='cpu') - student = models.get('resnet18', arch_params={'num_classes': 1000}) - teacher = models.get('resnet50', arch_params={'num_classes': 1000}, - pretrained_weights="imagenet") - - kd_model.train(training_params=self.kd_train_params, student=student, teacher=teacher, - train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + kd_model = KDTrainer("test_kd_ema_ckpt_reload", device="cpu") + student = models.get("resnet18", arch_params={"num_classes": 1000}) + teacher = models.get("resnet50", arch_params={"num_classes": 1000}, pretrained_weights="imagenet") + + kd_model.train( + training_params=self.kd_train_params, + student=student, + teacher=teacher, + train_loader=classification_test_dataloader(), + valid_loader=classification_test_dataloader(), + ) ema_model = kd_model.ema_model.ema net = kd_model.net # Load the trained KD trainer - kd_model = KDTrainer("test_kd_ema_ckpt_reload", device='cpu') - student = models.get('resnet18', arch_params={'num_classes': 1000}) - teacher = models.get('resnet50', arch_params={'num_classes': 1000}, - pretrained_weights="imagenet") + kd_model = KDTrainer("test_kd_ema_ckpt_reload", device="cpu") + student = models.get("resnet18", arch_params={"num_classes": 1000}) + teacher = models.get("resnet50", arch_params={"num_classes": 1000}, pretrained_weights="imagenet") train_params["resume"] = True - kd_model.train(training_params=train_params, student=student, teacher=teacher, - train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + kd_model.train( + training_params=train_params, + student=student, + teacher=teacher, + train_loader=classification_test_dataloader(), + valid_loader=classification_test_dataloader(), + ) reloaded_ema_model = kd_model.ema_model.ema reloaded_net = kd_model.net @@ -80,12 +98,11 @@ def test_kd_ckpt_reload_net(self): self.assertTrue(not check_models_have_same_weights(reloaded_net, ema_model)) # loaded student ema == loaded student net (since load_ema_as_net = False) - self.assertTrue( - not check_models_have_same_weights(reloaded_ema_model.module.student, reloaded_net.module.student)) + self.assertTrue(not check_models_have_same_weights(reloaded_ema_model.module.student, reloaded_net.module.student)) # loaded teacher ema == loaded teacher net (teacher always loads ema) self.assertTrue(check_models_have_same_weights(reloaded_ema_model.module.teacher, reloaded_net.module.teacher)) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/kd_trainer_test.py b/tests/unit_tests/kd_trainer_test.py index 8ff4385584..b19c908762 100644 --- a/tests/unit_tests/kd_trainer_test.py +++ b/tests/unit_tests/kd_trainer_test.py @@ -37,133 +37,147 @@ def __call__(self, context: PhaseContext): class KDTrainerTest(unittest.TestCase): @classmethod def setUp(cls): - cls.kd_train_params = {"max_epochs": 3, "lr_updates": [1], "lr_decay_factor": 0.1, "lr_mode": "step", - "lr_warmup_epochs": 0, "initial_lr": 0.1, - "loss": KDLogitsLoss(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", - 'loss_logging_items_names': ["Loss", "Task Loss", "Distillation Loss"], - "greater_metric_to_watch_is_better": True, "average_best_models": False} + cls.kd_train_params = { + "max_epochs": 3, + "lr_updates": [1], + "lr_decay_factor": 0.1, + "lr_mode": "step", + "lr_warmup_epochs": 0, + "initial_lr": 0.1, + "loss": KDLogitsLoss(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", + "loss_logging_items_names": ["Loss", "Task Loss", "Distillation Loss"], + "greater_metric_to_watch_is_better": True, + "average_best_models": False, + } def test_teacher_sg_module_methods(self): - student = models.get('resnet18', arch_params={'num_classes': 1000}) - teacher = models.get('resnet50', arch_params={'num_classes': 1000}, - pretrained_weights="imagenet") - kd_module = KDModule(arch_params={}, - student=student, - teacher=teacher - ) + student = models.get("resnet18", arch_params={"num_classes": 1000}) + teacher = models.get("resnet50", arch_params={"num_classes": 1000}, pretrained_weights="imagenet") + kd_module = KDModule(arch_params={}, student=student, teacher=teacher) initial_param_groups = kd_module.initialize_param_groups(lr=0.1, training_params={}) - updated_param_groups = kd_module.update_param_groups(param_groups=initial_param_groups, lr=0.2, - epoch=0, iter=0, training_params={}, - total_batch=None) + updated_param_groups = kd_module.update_param_groups(param_groups=initial_param_groups, lr=0.2, epoch=0, iter=0, training_params={}, total_batch=None) - self.assertTrue(initial_param_groups[0]['lr'] == 0.2 == updated_param_groups[0]['lr']) + self.assertTrue(initial_param_groups[0]["lr"] == 0.2 == updated_param_groups[0]["lr"]) def test_train_kd_module_external_models(self): - sg_model = KDTrainer("test_train_kd_module_external_models", device='cpu') + sg_model = KDTrainer("test_train_kd_module_external_models", device="cpu") teacher_model = ResNet50(arch_params={}, num_classes=5) student_model = ResNet18(arch_params={}, num_classes=5) - sg_model.train(training_params=self.kd_train_params, student=deepcopy(student_model), teacher=teacher_model, - train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + sg_model.train( + training_params=self.kd_train_params, + student=deepcopy(student_model), + teacher=teacher_model, + train_loader=classification_test_dataloader(), + valid_loader=classification_test_dataloader(), + ) # TEACHER WEIGHT'S SHOULD REMAIN THE SAME - self.assertTrue( - check_models_have_same_weights(teacher_model, sg_model.net.module.teacher)) + self.assertTrue(check_models_have_same_weights(teacher_model, sg_model.net.module.teacher)) # STUDENT WEIGHT'S SHOULD NOT REMAIN THE SAME - self.assertFalse( - check_models_have_same_weights(student_model, sg_model.net.module.student)) + self.assertFalse(check_models_have_same_weights(student_model, sg_model.net.module.student)) def test_train_model_with_input_adapter(self): - kd_trainer = KDTrainer("train_kd_module_with_with_input_adapter", device='cpu') - student = models.get('resnet18', arch_params={'num_classes': 5}) - teacher = models.get('resnet50', arch_params={'num_classes': 5}, - pretrained_weights="imagenet") - - adapter = NormalizationAdapter(mean_original=[0.485, 0.456, 0.406], - std_original=[0.229, 0.224, 0.225], - mean_required=[0.5, 0.5, 0.5], - std_required=[0.5, 0.5, 0.5]) - - kd_arch_params = { - "teacher_input_adapter": adapter} - kd_trainer.train(training_params=self.kd_train_params, student=student, teacher=teacher, - kd_arch_params=kd_arch_params, train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + kd_trainer = KDTrainer("train_kd_module_with_with_input_adapter", device="cpu") + student = models.get("resnet18", arch_params={"num_classes": 5}) + teacher = models.get("resnet50", arch_params={"num_classes": 5}, pretrained_weights="imagenet") + + adapter = NormalizationAdapter( + mean_original=[0.485, 0.456, 0.406], std_original=[0.229, 0.224, 0.225], mean_required=[0.5, 0.5, 0.5], std_required=[0.5, 0.5, 0.5] + ) + + kd_arch_params = {"teacher_input_adapter": adapter} + kd_trainer.train( + training_params=self.kd_train_params, + student=student, + teacher=teacher, + kd_arch_params=kd_arch_params, + train_loader=classification_test_dataloader(), + valid_loader=classification_test_dataloader(), + ) self.assertEqual(kd_trainer.net.module.teacher_input_adapter, adapter) def test_load_ckpt_best_for_student(self): - kd_trainer = KDTrainer("test_load_ckpt_best", device='cpu') - student = models.get('resnet18', arch_params={'num_classes': 5}) - teacher = models.get('resnet50', arch_params={'num_classes': 5}, - pretrained_weights="imagenet") + kd_trainer = KDTrainer("test_load_ckpt_best", device="cpu") + student = models.get("resnet18", arch_params={"num_classes": 5}) + teacher = models.get("resnet50", arch_params={"num_classes": 5}, pretrained_weights="imagenet") train_params = self.kd_train_params.copy() train_params["max_epochs"] = 1 - kd_trainer.train(training_params=train_params, student=student, teacher=teacher, - train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + kd_trainer.train( + training_params=train_params, + student=student, + teacher=teacher, + train_loader=classification_test_dataloader(), + valid_loader=classification_test_dataloader(), + ) best_student_ckpt = os.path.join(kd_trainer.checkpoints_dir_path, "ckpt_best.pth") - student_reloaded = models.get('resnet18', arch_params={'num_classes': 5}, - checkpoint_path=best_student_ckpt) + student_reloaded = models.get("resnet18", arch_params={"num_classes": 5}, checkpoint_path=best_student_ckpt) - self.assertTrue( - check_models_have_same_weights(student_reloaded, kd_trainer.net.module.student)) + self.assertTrue(check_models_have_same_weights(student_reloaded, kd_trainer.net.module.student)) def test_load_ckpt_best_for_student_with_ema(self): - kd_trainer = KDTrainer("test_load_ckpt_best", device='cpu') - student = models.get('resnet18', arch_params={'num_classes': 5}) - teacher = models.get('resnet50', arch_params={'num_classes': 5}, - pretrained_weights="imagenet") + kd_trainer = KDTrainer("test_load_ckpt_best", device="cpu") + student = models.get("resnet18", arch_params={"num_classes": 5}) + teacher = models.get("resnet50", arch_params={"num_classes": 5}, pretrained_weights="imagenet") train_params = self.kd_train_params.copy() train_params["max_epochs"] = 1 train_params["ema"] = True - kd_trainer.train(training_params=train_params, student=student, teacher=teacher, - train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + kd_trainer.train( + training_params=train_params, + student=student, + teacher=teacher, + train_loader=classification_test_dataloader(), + valid_loader=classification_test_dataloader(), + ) best_student_ckpt = os.path.join(kd_trainer.checkpoints_dir_path, "ckpt_best.pth") - student_reloaded = models.get('resnet18', arch_params={'num_classes': 5}, - checkpoint_path=best_student_ckpt) + student_reloaded = models.get("resnet18", arch_params={"num_classes": 5}, checkpoint_path=best_student_ckpt) - self.assertTrue( - check_models_have_same_weights(student_reloaded, kd_trainer.ema_model.ema.module.student)) + self.assertTrue(check_models_have_same_weights(student_reloaded, kd_trainer.ema_model.ema.module.student)) def test_resume_kd_training(self): - kd_trainer = KDTrainer("test_resume_training_start", device='cpu') - student = models.get('resnet18', arch_params={'num_classes': 5}) - teacher = models.get('resnet50', arch_params={'num_classes': 5}, - pretrained_weights="imagenet") + kd_trainer = KDTrainer("test_resume_training_start", device="cpu") + student = models.get("resnet18", arch_params={"num_classes": 5}) + teacher = models.get("resnet50", arch_params={"num_classes": 5}, pretrained_weights="imagenet") train_params = self.kd_train_params.copy() train_params["max_epochs"] = 1 - kd_trainer.train(training_params=train_params, student=student, teacher=teacher, - train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + kd_trainer.train( + training_params=train_params, + student=student, + teacher=teacher, + train_loader=classification_test_dataloader(), + valid_loader=classification_test_dataloader(), + ) latest_net = deepcopy(kd_trainer.net) - kd_trainer = KDTrainer("test_resume_training_start", device='cpu') - student = models.get('resnet18', arch_params={'num_classes': 5}) - teacher = models.get('resnet50', arch_params={'num_classes': 5}, - pretrained_weights="imagenet") + kd_trainer = KDTrainer("test_resume_training_start", device="cpu") + student = models.get("resnet18", arch_params={"num_classes": 5}) + teacher = models.get("resnet50", arch_params={"num_classes": 5}, pretrained_weights="imagenet") train_params["max_epochs"] = 2 train_params["resume"] = True collector = PreTrainingNetCollector() train_params["phase_callbacks"] = [collector] - kd_trainer.train(training_params=train_params, student=student, teacher=teacher, - train_loader=classification_test_dataloader(), - valid_loader=classification_test_dataloader()) + kd_trainer.train( + training_params=train_params, + student=student, + teacher=teacher, + train_loader=classification_test_dataloader(), + valid_loader=classification_test_dataloader(), + ) - self.assertTrue( - check_models_have_same_weights(collector.net, latest_net)) + self.assertTrue(check_models_have_same_weights(collector.net, latest_net)) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/pretrained_models_unit_test.py b/tests/unit_tests/pretrained_models_unit_test.py index b040010483..1ee14f511e 100644 --- a/tests/unit_tests/pretrained_models_unit_test.py +++ b/tests/unit_tests/pretrained_models_unit_test.py @@ -14,30 +14,24 @@ def setUp(self) -> None: self.imagenet_pretrained_models = ["resnet50", "repvgg_a0", "regnetY800"] def test_pretrained_resnet50_imagenet(self): - trainer = Trainer('imagenet_pretrained_resnet50_unit_test', - multi_gpu=MultiGPUMode.OFF) + trainer = Trainer("imagenet_pretrained_resnet50_unit_test", multi_gpu=MultiGPUMode.OFF) model = models.get("resnet50", pretrained_weights="imagenet") - trainer.test(model=model, test_loader=classification_test_dataloader(), test_metrics_list=[Accuracy()], - metrics_progress_verbose=True) + trainer.test(model=model, test_loader=classification_test_dataloader(), test_metrics_list=[Accuracy()], metrics_progress_verbose=True) def test_pretrained_regnetY800_imagenet(self): - trainer = Trainer('imagenet_pretrained_regnetY800_unit_test', - multi_gpu=MultiGPUMode.OFF) + trainer = Trainer("imagenet_pretrained_regnetY800_unit_test", multi_gpu=MultiGPUMode.OFF) model = models.get("regnetY800", pretrained_weights="imagenet") - trainer.test(model=model, test_loader=classification_test_dataloader(), test_metrics_list=[Accuracy()], - metrics_progress_verbose=True) + trainer.test(model=model, test_loader=classification_test_dataloader(), test_metrics_list=[Accuracy()], metrics_progress_verbose=True) def test_pretrained_repvgg_a0_imagenet(self): - trainer = Trainer('imagenet_pretrained_repvgg_a0_unit_test', - multi_gpu=MultiGPUMode.OFF) + trainer = Trainer("imagenet_pretrained_repvgg_a0_unit_test", multi_gpu=MultiGPUMode.OFF) model = models.get("repvgg_a0", pretrained_weights="imagenet", arch_params={"build_residual_branches": True}) - trainer.test(model=model, test_loader=classification_test_dataloader(), test_metrics_list=[Accuracy()], - metrics_progress_verbose=True) + trainer.test(model=model, test_loader=classification_test_dataloader(), test_metrics_list=[Accuracy()], metrics_progress_verbose=True) def tearDown(self) -> None: - if os.path.exists('~/.cache/torch/hub/'): - shutil.rmtree('~/.cache/torch/hub/') + if os.path.exists("~/.cache/torch/hub/"): + shutil.rmtree("~/.cache/torch/hub/") -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/test_without_train_test.py b/tests/unit_tests/test_without_train_test.py index 4f4c0a99d6..3a93e27fad 100644 --- a/tests/unit_tests/test_without_train_test.py +++ b/tests/unit_tests/test_without_train_test.py @@ -2,8 +2,7 @@ import unittest import os from super_gradients import Trainer -from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader, \ - detection_test_dataloader, segmentation_test_dataloader +from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader, detection_test_dataloader, segmentation_test_dataloader from super_gradients.training.metrics import Accuracy, Top5 from super_gradients.training import MultiGPUMode, models from super_gradients.training.metrics.detection_metrics import DetectionMetrics @@ -15,51 +14,53 @@ class TestWithoutTrainTest(unittest.TestCase): @classmethod def setUp(cls): # NAMES FOR THE EXPERIMENTS TO LATER DELETE - cls.folder_names = ['test_classification_model', 'test_detection_model', 'test_segmentation_model'] + cls.folder_names = ["test_classification_model", "test_detection_model", "test_segmentation_model"] @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 @staticmethod - def get_detection_trainer(name=''): - trainer = Trainer(name, - multi_gpu=MultiGPUMode.OFF) + def get_detection_trainer(name=""): + trainer = Trainer(name, multi_gpu=MultiGPUMode.OFF) model = models.get("yolox_s", num_classes=5) return trainer, model @staticmethod - def get_segmentation_trainer(name=''): + def get_segmentation_trainer(name=""): shelfnet_lw_arch_params = {"num_classes": 5} trainer = Trainer(name) - model = models.get('shelfnet34_lw', arch_params=shelfnet_lw_arch_params) + model = models.get("shelfnet34_lw", arch_params=shelfnet_lw_arch_params) return trainer, model def test_test_without_train(self): trainer, model = self.get_classification_trainer(self.folder_names[0]) - assert isinstance(trainer.test(model=model, silent_mode=True, - test_metrics_list=[Accuracy(), Top5()], test_loader=classification_test_dataloader()), tuple) + assert isinstance( + trainer.test(model=model, silent_mode=True, test_metrics_list=[Accuracy(), Top5()], test_loader=classification_test_dataloader()), tuple + ) trainer, model = self.get_detection_trainer(self.folder_names[1]) test_metrics = [DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=5)] - assert isinstance(trainer.test(model=model, silent_mode=True, - test_metrics_list=test_metrics, test_loader=detection_test_dataloader(image_size=320)), tuple) + assert isinstance( + trainer.test(model=model, silent_mode=True, test_metrics_list=test_metrics, test_loader=detection_test_dataloader(image_size=320)), tuple + ) trainer, model = self.get_segmentation_trainer(self.folder_names[2]) - assert isinstance(trainer.test(model=model, silent_mode=True, - test_metrics_list=[IoU(21), PixelAccuracy()], test_loader=segmentation_test_dataloader()), tuple) + assert isinstance( + trainer.test(model=model, silent_mode=True, test_metrics_list=[IoU(21), PixelAccuracy()], test_loader=segmentation_test_dataloader()), tuple + ) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/vit_unit_test.py b/tests/unit_tests/vit_unit_test.py index d61fada7af..928385b61b 100644 --- a/tests/unit_tests/vit_unit_test.py +++ b/tests/unit_tests/vit_unit_test.py @@ -7,25 +7,35 @@ class TestViT(unittest.TestCase): - def setUp(self): self.arch_params = {"image_size": (224, 224), "patch_size": (16, 16), "num_classes": 10} - 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"} + 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", + } def test_train_vit(self): """ Validate vit_base """ - trainer = Trainer("test_vit_base", device='cpu') - model = models.get('vit_base', arch_params=self.arch_params, num_classes=5) - trainer.train(model=model, training_params=self.train_params, - train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader()) + trainer = Trainer("test_vit_base", device="cpu") + model = models.get("vit_base", arch_params=self.arch_params, num_classes=5) + trainer.train( + model=model, training_params=self.train_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader() + ) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main()