From 5e48cf400b645459e1981f30532159df3e72431e Mon Sep 17 00:00:00 2001 From: shayaharon Date: Thu, 22 Dec 2022 11:52:00 +0200 Subject: [PATCH 1/7] refs in unit tests updated --- tests/end_to_end_tests/cifar_trainer_test.py | 5 ++- tests/unit_tests/factories_test.py | 3 +- .../initialize_with_dataloaders_test.py | 40 ++++++++++--------- tests/unit_tests/kd_ema_test.py | 13 +++--- tests/unit_tests/kd_trainer_test.py | 13 +++--- .../unit_tests/pretrained_models_unit_test.py | 5 ++- tests/unit_tests/save_ckpt_test.py | 37 ++++++++++------- tests/unit_tests/strictload_enum_test.py | 23 ++++++----- tests/unit_tests/test_without_train_test.py | 3 +- 9 files changed, 81 insertions(+), 61 deletions(-) diff --git a/tests/end_to_end_tests/cifar_trainer_test.py b/tests/end_to_end_tests/cifar_trainer_test.py index 9e5e509cd9..7b91a59cba 100644 --- a/tests/end_to_end_tests/cifar_trainer_test.py +++ b/tests/end_to_end_tests/cifar_trainer_test.py @@ -1,5 +1,6 @@ import unittest +from super_gradients.common.object_names import Models from super_gradients.training import models import super_gradients @@ -18,7 +19,7 @@ def test_train_cifar10_dataloader(self): super_gradients.init_trainer() trainer = Trainer("test") cifar10_train_dl, cifar10_val_dl = cifar10_train(), cifar10_val() - model = models.get("resnet18_cifar", arch_params={"num_classes": 10}) + model = models.get(Models.RESNET18_CIFAR, arch_params={"num_classes": 10}) trainer.train( model=model, training_params={ @@ -37,7 +38,7 @@ def test_train_cifar100_dataloader(self): super_gradients.init_trainer() trainer = Trainer("test") cifar100_train_dl, cifar100_val_dl = cifar100_train(), cifar100_val() - model = models.get("resnet18_cifar", arch_params={"num_classes": 100}) + model = models.get(Models.RESNET18_CIFAR, arch_params={"num_classes": 100}) trainer.train( model=model, training_params={ diff --git a/tests/unit_tests/factories_test.py b/tests/unit_tests/factories_test.py index 6fa60a5d46..73352caf23 100644 --- a/tests/unit_tests/factories_test.py +++ b/tests/unit_tests/factories_test.py @@ -5,6 +5,7 @@ from super_gradients import Trainer from super_gradients.common.decorators.factory_decorator import resolve_param from super_gradients.common.factories.activations_type_factory import ActivationsTypeFactory +from super_gradients.common.object_names import Models from super_gradients.training import models from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader from super_gradients.training.metrics import Accuracy, Top5 @@ -14,7 +15,7 @@ class FactoriesTest(unittest.TestCase): def test_training_with_factories(self): trainer = Trainer("test_train_with_factories") - net = models.get("resnet18", num_classes=5) + net = models.get(Models.RESNET18, num_classes=5) train_params = { "max_epochs": 2, "lr_updates": [1], diff --git a/tests/unit_tests/initialize_with_dataloaders_test.py b/tests/unit_tests/initialize_with_dataloaders_test.py index c1efce8294..b81a5c6139 100644 --- a/tests/unit_tests/initialize_with_dataloaders_test.py +++ b/tests/unit_tests/initialize_with_dataloaders_test.py @@ -1,5 +1,6 @@ import unittest +from super_gradients.common.object_names import Models from super_gradients.training import models from super_gradients import Trainer @@ -27,25 +28,28 @@ 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, - ) + model = models.get(Models.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, + ) self.assertTrue(0 < trainer.best_metric.item() < 1) -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 dce5dab97c..a940268807 100644 --- a/tests/unit_tests/kd_ema_test.py +++ b/tests/unit_tests/kd_ema_test.py @@ -1,5 +1,6 @@ import unittest +from super_gradients.common.object_names import Models from super_gradients.training import models from super_gradients.training import Trainer from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader @@ -39,8 +40,8 @@ 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") + student = models.get(Models.RESNET18, arch_params={"num_classes": 1000}) + teacher = models.get(Models.RESNET50, arch_params={"num_classes": 1000}, pretrained_weights="imagenet") kd_model.train( training_params=self.kd_train_params, @@ -59,8 +60,8 @@ 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") + student = models.get(Models.RESNET18, arch_params={"num_classes": 1000}) + teacher = models.get(Models.RESNET50, arch_params={"num_classes": 1000}, pretrained_weights="imagenet") kd_model.train( training_params=self.kd_train_params, @@ -74,8 +75,8 @@ def test_kd_ckpt_reload_net(self): # 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") + student = models.get(Models.RESNET18, arch_params={"num_classes": 1000}) + teacher = models.get(Models.RESNET50, arch_params={"num_classes": 1000}, pretrained_weights="imagenet") train_params["resume"] = True kd_model.train( diff --git a/tests/unit_tests/kd_trainer_test.py b/tests/unit_tests/kd_trainer_test.py index b19c908762..19debbb043 100644 --- a/tests/unit_tests/kd_trainer_test.py +++ b/tests/unit_tests/kd_trainer_test.py @@ -2,6 +2,7 @@ import unittest from copy import deepcopy +from super_gradients.common.object_names import Models from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader from super_gradients.training.kd_trainer.kd_trainer import KDTrainer import torch @@ -58,7 +59,7 @@ def setUp(cls): 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") + teacher = models.get(Models.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={}) @@ -88,7 +89,7 @@ def test_train_kd_module_external_models(self): 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") + teacher = models.get(Models.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] @@ -109,7 +110,7 @@ def test_train_model_with_input_adapter(self): 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") + teacher = models.get(Models.RESNET50, arch_params={"num_classes": 5}, pretrained_weights="imagenet") train_params = self.kd_train_params.copy() train_params["max_epochs"] = 1 kd_trainer.train( @@ -128,7 +129,7 @@ def test_load_ckpt_best_for_student(self): 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") + teacher = models.get(Models.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 @@ -148,7 +149,7 @@ def test_load_ckpt_best_for_student_with_ema(self): 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") + teacher = models.get(Models.RESNET50, arch_params={"num_classes": 5}, pretrained_weights="imagenet") train_params = self.kd_train_params.copy() train_params["max_epochs"] = 1 kd_trainer.train( @@ -162,7 +163,7 @@ 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") + teacher = models.get(Models.RESNET50, arch_params={"num_classes": 5}, pretrained_weights="imagenet") train_params["max_epochs"] = 2 train_params["resume"] = True diff --git a/tests/unit_tests/pretrained_models_unit_test.py b/tests/unit_tests/pretrained_models_unit_test.py index 1ee14f511e..9478f4fb33 100644 --- a/tests/unit_tests/pretrained_models_unit_test.py +++ b/tests/unit_tests/pretrained_models_unit_test.py @@ -1,5 +1,6 @@ import unittest import super_gradients +from super_gradients.common.object_names import Models from super_gradients.training import MultiGPUMode, models from super_gradients.training import Trainer from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader @@ -11,11 +12,11 @@ class PretrainedModelsUnitTest(unittest.TestCase): def setUp(self) -> None: super_gradients.init_trainer() - self.imagenet_pretrained_models = ["resnet50", "repvgg_a0", "regnetY800"] + self.imagenet_pretrained_models = [Models.RESNET50, "repvgg_a0", "regnetY800"] def test_pretrained_resnet50_imagenet(self): trainer = Trainer("imagenet_pretrained_resnet50_unit_test", multi_gpu=MultiGPUMode.OFF) - model = models.get("resnet50", pretrained_weights="imagenet") + model = models.get(Models.RESNET50, pretrained_weights="imagenet") trainer.test(model=model, test_loader=classification_test_dataloader(), test_metrics_list=[Accuracy()], metrics_progress_verbose=True) def test_pretrained_regnetY800_imagenet(self): diff --git a/tests/unit_tests/save_ckpt_test.py b/tests/unit_tests/save_ckpt_test.py index 8b64dfa1b0..dcb5208744 100644 --- a/tests/unit_tests/save_ckpt_test.py +++ b/tests/unit_tests/save_ckpt_test.py @@ -3,38 +3,47 @@ from super_gradients.training import Trainer, models from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader from super_gradients.training.metrics import Accuracy, Top5 +from super_gradients.common.object_names import Models 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") # Build Model - model = models.get("resnet18_cifar", arch_params={"num_classes": 10}) + model = models.get(Models.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/strictload_enum_test.py b/tests/unit_tests/strictload_enum_test.py index 406cb3c057..20257b1922 100644 --- a/tests/unit_tests/strictload_enum_test.py +++ b/tests/unit_tests/strictload_enum_test.py @@ -3,6 +3,7 @@ import tempfile import unittest +from super_gradients.common.object_names import Models from super_gradients.common.sg_loggers import BaseSGLogger from super_gradients.training import Trainer import torch @@ -98,8 +99,8 @@ def check_models_have_same_weights(self, model_1, model_2): def test_strict_load_on(self): # Define Model - model = models.get("resnet18", arch_params={"num_classes": 1000}) - pretrained_model = models.get("resnet18", arch_params={"num_classes": 1000}, pretrained_weights="imagenet") + model = models.get(Models.RESNET18, arch_params={"num_classes": 1000}) + pretrained_model = models.get(Models.RESNET18, arch_params={"num_classes": 1000}, pretrained_weights="imagenet") # Make sure we initialized a model with different weights assert not self.check_models_have_same_weights(model, pretrained_model) @@ -107,15 +108,15 @@ def test_strict_load_on(self): pretrained_sd_path = os.path.join(self.temp_working_file_dir, "pretrained_net_strict_load_on.pth") torch.save(pretrained_model.state_dict(), pretrained_sd_path) - model = models.get("resnet18", arch_params={"num_classes": 1000}, checkpoint_path=pretrained_sd_path, strict_load=StrictLoad.ON) + model = models.get(Models.RESNET18, arch_params={"num_classes": 1000}, checkpoint_path=pretrained_sd_path, strict_load=StrictLoad.ON) # Assert the weights were loaded correctly assert self.check_models_have_same_weights(model, pretrained_model) def test_strict_load_off(self): # Define Model - model = models.get("resnet18", arch_params={"num_classes": 1000}) - pretrained_model = models.get("resnet18", arch_params={"num_classes": 1000}, pretrained_weights="imagenet") + model = models.get(Models.RESNET18, arch_params={"num_classes": 1000}) + pretrained_model = models.get(Models.RESNET18, arch_params={"num_classes": 1000}, pretrained_weights="imagenet") # Make sure we initialized a model with different weights assert not self.check_models_have_same_weights(model, pretrained_model) @@ -125,17 +126,17 @@ def test_strict_load_off(self): torch.save(pretrained_model.state_dict(), pretrained_sd_path) with self.assertRaises(RuntimeError): - models.get("resnet18", arch_params={"num_classes": 1000}, checkpoint_path=pretrained_sd_path, strict_load=StrictLoad.ON) + models.get(Models.RESNET18, arch_params={"num_classes": 1000}, checkpoint_path=pretrained_sd_path, strict_load=StrictLoad.ON) - model = models.get("resnet18", arch_params={"num_classes": 1000}, checkpoint_path=pretrained_sd_path, strict_load=StrictLoad.OFF) + model = models.get(Models.RESNET18, arch_params={"num_classes": 1000}, checkpoint_path=pretrained_sd_path, strict_load=StrictLoad.OFF) del model.linear # Assert the weights were loaded correctly assert self.check_models_have_same_weights(model, pretrained_model) def test_strict_load_no_key_matching_sg_checkpoint(self): # Define Model - model = models.get("resnet18", arch_params={"num_classes": 1000}) - pretrained_model = models.get("resnet18", arch_params={"num_classes": 1000}, pretrained_weights="imagenet") + model = models.get(Models.RESNET18, arch_params={"num_classes": 1000}) + pretrained_model = models.get(Models.RESNET18, arch_params={"num_classes": 1000}, pretrained_weights="imagenet") # Make sure we initialized a model with different weights assert not self.check_models_have_same_weights(model, pretrained_model) @@ -144,9 +145,9 @@ def test_strict_load_no_key_matching_sg_checkpoint(self): torch.save(self.change_state_dict_keys(pretrained_model.state_dict()), pretrained_sd_path) with self.assertRaises(RuntimeError): - models.get("resnet18", arch_params={"num_classes": 1000}, checkpoint_path=pretrained_sd_path, strict_load=StrictLoad.ON) + models.get(Models.RESNET18, arch_params={"num_classes": 1000}, checkpoint_path=pretrained_sd_path, strict_load=StrictLoad.ON) - model = models.get("resnet18", arch_params={"num_classes": 1000}, checkpoint_path=pretrained_sd_path, strict_load=StrictLoad.NO_KEY_MATCHING) + model = models.get(Models.RESNET18, arch_params={"num_classes": 1000}, checkpoint_path=pretrained_sd_path, strict_load=StrictLoad.NO_KEY_MATCHING) # Assert the weights were loaded correctly assert self.check_models_have_same_weights(model, pretrained_model) diff --git a/tests/unit_tests/test_without_train_test.py b/tests/unit_tests/test_without_train_test.py index 3a93e27fad..cfc276c123 100644 --- a/tests/unit_tests/test_without_train_test.py +++ b/tests/unit_tests/test_without_train_test.py @@ -2,6 +2,7 @@ import unittest import os from super_gradients import Trainer +from super_gradients.common.object_names import Models 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 @@ -26,7 +27,7 @@ def tearDownClass(cls) -> None: @staticmethod def get_classification_trainer(name=""): trainer = Trainer(name) - model = models.get("resnet18", num_classes=5) + model = models.get(Models.RESNET18, num_classes=5) return trainer, model @staticmethod From 531c23e147097a43db5579d5519cd23287d80db5 Mon Sep 17 00:00:00 2001 From: shayaharon Date: Thu, 22 Dec 2022 14:58:17 +0200 Subject: [PATCH 2/7] formatting --- .../ddrnet_classification_example.py | 87 ++--- .../examples/early_stop/early_stop_example.py | 29 +- .../clearml_logger_example.py | 3 +- .../deci_platform_logger_example.py | 4 +- .../quantization/resnet_qat_example.py | 3 +- .../regseg_transfer_learning_example.py | 79 ++--- tests/end_to_end_tests/trainer_test.py | 65 ++-- .../ema_train_integration_test.py | 3 +- tests/integration_tests/lr_test.py | 50 +-- .../pretrained_models_test.py | 307 ++++++++++-------- tests/unit_tests/config_inspector_test.py | 3 +- tests/unit_tests/dataset_statistics_test.py | 38 +-- tests/unit_tests/detection_utils_test.py | 12 +- tests/unit_tests/double_training_test.py | 3 +- tests/unit_tests/forward_pass_prep_fn_test.py | 39 ++- tests/unit_tests/kd_trainer_test.py | 16 +- .../local_ckpt_head_replacement_test.py | 5 +- tests/unit_tests/loss_loggings_test.py | 90 +++-- .../unit_tests/pretrained_models_unit_test.py | 4 +- .../unit_tests/quantization_utility_tests.py | 6 +- tests/unit_tests/test_without_train_test.py | 4 +- tests/unit_tests/train_after_test_test.py | 5 +- .../train_with_intialized_param_args_test.py | 169 ++++++---- tests/unit_tests/vit_unit_test.py | 3 +- 24 files changed, 595 insertions(+), 432 deletions(-) diff --git a/src/super_gradients/examples/ddrnet_imagenet/ddrnet_classification_example.py b/src/super_gradients/examples/ddrnet_imagenet/ddrnet_classification_example.py index 1060ef4dab..58a89657be 100644 --- a/src/super_gradients/examples/ddrnet_imagenet/ddrnet_classification_example.py +++ b/src/super_gradients/examples/ddrnet_imagenet/ddrnet_classification_example.py @@ -15,6 +15,7 @@ import torch from super_gradients.common import MultiGPUMode +from super_gradients.common.object_names import Models from super_gradients.training.datasets.datasets_utils import RandomResizedCropAndInterpolation from torchvision.transforms import RandomHorizontalFlip, ColorJitter, ToTensor, Normalize import super_gradients @@ -22,6 +23,7 @@ import argparse from super_gradients.training.metrics import Accuracy, Top5 from super_gradients.training.datasets.data_augmentation import RandomErase + parser = argparse.ArgumentParser() super_gradients.init_trainer() @@ -29,53 +31,56 @@ parser.add_argument("--max_epochs", type=int, default=100) parser.add_argument("--batch", type=int, default=3) parser.add_argument("--experiment_name", type=str, default="ddrnet_23") -parser.add_argument("-s", "--slim", action="store_true", help='train the slim version of DDRNet23') +parser.add_argument("-s", "--slim", action="store_true", help="train the slim version of DDRNet23") args, _ = parser.parse_known_args() distributed = super_gradients.is_distributed() devices = torch.cuda.device_count() if not distributed else 1 -train_params_ddr = {"max_epochs": args.max_epochs, - "lr_mode": "step", - "lr_updates": [30, 60, 90], - "lr_decay_factor": 0.1, - "initial_lr": 0.1 * devices, - "optimizer": "SGD", - "optimizer_params": {"weight_decay": 0.0001, "momentum": 0.9, "nesterov": True}, - "loss": "cross_entropy", - "train_metrics_list": [Accuracy(), Top5()], - "valid_metrics_list": [Accuracy(), Top5()], - - "metric_to_watch": "Accuracy", - "greater_metric_to_watch_is_better": True - } - -dataset_params = {"batch_size": args.batch, - "color_jitter": 0.4, - "random_erase_prob": 0.2, - "random_erase_value": 'random', - "train_interpolation": 'random', - } - - -train_transforms = [RandomResizedCropAndInterpolation(size=224, interpolation="random"), - RandomHorizontalFlip(), - ColorJitter(0.4, 0.4, 0.4), - ToTensor(), - Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), - RandomErase(0.2, "random") - ] - -trainer = Trainer(experiment_name=args.experiment_name, - multi_gpu=MultiGPUMode.DISTRIBUTED_DATA_PARALLEL if distributed else MultiGPUMode.DATA_PARALLEL, - device='cuda') - -train_loader = dataloaders.imagenet_train(dataset_params={"transforms": train_transforms}, - dataloader_params={"batch_size": args.batch}) +train_params_ddr = { + "max_epochs": args.max_epochs, + "lr_mode": "step", + "lr_updates": [30, 60, 90], + "lr_decay_factor": 0.1, + "initial_lr": 0.1 * devices, + "optimizer": "SGD", + "optimizer_params": {"weight_decay": 0.0001, "momentum": 0.9, "nesterov": True}, + "loss": "cross_entropy", + "train_metrics_list": [Accuracy(), Top5()], + "valid_metrics_list": [Accuracy(), Top5()], + "metric_to_watch": "Accuracy", + "greater_metric_to_watch_is_better": True, +} + +dataset_params = { + "batch_size": args.batch, + "color_jitter": 0.4, + "random_erase_prob": 0.2, + "random_erase_value": "random", + "train_interpolation": "random", +} + + +train_transforms = [ + RandomResizedCropAndInterpolation(size=224, interpolation="random"), + RandomHorizontalFlip(), + ColorJitter(0.4, 0.4, 0.4), + ToTensor(), + Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]), + RandomErase(0.2, "random"), +] + +trainer = Trainer( + experiment_name=args.experiment_name, multi_gpu=MultiGPUMode.DISTRIBUTED_DATA_PARALLEL if distributed else MultiGPUMode.DATA_PARALLEL, device="cuda" +) + +train_loader = dataloaders.imagenet_train(dataset_params={"transforms": train_transforms}, dataloader_params={"batch_size": args.batch}) valid_loader = dataloaders.imagenet_val() -model = models.get("ddrnet_23_slim" if args.slim else "ddrnet_23", - arch_params={"aux_head": False, "classification_mode": True, 'dropout_prob': 0.3}, - num_classes=1000) +model = models.get( + Models.DDRNET_23_SLIM if args.slim else Models.DDRNET_23, + arch_params={"aux_head": False, "classification_mode": True, "dropout_prob": 0.3}, + num_classes=1000, +) trainer.train(model=model, training_params=train_params_ddr, train_loader=train_loader, valid_loader=valid_loader) diff --git a/src/super_gradients/examples/early_stop/early_stop_example.py b/src/super_gradients/examples/early_stop/early_stop_example.py index 1e2dda7f40..e2cbb782cf 100644 --- a/src/super_gradients/examples/early_stop/early_stop_example.py +++ b/src/super_gradients/examples/early_stop/early_stop_example.py @@ -2,6 +2,7 @@ # Reaches ~94.9 Accuracy after 250 Epochs import super_gradients from super_gradients import Trainer +from super_gradients.common.object_names import Models from super_gradients.training import models, dataloaders from super_gradients.training.metrics.classification_metrics import Accuracy, Top5 from super_gradients.training.utils.early_stopping import EarlyStop @@ -13,18 +14,28 @@ early_stop_acc = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="Accuracy", mode="max", patience=3, verbose=True) early_stop_val_loss = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="LabelSmoothingCrossEntropyLoss", mode="min", patience=3, verbose=True) -train_params = {"max_epochs": 250, "lr_updates": [100, 150, 200], "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, "phase_callbacks": [early_stop_acc, early_stop_val_loss]} +train_params = { + "max_epochs": 250, + "lr_updates": [100, 150, 200], + "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, + "phase_callbacks": [early_stop_acc, early_stop_val_loss], +} # Define Model trainer = Trainer("Callback_Example") # Build Model -model = models.get("resnet18_cifar", num_classes=10) +model = models.get(Models.RESNET18_CIFAR, num_classes=10) -trainer.train(model=model, training_params=train_params, - train_loader=dataloaders.cifar10_train(), valid_loader=dataloaders.cifar10_val()) +trainer.train(model=model, training_params=train_params, train_loader=dataloaders.cifar10_train(), valid_loader=dataloaders.cifar10_val()) diff --git a/src/super_gradients/examples/loggers_examples/clearml_logger_example.py b/src/super_gradients/examples/loggers_examples/clearml_logger_example.py index 8292cbac1d..700bb6f584 100644 --- a/src/super_gradients/examples/loggers_examples/clearml_logger_example.py +++ b/src/super_gradients/examples/loggers_examples/clearml_logger_example.py @@ -1,10 +1,11 @@ +from super_gradients.common.object_names import Models from super_gradients.training import Trainer, models from super_gradients.training.metrics.classification_metrics import Accuracy, Top5 from super_gradients.training.dataloaders.dataloaders import cifar10_train, cifar10_val trainer = Trainer(experiment_name="demo-clearml-logger") -model = models.get("resnet18", num_classes=10) +model = models.get(Models.RESNET18, num_classes=10) training_params = { "max_epochs": 20, diff --git a/src/super_gradients/examples/loggers_examples/deci_platform_logger_example.py b/src/super_gradients/examples/loggers_examples/deci_platform_logger_example.py index 53c51d7f9b..988373e9ee 100644 --- a/src/super_gradients/examples/loggers_examples/deci_platform_logger_example.py +++ b/src/super_gradients/examples/loggers_examples/deci_platform_logger_example.py @@ -1,4 +1,6 @@ import os + +from super_gradients.common.object_names import Models from super_gradients.training import Trainer, models from super_gradients.training.metrics.classification_metrics import Accuracy, Top5 from super_gradients.training.dataloaders.dataloaders import cifar10_train, cifar10_val @@ -7,7 +9,7 @@ trainer = Trainer(experiment_name="demo-deci-platform-logger") -model = models.get("resnet18", num_classes=10) +model = models.get(Models.RESNET18, num_classes=10) training_params = { "max_epochs": 20, "lr_updates": [5, 10, 15], diff --git a/src/super_gradients/examples/quantization/resnet_qat_example.py b/src/super_gradients/examples/quantization/resnet_qat_example.py index 66a7e6145f..e8df8fdd57 100644 --- a/src/super_gradients/examples/quantization/resnet_qat_example.py +++ b/src/super_gradients/examples/quantization/resnet_qat_example.py @@ -4,6 +4,7 @@ import super_gradients from super_gradients import Trainer +from super_gradients.common.object_names import Models from super_gradients.training import MultiGPUMode from super_gradients.training import models as sg_models from super_gradients.training.dataloaders import imagenet_train, imagenet_val @@ -40,7 +41,7 @@ def selective_quantize(model: nn.Module): def sg_vanilla_resnet50(): - return sg_models.get("resnet50", pretrained_weights="imagenet", num_classes=1000) + return sg_models.get(Models.RESNET50, pretrained_weights="imagenet", num_classes=1000) def sg_naive_qdq_resnet50(): diff --git a/src/super_gradients/examples/regseg_transfer_learning_example/regseg_transfer_learning_example.py b/src/super_gradients/examples/regseg_transfer_learning_example/regseg_transfer_learning_example.py index 555aa8d92b..aed420ece0 100644 --- a/src/super_gradients/examples/regseg_transfer_learning_example/regseg_transfer_learning_example.py +++ b/src/super_gradients/examples/regseg_transfer_learning_example/regseg_transfer_learning_example.py @@ -1,20 +1,31 @@ +from super_gradients.common.object_names import Models from super_gradients.training import models, dataloaders from super_gradients.training.sg_trainer import Trainer from super_gradients.training.metrics import BinaryIOU -from super_gradients.training.transforms.transforms import SegResize, SegRandomFlip, SegRandomRescale, SegCropImageAndMask, \ - SegPadShortToCropSize, SegColorJitter +from super_gradients.training.transforms.transforms import ( + SegResize, + SegRandomFlip, + SegRandomRescale, + SegCropImageAndMask, + SegPadShortToCropSize, + SegColorJitter, +) from super_gradients.training.utils.callbacks import BinarySegmentationVisualizationCallback, Phase # DEFINE DATA TRANSFORMATIONS dl_train = dataloaders.supervisely_persons_train( - dataset_params={"transforms": [SegColorJitter(brightness=0.5, contrast=0.5, saturation=0.5), - SegRandomFlip(), - SegRandomRescale(scales=[0.25, 1.]), - SegPadShortToCropSize([320, 480]), - SegCropImageAndMask(crop_size=[320, 480], - mode="random")]}) + dataset_params={ + "transforms": [ + SegColorJitter(brightness=0.5, contrast=0.5, saturation=0.5), + SegRandomFlip(), + SegRandomRescale(scales=[0.25, 1.0]), + SegPadShortToCropSize([320, 480]), + SegCropImageAndMask(crop_size=[320, 480], mode="random"), + ] + } +) dl_val = dataloaders.supervisely_persons_val(dataset_params={"transforms": [SegResize(h=480, w=320)]}) @@ -23,35 +34,27 @@ # THIS IS WHERE THE MAGIC HAPPENS- SINCE TRAINER'S CLASSES ATTRIBUTE WAS SET TO BE DIFFERENT FROM CITYSCAPES'S, AFTER # LOADING THE PRETRAINED REGSET, IT WILL CALL IT'S REPLACE_HEAD METHOD AND CHANGE IT'S SEGMENTATION HEAD LAYER ACCORDING # TO OUR BINARY SEGMENTATION DATASET -model = models.get("regseg48", pretrained_weights="cityscapes", num_classes=1) +model = models.get(Models.RREGSEG48, pretrained_weights="cityscapes", num_classes=1) # DEFINE TRAINING PARAMS. SEE DOCS FOR THE FULL LIST. -train_params = {"max_epochs": 50, - "lr_mode": "cosine", - "initial_lr": 0.0064, # for batch_size=16 - "optimizer_params": {"momentum": 0.843, - "weight_decay": 0.00036, - "nesterov": True}, - - "cosine_final_lr_ratio": 0.1, - "multiply_head_lr": 10, - "optimizer": "SGD", - "loss": "bce_dice_loss", - "ema": True, - "zero_weight_decay_on_bias_and_bn": True, - "average_best_models": True, - "mixed_precision": False, - "metric_to_watch": "mean_IOU", - "greater_metric_to_watch_is_better": True, - "train_metrics_list": [BinaryIOU()], - "valid_metrics_list": [BinaryIOU()], - - "phase_callbacks": [BinarySegmentationVisualizationCallback(phase=Phase.VALIDATION_BATCH_END, - freq=1, - last_img_idx_in_batch=4)], - } - -trainer.train(model=model, - training_params=train_params, - train_loader=dl_train, - valid_loader=dl_val) +train_params = { + "max_epochs": 50, + "lr_mode": "cosine", + "initial_lr": 0.0064, # for batch_size=16 + "optimizer_params": {"momentum": 0.843, "weight_decay": 0.00036, "nesterov": True}, + "cosine_final_lr_ratio": 0.1, + "multiply_head_lr": 10, + "optimizer": "SGD", + "loss": "bce_dice_loss", + "ema": True, + "zero_weight_decay_on_bias_and_bn": True, + "average_best_models": True, + "mixed_precision": False, + "metric_to_watch": "mean_IOU", + "greater_metric_to_watch_is_better": True, + "train_metrics_list": [BinaryIOU()], + "valid_metrics_list": [BinaryIOU()], + "phase_callbacks": [BinarySegmentationVisualizationCallback(phase=Phase.VALIDATION_BATCH_END, freq=1, last_img_idx_in_batch=4)], +} + +trainer.train(model=model, training_params=train_params, train_loader=dl_train, valid_loader=dl_val) diff --git a/tests/end_to_end_tests/trainer_test.py b/tests/end_to_end_tests/trainer_test.py index 21445bd3aa..67b7ab817e 100644 --- a/tests/end_to_end_tests/trainer_test.py +++ b/tests/end_to_end_tests/trainer_test.py @@ -1,6 +1,7 @@ import shutil import unittest +from super_gradients.common.object_names import Models from super_gradients.training import models import super_gradients @@ -16,66 +17,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) + model = models.get(Models.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/ema_train_integration_test.py b/tests/integration_tests/ema_train_integration_test.py index 233b9b5221..f7722f326d 100644 --- a/tests/integration_tests/ema_train_integration_test.py +++ b/tests/integration_tests/ema_train_integration_test.py @@ -1,3 +1,4 @@ +from super_gradients.common.object_names import Models from super_gradients.training import MultiGPUMode, models from super_gradients.training import Trainer from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader @@ -22,7 +23,7 @@ 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.model = models.get("resnet18_cifar", arch_params={"num_classes": 5}) + self.model = models.get(Models.RESNET18_CIFAR, arch_params={"num_classes": 5}) @classmethod def tearDownClass(cls) -> None: diff --git a/tests/integration_tests/lr_test.py b/tests/integration_tests/lr_test.py index cee92f0c10..82bdaec94d 100644 --- a/tests/integration_tests/lr_test.py +++ b/tests/integration_tests/lr_test.py @@ -2,6 +2,7 @@ import unittest import os +from super_gradients.common.object_names import Models from super_gradients.training import models from super_gradients import Trainer @@ -13,25 +14,28 @@ 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) + model = models.get(Models.RESNET18_CIFAR, num_classes=5) return trainer, model def test_function_lr(self): @@ -42,26 +46,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/integration_tests/pretrained_models_test.py b/tests/integration_tests/pretrained_models_test.py index 94e544bfa1..ced39aea87 100644 --- a/tests/integration_tests/pretrained_models_test.py +++ b/tests/integration_tests/pretrained_models_test.py @@ -1,5 +1,6 @@ import unittest +from super_gradients.common.object_names import Models from super_gradients.training import MultiGPUMode from super_gradients.training import Trainer from super_gradients.training.dataloaders import imagenet_val, imagenet_vit_base_val @@ -35,19 +36,19 @@ def setUp(self) -> None: self.imagenet_pretrained_arch_params = { "resnet": {}, "regnet": {}, - "repvgg_a0": {"build_residual_branches": True}, - "efficientnet_b0": {}, + Models.REPVGG_A0: {"build_residual_branches": True}, + Models.EFFICIENTNET_B0: {}, "mobilenet": {}, - "vit_base": {"image_size": (224, 224), "patch_size": (16, 16)}, + Models.VIT_BASE: {"image_size": (224, 224), "patch_size": (16, 16)}, } self.imagenet_pretrained_trainsfer_learning_arch_params = { "resnet": {}, "regnet": {}, - "repvgg_a0": {"build_residual_branches": True}, - "efficientnet_b0": {}, + Models.REPVGG_A0: {"build_residual_branches": True}, + Models.EFFICIENTNET_B0: {}, "mobilenet": {}, - "vit_base": {"image_size": (224, 224), "patch_size": (16, 16)}, + Models.VIT_BASE: {"image_size": (224, 224), "patch_size": (16, 16)}, } self.imagenet_pretrained_ckpt_params = {"pretrained_weights": "imagenet"} @@ -55,21 +56,21 @@ def setUp(self) -> None: self.imagenet21k_pretrained_ckpt_params = {"pretrained_weights": "imagenet21k"} self.imagenet_pretrained_accuracies = { - "resnet50": 0.8191, - "resnet34": 0.7413, - "resnet18": 0.706, - "repvgg_a0": 0.7205, - "regnetY800": 0.7707, - "regnetY600": 0.7618, - "regnetY400": 0.7474, - "regnetY200": 0.7088, - "efficientnet_b0": 0.7762, - "mobilenet_v3_large": 0.7452, - "mobilenet_v3_small": 0.6745, - "mobilenet_v2": 0.7308, - "vit_base": 0.8415, - "vit_large": 0.8564, - "beit_base_patch16_224": 0.85, + Models.RESNET50: 0.8191, + Models.RESNET34: 0.7413, + Models.RESNET18: 0.706, + Models.REPVGG_A0: 0.7205, + Models.REGNETY800: 0.7707, + Models.REGNETY600: 0.7618, + Models.REGNETY400: 0.7474, + Models.REGNETY200: 0.7088, + Models.EFFICIENTNET_B0: 0.7762, + Models.MOBILENET_V3_LARGE: 0.7452, + Models.MOBILENET_V3_SMALL: 0.6745, + Models.MOBILENET_V2: 0.7308, + Models.VIT_BASE: 0.8415, + Models.VIT_LARGE: 0.8564, + Models.BEIT_BASE_PATCH16_224: 0.85, } self.imagenet_dataset = imagenet_val(dataloader_params={"batch_size": 128}) @@ -91,7 +92,7 @@ def setUp(self) -> None: "greater_metric_to_watch_is_better": True, } self.coco_pretrained_ckpt_params = {"pretrained_weights": "coco"} - self.coco_pretrained_arch_params = {"ssd_lite_mobilenet_v2": {"num_classes": 80}, "coco_ssd_mobilenet_v1": {"num_classes": 80}} + self.coco_pretrained_arch_params = {Models.SSD_LITE_MOBILENET_V2: {"num_classes": 80}, "coco_ssd_mobilenet_v1": {"num_classes": 80}} self.coco_pretrained_ckpt_params = {"pretrained_weights": "coco"} self.coco_dataset = { @@ -102,13 +103,13 @@ def setUp(self) -> None: } self.coco_pretrained_maps = { - "ssd_lite_mobilenet_v2": 0.2041, + Models.SSD_LITE_MOBILENET_V2: 0.2041, "coco_ssd_mobilenet_v1": 0.243, - "yolox_s": 0.4047, - "yolox_m": 0.4640, - "yolox_l": 0.4925, - "yolox_n": 0.2677, - "yolox_t": 0.3718, + Models.YOLOX_S: 0.4047, + Models.YOLOX_M: 0.4640, + Models.YOLOX_L: 0.4925, + Models.YOLOX_N: 0.2677, + Models.YOLOX_T: 0.3718, } self.transfer_detection_dataset = detection_test_dataloader() @@ -151,27 +152,27 @@ def setUp(self) -> None: self.coco_segmentation_subclass_pretrained_mious = {"shelfnet34_lw": 0.651} self.coco_segmentation_dataset = coco_segmentation_val() - self.cityscapes_pretrained_models = ["ddrnet_23", "ddrnet_23_slim", "stdc1_seg50", "regseg48"] + self.cityscapes_pretrained_models = [Models.DDRNET_23, Models.DDRNET_23_SLIM, Models.STDC1_SEG50, Models.REGSEG48] self.cityscapes_pretrained_arch_params = { - "ddrnet_23": {"aux_head": True}, - "regseg48": {}, + Models.DDRNET_23: {"aux_head": True}, + Models.REGSEG48: {}, "stdc": {"use_aux_heads": True, "aux_head": True}, "pplite_seg": {"use_aux_heads": True}, } self.cityscapes_pretrained_ckpt_params = {"pretrained_weights": "cityscapes"} self.cityscapes_pretrained_mious = { - "ddrnet_23": 0.8026, - "ddrnet_23_slim": 0.7801, - "stdc1_seg50": 0.7511, - "stdc1_seg75": 0.7687, - "stdc2_seg50": 0.7644, - "stdc2_seg75": 0.7893, - "regseg48": 0.7815, - "pp_lite_t_seg50": 0.7492, - "pp_lite_t_seg75": 0.7756, - "pp_lite_b_seg50": 0.7648, - "pp_lite_b_seg75": 0.7852, + Models.DDRNET_23: 0.8026, + Models.DDRNET_23_SLIM: 0.7801, + Models.STDC1_SEG50: 0.7511, + Models.STDC1_SEG75: 0.7687, + Models.STDC2_SEG50: 0.7644, + Models.STDC2_SEG75: 0.7893, + Models.REGSEG48: 0.7815, + Models.PP_LITE_T_SEG50: 0.7492, + Models.PP_LITE_T_SEG75: 0.7756, + Models.PP_LITE_B_SEG50: 0.7648, + Models.PP_LITE_B_SEG75: 0.7852, } self.cityscapes_dataset = cityscapes_val() @@ -229,13 +230,13 @@ def setUp(self) -> None: def test_pretrained_resnet50_imagenet(self): trainer = Trainer("imagenet_pretrained_resnet50", multi_gpu=MultiGPUMode.OFF) - model = models.get("resnet50", arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.RESNET50, arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["resnet50"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.RESNET50], delta=0.001) def test_transfer_learning_resnet50_imagenet(self): trainer = Trainer("imagenet_pretrained_resnet50_transfer_learning", multi_gpu=MultiGPUMode.OFF) - model = models.get("resnet50", arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) + model = models.get(Models.RESNET50, arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -246,13 +247,13 @@ def test_transfer_learning_resnet50_imagenet(self): def test_pretrained_resnet34_imagenet(self): trainer = Trainer("imagenet_pretrained_resnet34", multi_gpu=MultiGPUMode.OFF) - model = models.get("resnet34", arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.RESNET34, arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["resnet34"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.RESNET34], delta=0.001) def test_transfer_learning_resnet34_imagenet(self): trainer = Trainer("imagenet_pretrained_resnet34_transfer_learning", multi_gpu=MultiGPUMode.OFF) - model = models.get("resnet34", arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) + model = models.get(Models.RESNET34, arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -263,13 +264,13 @@ def test_transfer_learning_resnet34_imagenet(self): def test_pretrained_resnet18_imagenet(self): trainer = Trainer("imagenet_pretrained_resnet18", multi_gpu=MultiGPUMode.OFF) - model = models.get("resnet18", arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.RESNET18, arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["resnet18"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.RESNET18], delta=0.001) def test_transfer_learning_resnet18_imagenet(self): trainer = Trainer("imagenet_pretrained_resnet18_transfer_learning", multi_gpu=MultiGPUMode.OFF) - model = models.get("resnet18", arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) + model = models.get(Models.RESNET18, arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -280,13 +281,13 @@ def test_transfer_learning_resnet18_imagenet(self): def test_pretrained_regnetY800_imagenet(self): trainer = Trainer("imagenet_pretrained_regnetY800", multi_gpu=MultiGPUMode.OFF) - model = models.get("regnetY800", arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.REGNETY800, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY800"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.REGNETY800], delta=0.001) def test_transfer_learning_regnetY800_imagenet(self): trainer = Trainer("imagenet_pretrained_regnetY800_transfer_learning", multi_gpu=MultiGPUMode.OFF) - model = models.get("regnetY800", arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) + model = models.get(Models.REGNETY800, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -297,13 +298,13 @@ def test_transfer_learning_regnetY800_imagenet(self): def test_pretrained_regnetY600_imagenet(self): trainer = Trainer("imagenet_pretrained_regnetY600", multi_gpu=MultiGPUMode.OFF) - model = models.get("regnetY600", arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.REGNETY600, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY600"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.REGNETY600], delta=0.001) def test_transfer_learning_regnetY600_imagenet(self): trainer = Trainer("imagenet_pretrained_regnetY600_transfer_learning", multi_gpu=MultiGPUMode.OFF) - model = models.get("regnetY600", arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) + model = models.get(Models.REGNETY600, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -314,13 +315,13 @@ def test_transfer_learning_regnetY600_imagenet(self): def test_pretrained_regnetY400_imagenet(self): trainer = Trainer("imagenet_pretrained_regnetY400", multi_gpu=MultiGPUMode.OFF) - model = models.get("regnetY400", arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.REGNETY400, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY400"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.REGNETY400], delta=0.001) def test_transfer_learning_regnetY400_imagenet(self): trainer = Trainer("imagenet_pretrained_regnetY400_transfer_learning", multi_gpu=MultiGPUMode.OFF) - model = models.get("regnetY400", arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) + model = models.get(Models.REGNETY400, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -331,13 +332,13 @@ def test_transfer_learning_regnetY400_imagenet(self): def test_pretrained_regnetY200_imagenet(self): trainer = Trainer("imagenet_pretrained_regnetY200", multi_gpu=MultiGPUMode.OFF) - model = models.get("regnetY200", arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.REGNETY200, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY200"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.REGNETY200], delta=0.001) def test_transfer_learning_regnetY200_imagenet(self): trainer = Trainer("imagenet_pretrained_regnetY200_transfer_learning", multi_gpu=MultiGPUMode.OFF) - model = models.get("regnetY200", arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) + model = models.get(Models.REGNETY200, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -348,13 +349,15 @@ def test_transfer_learning_regnetY200_imagenet(self): def test_pretrained_repvgg_a0_imagenet(self): trainer = Trainer("imagenet_pretrained_repvgg_a0", multi_gpu=MultiGPUMode.OFF) - model = models.get("repvgg_a0", arch_params=self.imagenet_pretrained_arch_params["repvgg_a0"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.REPVGG_A0, arch_params=self.imagenet_pretrained_arch_params[Models.REPVGG_A0], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["repvgg_a0"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.REPVGG_A0], delta=0.001) def test_transfer_learning_repvgg_a0_imagenet(self): trainer = Trainer("imagenet_pretrained_repvgg_a0_transfer_learning", multi_gpu=MultiGPUMode.OFF) - model = models.get("repvgg_a0", arch_params=self.imagenet_pretrained_arch_params["repvgg_a0"], **self.imagenet_pretrained_ckpt_params, num_classes=5) + model = models.get( + Models.REPVGG_A0, arch_params=self.imagenet_pretrained_arch_params[Models.REPVGG_A0], **self.imagenet_pretrained_ckpt_params, num_classes=5 + ) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -364,7 +367,7 @@ def test_transfer_learning_repvgg_a0_imagenet(self): def test_pretrained_regseg48_cityscapes(self): trainer = Trainer("cityscapes_pretrained_regseg48", multi_gpu=MultiGPUMode.OFF) - model = models.get("regseg48", arch_params=self.cityscapes_pretrained_arch_params["regseg48"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.REGSEG48, arch_params=self.cityscapes_pretrained_arch_params[Models.REGSEG48], **self.cityscapes_pretrained_ckpt_params) res = ( trainer.test( model=model, test_loader=self.cityscapes_dataset, test_metrics_list=[IoU(num_classes=20, ignore_index=19)], metrics_progress_verbose=True @@ -372,11 +375,11 @@ def test_pretrained_regseg48_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["regseg48"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.REGSEG48], delta=0.001) def test_transfer_learning_regseg48_cityscapes(self): trainer = Trainer("regseg48_cityscapes_transfer_learning", multi_gpu=MultiGPUMode.OFF) - model = models.get("regseg48", arch_params=self.cityscapes_pretrained_arch_params["regseg48"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.REGSEG48, arch_params=self.cityscapes_pretrained_arch_params[Models.REGSEG48], **self.cityscapes_pretrained_ckpt_params) trainer.train( model=model, train_loader=self.transfer_segmentation_dataset, @@ -386,7 +389,7 @@ def test_transfer_learning_regseg48_cityscapes(self): def test_pretrained_ddrnet23_cityscapes(self): trainer = Trainer("cityscapes_pretrained_ddrnet23", multi_gpu=MultiGPUMode.OFF) - model = models.get("ddrnet_23", arch_params=self.cityscapes_pretrained_arch_params["ddrnet_23"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.DDRNET_23, arch_params=self.cityscapes_pretrained_arch_params[Models.DDRNET_23], **self.cityscapes_pretrained_ckpt_params) res = ( trainer.test( model=model, test_loader=self.cityscapes_dataset, test_metrics_list=[IoU(num_classes=20, ignore_index=19)], metrics_progress_verbose=True @@ -394,11 +397,13 @@ def test_pretrained_ddrnet23_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["ddrnet_23"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.DDRNET_23], delta=0.001) def test_pretrained_ddrnet23_slim_cityscapes(self): trainer = Trainer("cityscapes_pretrained_ddrnet23_slim", multi_gpu=MultiGPUMode.OFF) - model = models.get("ddrnet_23_slim", arch_params=self.cityscapes_pretrained_arch_params["ddrnet_23"], **self.cityscapes_pretrained_ckpt_params) + model = models.get( + Models.DDRNET_23_SLIM, arch_params=self.cityscapes_pretrained_arch_params[Models.DDRNET_23], **self.cityscapes_pretrained_ckpt_params + ) res = ( trainer.test( model=model, test_loader=self.cityscapes_dataset, test_metrics_list=[IoU(num_classes=20, ignore_index=19)], metrics_progress_verbose=True @@ -406,11 +411,11 @@ def test_pretrained_ddrnet23_slim_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["ddrnet_23_slim"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.DDRNET_23_SLIM], delta=0.001) def test_transfer_learning_ddrnet23_cityscapes(self): trainer = Trainer("cityscapes_pretrained_ddrnet23_transfer_learning", multi_gpu=MultiGPUMode.OFF) - model = models.get("ddrnet_23", arch_params=self.cityscapes_pretrained_arch_params["ddrnet_23"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.DDRNET_23, arch_params=self.cityscapes_pretrained_arch_params[Models.DDRNET_23], **self.cityscapes_pretrained_ckpt_params) trainer.train( model=model, training_params=self.ddrnet_transfer_segmentation_train_params, @@ -420,7 +425,9 @@ def test_transfer_learning_ddrnet23_cityscapes(self): def test_transfer_learning_ddrnet23_slim_cityscapes(self): trainer = Trainer("cityscapes_pretrained_ddrnet23_slim_transfer_learning", multi_gpu=MultiGPUMode.OFF) - model = models.get("ddrnet_23_slim", arch_params=self.cityscapes_pretrained_arch_params["ddrnet_23"], **self.cityscapes_pretrained_ckpt_params) + model = models.get( + Models.DDRNET_23_SLIM, arch_params=self.cityscapes_pretrained_arch_params[Models.DDRNET_23], **self.cityscapes_pretrained_ckpt_params + ) trainer.train( model=model, training_params=self.ddrnet_transfer_segmentation_train_params, @@ -441,15 +448,20 @@ def test_pretrained_coco_segmentation_subclass_pretrained_shelfnet34_lw(self): def test_pretrained_efficientnet_b0_imagenet(self): trainer = Trainer("imagenet_pretrained_efficientnet_b0", multi_gpu=MultiGPUMode.OFF) - model = models.get("efficientnet_b0", arch_params=self.imagenet_pretrained_arch_params["efficientnet_b0"], **self.imagenet_pretrained_ckpt_params) + model = models.get( + Models.EFFICIENTNET_B0, arch_params=self.imagenet_pretrained_arch_params[Models.EFFICIENTNET_B0], **self.imagenet_pretrained_ckpt_params + ) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["efficientnet_b0"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.EFFICIENTNET_B0], delta=0.001) def test_transfer_learning_efficientnet_b0_imagenet(self): trainer = Trainer("imagenet_pretrained_efficientnet_b0_transfer_learning", multi_gpu=MultiGPUMode.OFF) model = models.get( - "efficientnet_b0", arch_params=self.imagenet_pretrained_arch_params["efficientnet_b0"], **self.imagenet_pretrained_ckpt_params, num_classes=5 + Models.EFFICIENTNET_B0, + arch_params=self.imagenet_pretrained_arch_params[Models.EFFICIENTNET_B0], + **self.imagenet_pretrained_ckpt_params, + num_classes=5, ) trainer.train( model=model, @@ -460,7 +472,9 @@ def test_transfer_learning_efficientnet_b0_imagenet(self): def test_pretrained_ssd_lite_mobilenet_v2_coco(self): trainer = Trainer("coco_ssd_lite_mobilenet_v2", multi_gpu=MultiGPUMode.OFF) - model = models.get("ssd_lite_mobilenet_v2", arch_params=self.coco_pretrained_arch_params["ssd_lite_mobilenet_v2"], **self.coco_pretrained_ckpt_params) + model = models.get( + Models.SSD_LITE_MOBILENET_V2, arch_params=self.coco_pretrained_arch_params[Models.SSD_LITE_MOBILENET_V2], **self.coco_pretrained_ckpt_params + ) ssd_post_prediction_callback = SSDPostPredictCallback() res = trainer.test( model=model, @@ -468,13 +482,13 @@ def test_pretrained_ssd_lite_mobilenet_v2_coco(self): test_metrics_list=[DetectionMetrics(post_prediction_callback=ssd_post_prediction_callback, num_cls=80)], metrics_progress_verbose=True, )[2] - self.assertAlmostEqual(res, self.coco_pretrained_maps["ssd_lite_mobilenet_v2"], delta=0.001) + self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.SSD_LITE_MOBILENET_V2], delta=0.001) def test_transfer_learning_ssd_lite_mobilenet_v2_coco(self): trainer = Trainer("coco_ssd_lite_mobilenet_v2_transfer_learning", multi_gpu=MultiGPUMode.OFF) - transfer_arch_params = self.coco_pretrained_arch_params["ssd_lite_mobilenet_v2"].copy() + transfer_arch_params = self.coco_pretrained_arch_params[Models.SSD_LITE_MOBILENET_V2].copy() transfer_arch_params["num_classes"] = 5 - model = models.get("ssd_lite_mobilenet_v2", arch_params=transfer_arch_params, **self.coco_pretrained_ckpt_params) + model = models.get(Models.SSD_LITE_MOBILENET_V2, arch_params=transfer_arch_params, **self.coco_pretrained_ckpt_params) trainer.train( model=model, training_params=self.transfer_detection_train_params_ssd, @@ -484,7 +498,7 @@ def test_transfer_learning_ssd_lite_mobilenet_v2_coco(self): def test_pretrained_ssd_mobilenet_v1_coco(self): trainer = Trainer("coco_ssd_mobilenet_v1", multi_gpu=MultiGPUMode.OFF) - model = models.get("ssd_mobilenet_v1", arch_params=self.coco_pretrained_arch_params["coco_ssd_mobilenet_v1"], **self.coco_pretrained_ckpt_params) + model = models.get(Models.SSD_MOBILENET_V1, arch_params=self.coco_pretrained_arch_params["coco_ssd_mobilenet_v1"], **self.coco_pretrained_ckpt_params) ssd_post_prediction_callback = SSDPostPredictCallback() res = trainer.test( model=model, @@ -495,60 +509,60 @@ def test_pretrained_ssd_mobilenet_v1_coco(self): self.assertAlmostEqual(res, self.coco_pretrained_maps["coco_ssd_mobilenet_v1"], delta=0.001) def test_pretrained_yolox_s_coco(self): - trainer = Trainer("yolox_s", multi_gpu=MultiGPUMode.OFF) + trainer = Trainer(Models.YOLOX_S, multi_gpu=MultiGPUMode.OFF) - model = models.get("yolox_s", **self.coco_pretrained_ckpt_params) + model = models.get(Models.YOLOX_S, **self.coco_pretrained_ckpt_params) res = trainer.test( model=model, test_loader=self.coco_dataset["yolox"], test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=80, normalize_targets=True)], )[2] - self.assertAlmostEqual(res, self.coco_pretrained_maps["yolox_s"], delta=0.001) + self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.YOLOX_S], delta=0.001) def test_pretrained_yolox_m_coco(self): - trainer = Trainer("yolox_m", multi_gpu=MultiGPUMode.OFF) - model = models.get("yolox_m", **self.coco_pretrained_ckpt_params) + trainer = Trainer(Models.YOLOX_M, multi_gpu=MultiGPUMode.OFF) + model = models.get(Models.YOLOX_M, **self.coco_pretrained_ckpt_params) res = trainer.test( model=model, test_loader=self.coco_dataset["yolox"], test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=80, normalize_targets=True)], )[2] - self.assertAlmostEqual(res, self.coco_pretrained_maps["yolox_m"], delta=0.001) + self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.YOLOX_M], delta=0.001) def test_pretrained_yolox_l_coco(self): - trainer = Trainer("yolox_l", multi_gpu=MultiGPUMode.OFF) - model = models.get("yolox_l", **self.coco_pretrained_ckpt_params) + trainer = Trainer(Models.YOLOX_L, multi_gpu=MultiGPUMode.OFF) + model = models.get(Models.YOLOX_L, **self.coco_pretrained_ckpt_params) res = trainer.test( model=model, test_loader=self.coco_dataset["yolox"], test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=80, normalize_targets=True)], )[2] - self.assertAlmostEqual(res, self.coco_pretrained_maps["yolox_l"], delta=0.001) + self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.YOLOX_L], delta=0.001) def test_pretrained_yolox_n_coco(self): - trainer = Trainer("yolox_n", multi_gpu=MultiGPUMode.OFF) + trainer = Trainer(Models.YOLOX_N, multi_gpu=MultiGPUMode.OFF) - model = models.get("yolox_n", **self.coco_pretrained_ckpt_params) + model = models.get(Models.YOLOX_N, **self.coco_pretrained_ckpt_params) res = trainer.test( model=model, test_loader=self.coco_dataset["yolox"], test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=80, normalize_targets=True)], )[2] - self.assertAlmostEqual(res, self.coco_pretrained_maps["yolox_n"], delta=0.001) + self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.YOLOX_N], delta=0.001) def test_pretrained_yolox_t_coco(self): - trainer = Trainer("yolox_t", multi_gpu=MultiGPUMode.OFF) - model = models.get("yolox_t", **self.coco_pretrained_ckpt_params) + trainer = Trainer(Models.YOLOX_T, multi_gpu=MultiGPUMode.OFF) + model = models.get(Models.YOLOX_T, **self.coco_pretrained_ckpt_params) res = trainer.test( model=model, test_loader=self.coco_dataset["yolox"], test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=80, normalize_targets=True)], )[2] - self.assertAlmostEqual(res, self.coco_pretrained_maps["yolox_t"], delta=0.001) + self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.YOLOX_T], delta=0.001) def test_transfer_learning_yolox_n_coco(self): trainer = Trainer("test_transfer_learning_yolox_n_coco", multi_gpu=MultiGPUMode.OFF) - model = models.get("yolox_n", **self.coco_pretrained_ckpt_params, num_classes=5) + model = models.get(Models.YOLOX_N, **self.coco_pretrained_ckpt_params, num_classes=5) trainer.train( model=model, training_params=self.transfer_detection_train_params_yolox, @@ -560,7 +574,7 @@ def test_transfer_learning_mobilenet_v3_large_imagenet(self): trainer = Trainer("imagenet_pretrained_mobilenet_v3_large_transfer_learning", multi_gpu=MultiGPUMode.OFF) model = models.get( - "mobilenet_v3_large", arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params, num_classes=5 + Models.MOBILENET_V3_LARGE, arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params, num_classes=5 ) trainer.train( model=model, @@ -572,15 +586,15 @@ def test_transfer_learning_mobilenet_v3_large_imagenet(self): def test_pretrained_mobilenet_v3_large_imagenet(self): trainer = Trainer("imagenet_mobilenet_v3_large", multi_gpu=MultiGPUMode.OFF) - model = models.get("mobilenet_v3_large", arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.MOBILENET_V3_LARGE, arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["mobilenet_v3_large"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.MOBILENET_V3_LARGE], delta=0.001) def test_transfer_learning_mobilenet_v3_small_imagenet(self): trainer = Trainer("imagenet_pretrained_mobilenet_v3_small_transfer_learning", multi_gpu=MultiGPUMode.OFF) model = models.get( - "mobilenet_v3_small", arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params, num_classes=5 + Models.MOBILENET_V3_SMALL, arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params, num_classes=5 ) trainer.train( model=model, @@ -592,14 +606,16 @@ def test_transfer_learning_mobilenet_v3_small_imagenet(self): def test_pretrained_mobilenet_v3_small_imagenet(self): trainer = Trainer("imagenet_mobilenet_v3_small", multi_gpu=MultiGPUMode.OFF) - model = models.get("mobilenet_v3_small", arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.MOBILENET_V3_SMALL, arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["mobilenet_v3_small"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.MOBILENET_V3_SMALL], delta=0.001) def test_transfer_learning_mobilenet_v2_imagenet(self): trainer = Trainer("imagenet_pretrained_mobilenet_v2_transfer_learning", multi_gpu=MultiGPUMode.OFF) - model = models.get("mobilenet_v2", arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) + model = models.get( + Models.MOBILENET_V2, arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params, num_classes=5 + ) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -610,13 +626,13 @@ def test_transfer_learning_mobilenet_v2_imagenet(self): def test_pretrained_mobilenet_v2_imagenet(self): trainer = Trainer("imagenet_mobilenet_v2", multi_gpu=MultiGPUMode.OFF) - model = models.get("mobilenet_v2", arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.MOBILENET_V2, arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["mobilenet_v2"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.MOBILENET_V2], delta=0.001) def test_pretrained_stdc1_seg50_cityscapes(self): trainer = Trainer("cityscapes_pretrained_stdc1_seg50", multi_gpu=MultiGPUMode.OFF) - model = models.get("stdc1_seg50", arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.STDC1_SEG50, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params) res = ( trainer.test( model=model, @@ -627,11 +643,13 @@ def test_pretrained_stdc1_seg50_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["stdc1_seg50"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.STDC1_SEG50], delta=0.001) def test_transfer_learning_stdc1_seg50_cityscapes(self): trainer = Trainer("cityscapes_pretrained_stdc1_seg50_transfer_learning", multi_gpu=MultiGPUMode.OFF) - model = models.get("stdc1_seg50", arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params, num_classes=5) + model = models.get( + Models.STDC1_SEG50, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params, num_classes=5 + ) trainer.train( model=model, training_params=self.stdc_transfer_segmentation_train_params, @@ -641,7 +659,7 @@ def test_transfer_learning_stdc1_seg50_cityscapes(self): def test_pretrained_stdc1_seg75_cityscapes(self): trainer = Trainer("cityscapes_pretrained_stdc1_seg75", multi_gpu=MultiGPUMode.OFF) - model = models.get("stdc1_seg75", arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.STDC1_SEG75, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params) res = ( trainer.test( model=model, @@ -652,11 +670,13 @@ def test_pretrained_stdc1_seg75_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["stdc1_seg75"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.STDC1_SEG75], delta=0.001) def test_transfer_learning_stdc1_seg75_cityscapes(self): trainer = Trainer("cityscapes_pretrained_stdc1_seg75_transfer_learning", multi_gpu=MultiGPUMode.OFF) - model = models.get("stdc1_seg75", arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params, num_classes=5) + model = models.get( + Models.STDC1_SEG75, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params, num_classes=5 + ) trainer.train( model=model, training_params=self.stdc_transfer_segmentation_train_params, @@ -666,7 +686,7 @@ def test_transfer_learning_stdc1_seg75_cityscapes(self): def test_pretrained_stdc2_seg50_cityscapes(self): trainer = Trainer("cityscapes_pretrained_stdc2_seg50", multi_gpu=MultiGPUMode.OFF) - model = models.get("stdc2_seg50", arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.STDC2_SEG50, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params) res = ( trainer.test( model=model, @@ -677,11 +697,13 @@ def test_pretrained_stdc2_seg50_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["stdc2_seg50"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.STDC2_SEG50], delta=0.001) def test_transfer_learning_stdc2_seg50_cityscapes(self): trainer = Trainer("cityscapes_pretrained_stdc2_seg50_transfer_learning", multi_gpu=MultiGPUMode.OFF) - model = models.get("stdc2_seg50", arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params, num_classes=5) + model = models.get( + Models.STDC2_SEG50, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params, num_classes=5 + ) trainer.train( model=model, training_params=self.stdc_transfer_segmentation_train_params, @@ -691,7 +713,7 @@ def test_transfer_learning_stdc2_seg50_cityscapes(self): def test_pretrained_stdc2_seg75_cityscapes(self): trainer = Trainer("cityscapes_pretrained_stdc2_seg75", multi_gpu=MultiGPUMode.OFF) - model = models.get("stdc2_seg75", arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.STDC2_SEG75, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params) res = ( trainer.test( model=model, @@ -702,11 +724,13 @@ def test_pretrained_stdc2_seg75_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["stdc2_seg75"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.STDC2_SEG75], delta=0.001) def test_transfer_learning_stdc2_seg75_cityscapes(self): trainer = Trainer("cityscapes_pretrained_stdc2_seg75_transfer_learning", multi_gpu=MultiGPUMode.OFF) - model = models.get("stdc2_seg75", arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params, num_classes=5) + model = models.get( + Models.STDC2_SEG75, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params, num_classes=5 + ) trainer.train( model=model, training_params=self.stdc_transfer_segmentation_train_params, @@ -717,7 +741,9 @@ def test_transfer_learning_stdc2_seg75_cityscapes(self): def test_transfer_learning_vit_base_imagenet21k(self): trainer = Trainer("imagenet21k_pretrained_vit_base", multi_gpu=MultiGPUMode.OFF) - model = models.get("vit_base", arch_params=self.imagenet_pretrained_arch_params["vit_base"], **self.imagenet21k_pretrained_ckpt_params, num_classes=5) + model = models.get( + Models.VIT_BASE, arch_params=self.imagenet_pretrained_arch_params[Models.VIT_BASE], **self.imagenet21k_pretrained_ckpt_params, num_classes=5 + ) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -728,7 +754,9 @@ def test_transfer_learning_vit_base_imagenet21k(self): def test_transfer_learning_vit_large_imagenet21k(self): trainer = Trainer("imagenet21k_pretrained_vit_large", multi_gpu=MultiGPUMode.OFF) - model = models.get("vit_large", arch_params=self.imagenet_pretrained_arch_params["vit_base"], **self.imagenet21k_pretrained_ckpt_params, num_classes=5) + model = models.get( + Models.VIT_LARGE, arch_params=self.imagenet_pretrained_arch_params[Models.VIT_BASE], **self.imagenet21k_pretrained_ckpt_params, num_classes=5 + ) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -738,39 +766,44 @@ def test_transfer_learning_vit_large_imagenet21k(self): def test_pretrained_vit_base_imagenet(self): trainer = Trainer("imagenet_pretrained_vit_base", multi_gpu=MultiGPUMode.OFF) - model = models.get("vit_base", arch_params=self.imagenet_pretrained_arch_params["vit_base"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.VIT_BASE, arch_params=self.imagenet_pretrained_arch_params[Models.VIT_BASE], **self.imagenet_pretrained_ckpt_params) res = ( trainer.test(model=model, test_loader=self.imagenet_dataset_05_mean_std, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0] .cpu() .item() ) - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["vit_base"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.VIT_BASE], delta=0.001) def test_pretrained_vit_large_imagenet(self): trainer = Trainer("imagenet_pretrained_vit_large", multi_gpu=MultiGPUMode.OFF) - model = models.get("vit_large", arch_params=self.imagenet_pretrained_arch_params["vit_base"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.VIT_LARGE, arch_params=self.imagenet_pretrained_arch_params[Models.VIT_BASE], **self.imagenet_pretrained_ckpt_params) res = ( trainer.test(model=model, test_loader=self.imagenet_dataset_05_mean_std, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0] .cpu() .item() ) - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["vit_large"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.VIT_LARGE], delta=0.001) def test_pretrained_beit_base_imagenet(self): trainer = Trainer("imagenet_pretrained_beit_base", multi_gpu=MultiGPUMode.OFF) - model = models.get("beit_base_patch16_224", arch_params=self.imagenet_pretrained_arch_params["vit_base"], **self.imagenet_pretrained_ckpt_params) + model = models.get( + Models.BEIT_BASE_PATCH16_224, arch_params=self.imagenet_pretrained_arch_params[Models.VIT_BASE], **self.imagenet_pretrained_ckpt_params + ) res = ( trainer.test(model=model, test_loader=self.imagenet_dataset_05_mean_std, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0] .cpu() .item() ) - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["beit_base_patch16_224"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.BEIT_BASE_PATCH16_224], delta=0.001) def test_transfer_learning_beit_base_imagenet(self): trainer = Trainer("test_transfer_learning_beit_base_imagenet", multi_gpu=MultiGPUMode.OFF) model = models.get( - "beit_base_patch16_224", arch_params=self.imagenet_pretrained_arch_params["vit_base"], **self.imagenet_pretrained_ckpt_params, num_classes=5 + Models.BEIT_BASE_PATCH16_224, + arch_params=self.imagenet_pretrained_arch_params[Models.VIT_BASE], + **self.imagenet_pretrained_ckpt_params, + num_classes=5, ) trainer.train( model=model, @@ -781,7 +814,7 @@ def test_transfer_learning_beit_base_imagenet(self): def test_pretrained_pplite_t_seg50_cityscapes(self): trainer = Trainer("cityscapes_pretrained_pplite_t_seg50", multi_gpu=MultiGPUMode.OFF) - model = models.get("pp_lite_t_seg50", arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.PP_LITE_T_SEG50, arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params) res = ( trainer.test( @@ -793,11 +826,11 @@ def test_pretrained_pplite_t_seg50_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["pp_lite_t_seg50"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.PP_LITE_T_SEG50], delta=0.001) def test_pretrained_pplite_t_seg75_cityscapes(self): trainer = Trainer("cityscapes_pretrained_pplite_t_seg75", multi_gpu=MultiGPUMode.OFF) - model = models.get("pp_lite_t_seg75", arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.PP_LITE_T_SEG75, arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params) res = ( trainer.test( @@ -809,11 +842,11 @@ def test_pretrained_pplite_t_seg75_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["pp_lite_t_seg75"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.PP_LITE_T_SEG75], delta=0.001) def test_pretrained_pplite_b_seg50_cityscapes(self): trainer = Trainer("cityscapes_pretrained_pplite_b_seg50", multi_gpu=MultiGPUMode.OFF) - model = models.get("pp_lite_b_seg50", arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.PP_LITE_B_SEG50, arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params) res = ( trainer.test( @@ -825,11 +858,11 @@ def test_pretrained_pplite_b_seg50_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["pp_lite_b_seg50"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.PP_LITE_B_SEG50], delta=0.001) def test_pretrained_pplite_b_seg75_cityscapes(self): trainer = Trainer("cityscapes_pretrained_pplite_b_seg75", multi_gpu=MultiGPUMode.OFF) - model = models.get("pp_lite_b_seg75", arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.PP_LITE_B_SEG75, arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params) res = ( trainer.test( @@ -841,7 +874,7 @@ def test_pretrained_pplite_b_seg75_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["pp_lite_b_seg75"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.PP_LITE_B_SEG75], delta=0.001) def tearDown(self) -> None: if os.path.exists("~/.cache/torch/hub/"): diff --git a/tests/unit_tests/config_inspector_test.py b/tests/unit_tests/config_inspector_test.py index d64efc00f5..b2a241e9cd 100644 --- a/tests/unit_tests/config_inspector_test.py +++ b/tests/unit_tests/config_inspector_test.py @@ -5,6 +5,7 @@ import pkg_resources from omegaconf import OmegaConf +from super_gradients.common.object_names import Models from super_gradients.training.models import SgModule, get_arch_params from super_gradients.training.models.model_factory import get_architecture from super_gradients.training.utils import HpmStruct @@ -215,7 +216,7 @@ def get_all_arch_params_configs(self): def test_resnet18_cifar_arch_params(self): arch_params = get_arch_params("resnet18_cifar_arch_params") - architecture_cls, arch_params, pretrained_weights_path, is_remote = get_architecture("resnet18", HpmStruct(**arch_params)) + architecture_cls, arch_params, pretrained_weights_path, is_remote = get_architecture(Models.RESNET18, HpmStruct(**arch_params)) with raise_if_unused_params(arch_params) as tracked_arch_params: _ = architecture_cls(arch_params=tracked_arch_params) diff --git a/tests/unit_tests/dataset_statistics_test.py b/tests/unit_tests/dataset_statistics_test.py index 99b6ef3338..f610e5a764 100644 --- a/tests/unit_tests/dataset_statistics_test.py +++ b/tests/unit_tests/dataset_statistics_test.py @@ -1,5 +1,6 @@ import unittest +from super_gradients.common.object_names import Models from super_gradients.training.dataloaders.dataloaders import coco2017_train, coco2017_val from super_gradients.training.metrics.detection_metrics import DetectionMetrics @@ -8,7 +9,6 @@ class TestDatasetStatisticsTensorboardLogger(unittest.TestCase): - def test_dataset_statistics_tensorboard_logger(self): """ ** IMPORTANT NOTE ** @@ -18,27 +18,23 @@ def test_dataset_statistics_tensorboard_logger(self): """ # Create dataset - 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", - } + trainer = Trainer("dataset_statistics_visual_test") + + model = models.get(Models.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", + } 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/detection_utils_test.py b/tests/unit_tests/detection_utils_test.py index 1567867b91..f3d37138c3 100644 --- a/tests/unit_tests/detection_utils_test.py +++ b/tests/unit_tests/detection_utils_test.py @@ -1,6 +1,7 @@ import os import unittest +from super_gradients.common.object_names import Models from super_gradients.training import Trainer, utils as core_utils, models from super_gradients.training.dataloaders.dataloaders import coco2017_val from super_gradients.training.datasets.datasets_conf import COCO_DETECTION_CLASSES_LIST @@ -12,8 +13,8 @@ class TestDetectionUtils(unittest.TestCase): def test_visualization(self): # Create Yolo model - trainer = Trainer('visualization_test') - model = models.get("yolox_n", pretrained_weights="coco") + trainer = Trainer("visualization_test") + model = models.get(Models.YOLOX_N, pretrained_weights="coco") post_prediction_callback = YoloPostPredictionCallback() # Simulate one iteration of validation subset @@ -24,16 +25,15 @@ def test_visualization(self): output = model(imgs) output = post_prediction_callback(output) # Visualize the batch - DetectionVisualization.visualize_batch(imgs, output, targets, batch_i, - COCO_DETECTION_CLASSES_LIST, trainer.checkpoints_dir_path) + DetectionVisualization.visualize_batch(imgs, output, targets, batch_i, COCO_DETECTION_CLASSES_LIST, trainer.checkpoints_dir_path) # Assert images ware created and delete them - img_name = '{}/{}_{}.jpg' + img_name = "{}/{}_{}.jpg" for i in range(4): img_path = img_name.format(trainer.checkpoints_dir_path, batch_i, i) self.assertTrue(os.path.exists(img_path)) os.remove(img_path) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/unit_tests/double_training_test.py b/tests/unit_tests/double_training_test.py index 8c526e7f24..b556aaabc2 100644 --- a/tests/unit_tests/double_training_test.py +++ b/tests/unit_tests/double_training_test.py @@ -1,6 +1,7 @@ import unittest import torch from super_gradients import Trainer +from super_gradients.common.object_names import Models from super_gradients.training import models from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader from super_gradients.training.metrics import Accuracy @@ -18,7 +19,7 @@ def test_call_train_twice(self): trainer = Trainer("external_criterion_test") dataloader = classification_test_dataloader(batch_size=10) - model = models.get("resnet18", num_classes=5) + model = models.get(Models.RESNET18, num_classes=5) train_params = { "max_epochs": 2, "lr_updates": [1], diff --git a/tests/unit_tests/forward_pass_prep_fn_test.py b/tests/unit_tests/forward_pass_prep_fn_test.py index d7701f1277..97c8a27af7 100644 --- a/tests/unit_tests/forward_pass_prep_fn_test.py +++ b/tests/unit_tests/forward_pass_prep_fn_test.py @@ -1,4 +1,6 @@ import unittest + +from super_gradients.common.object_names import Models from super_gradients.training import Trainer, models from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader from super_gradients.training.metrics import Accuracy @@ -20,32 +22,39 @@ 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") - model = models.get("resnet18", num_classes=5) + model = models.get(Models.RESNET18, num_classes=5) 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/kd_trainer_test.py b/tests/unit_tests/kd_trainer_test.py index 19debbb043..61ae5ac045 100644 --- a/tests/unit_tests/kd_trainer_test.py +++ b/tests/unit_tests/kd_trainer_test.py @@ -58,7 +58,7 @@ def setUp(cls): } def test_teacher_sg_module_methods(self): - student = models.get("resnet18", arch_params={"num_classes": 1000}) + student = models.get(Models.RESNET18, arch_params={"num_classes": 1000}) teacher = models.get(Models.RESNET50, arch_params={"num_classes": 1000}, pretrained_weights="imagenet") kd_module = KDModule(arch_params={}, student=student, teacher=teacher) @@ -88,7 +88,7 @@ def test_train_kd_module_external_models(self): 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}) + student = models.get(Models.RESNET18, arch_params={"num_classes": 5}) teacher = models.get(Models.RESNET50, arch_params={"num_classes": 5}, pretrained_weights="imagenet") adapter = NormalizationAdapter( @@ -109,7 +109,7 @@ def test_train_model_with_input_adapter(self): 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}) + student = models.get(Models.RESNET18, arch_params={"num_classes": 5}) teacher = models.get(Models.RESNET50, arch_params={"num_classes": 5}, pretrained_weights="imagenet") train_params = self.kd_train_params.copy() train_params["max_epochs"] = 1 @@ -122,13 +122,13 @@ def test_load_ckpt_best_for_student(self): ) 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(Models.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)) 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}) + student = models.get(Models.RESNET18, arch_params={"num_classes": 5}) teacher = models.get(Models.RESNET50, arch_params={"num_classes": 5}, pretrained_weights="imagenet") train_params = self.kd_train_params.copy() train_params["max_epochs"] = 1 @@ -142,13 +142,13 @@ def test_load_ckpt_best_for_student_with_ema(self): ) 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(Models.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)) def test_resume_kd_training(self): kd_trainer = KDTrainer("test_resume_training_start", device="cpu") - student = models.get("resnet18", arch_params={"num_classes": 5}) + student = models.get(Models.RESNET18, arch_params={"num_classes": 5}) teacher = models.get(Models.RESNET50, arch_params={"num_classes": 5}, pretrained_weights="imagenet") train_params = self.kd_train_params.copy() train_params["max_epochs"] = 1 @@ -162,7 +162,7 @@ def test_resume_kd_training(self): latest_net = deepcopy(kd_trainer.net) kd_trainer = KDTrainer("test_resume_training_start", device="cpu") - student = models.get("resnet18", arch_params={"num_classes": 5}) + student = models.get(Models.RESNET18, arch_params={"num_classes": 5}) teacher = models.get(Models.RESNET50, arch_params={"num_classes": 5}, pretrained_weights="imagenet") train_params["max_epochs"] = 2 diff --git a/tests/unit_tests/local_ckpt_head_replacement_test.py b/tests/unit_tests/local_ckpt_head_replacement_test.py index 1884c84fef..8ba7371683 100644 --- a/tests/unit_tests/local_ckpt_head_replacement_test.py +++ b/tests/unit_tests/local_ckpt_head_replacement_test.py @@ -1,5 +1,6 @@ import unittest +from super_gradients.common.object_names import Models from super_gradients.training import Trainer, models from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader from super_gradients.training.metrics import Accuracy, Top5 @@ -27,12 +28,12 @@ def test_local_ckpt_head_replacement(self): } # Define Model - net = models.get("resnet18", num_classes=5) + net = models.get(Models.RESNET18, num_classes=5) trainer = Trainer("test_resume_training") trainer.train(model=net, training_params=train_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader()) ckpt_path = os.path.join(trainer.checkpoints_dir_path, "ckpt_latest.pth") - net2 = models.get("resnet18", num_classes=10, checkpoint_num_classes=5, checkpoint_path=ckpt_path) + net2 = models.get(Models.RESNET18, num_classes=10, checkpoint_num_classes=5, checkpoint_path=ckpt_path) self.assertFalse(check_models_have_same_weights(net, net2)) net.linear = None diff --git a/tests/unit_tests/loss_loggings_test.py b/tests/unit_tests/loss_loggings_test.py index 56eaa0d790..54d476f0ad 100644 --- a/tests/unit_tests/loss_loggings_test.py +++ b/tests/unit_tests/loss_loggings_test.py @@ -4,6 +4,7 @@ from torchmetrics import Accuracy import torch from super_gradients import Trainer +from super_gradients.common.object_names import Models from super_gradients.training import models from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader @@ -26,52 +27,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} + model = models.get(Models.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, + } 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} + model = models.get(Models.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, + } 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} + model = models.get(Models.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, + } 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/pretrained_models_unit_test.py b/tests/unit_tests/pretrained_models_unit_test.py index 9478f4fb33..1eeff502ba 100644 --- a/tests/unit_tests/pretrained_models_unit_test.py +++ b/tests/unit_tests/pretrained_models_unit_test.py @@ -21,12 +21,12 @@ def test_pretrained_resnet50_imagenet(self): def test_pretrained_regnetY800_imagenet(self): trainer = Trainer("imagenet_pretrained_regnetY800_unit_test", multi_gpu=MultiGPUMode.OFF) - model = models.get("regnetY800", pretrained_weights="imagenet") + model = models.get(Models.REGNETY800, pretrained_weights="imagenet") 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) - model = models.get("repvgg_a0", pretrained_weights="imagenet", arch_params={"build_residual_branches": True}) + model = models.get(Models.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) def tearDown(self) -> None: diff --git a/tests/unit_tests/quantization_utility_tests.py b/tests/unit_tests/quantization_utility_tests.py index 375055a2e9..ffe516503f 100644 --- a/tests/unit_tests/quantization_utility_tests.py +++ b/tests/unit_tests/quantization_utility_tests.py @@ -3,6 +3,8 @@ import torchvision from torch import nn +from super_gradients.common.object_names import Models + try: import super_gradients from pytorch_quantization import nn as quant_nn @@ -782,7 +784,7 @@ def test_sg_resnet_sg_vanilla_quantization_matches_pytorch_quantization(self): default_per_channel_quant_modules=True, ) - resnet_sg: nn.Module = super_gradients.training.models.get("resnet50", pretrained_weights="imagenet", num_classes=1000) + resnet_sg: nn.Module = super_gradients.training.models.get(Models.RESNET50, pretrained_weights="imagenet", num_classes=1000) sq.quantize_module(resnet_sg, preserve_state_dict=True) # PYTORCH-QUANTIZATION @@ -791,7 +793,7 @@ def test_sg_resnet_sg_vanilla_quantization_matches_pytorch_quantization(self): quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input) quant_modules.initialize() - resnet_pyquant: nn.Module = super_gradients.training.models.get("resnet50", pretrained_weights="imagenet", num_classes=1000) + resnet_pyquant: nn.Module = super_gradients.training.models.get(Models.RESNET50, pretrained_weights="imagenet", num_classes=1000) quant_modules.deactivate() diff --git a/tests/unit_tests/test_without_train_test.py b/tests/unit_tests/test_without_train_test.py index cfc276c123..544d6a6971 100644 --- a/tests/unit_tests/test_without_train_test.py +++ b/tests/unit_tests/test_without_train_test.py @@ -33,14 +33,14 @@ def get_classification_trainer(name=""): @staticmethod def get_detection_trainer(name=""): trainer = Trainer(name, multi_gpu=MultiGPUMode.OFF) - model = models.get("yolox_s", num_classes=5) + model = models.get(Models.YOLOX_S, num_classes=5) return trainer, model @staticmethod 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(Models.SHELFNET34_LW, arch_params=shelfnet_lw_arch_params) return trainer, model def test_test_without_train(self): diff --git a/tests/unit_tests/train_after_test_test.py b/tests/unit_tests/train_after_test_test.py index 56797c0e2a..3e8f2f2459 100644 --- a/tests/unit_tests/train_after_test_test.py +++ b/tests/unit_tests/train_after_test_test.py @@ -1,6 +1,7 @@ import unittest import torch from super_gradients import Trainer +from super_gradients.common.object_names import Models from super_gradients.training import models from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader from super_gradients.training.metrics import Accuracy @@ -18,7 +19,7 @@ def test_call_train_after_test(self): trainer = Trainer("test_call_train_after_test") dataloader = classification_test_dataloader(batch_size=10) - model = models.get("resnet18", num_classes=5) + model = models.get(Models.RESNET18, num_classes=5) train_params = { "max_epochs": 2, "lr_updates": [1], @@ -42,7 +43,7 @@ def test_call_train_after_test_with_loss(self): trainer = Trainer("test_call_train_after_test_with_loss") dataloader = classification_test_dataloader(batch_size=10) - model = models.get("resnet18", num_classes=5) + model = models.get(Models.RESNET18, num_classes=5) train_params = { "max_epochs": 2, "lr_updates": [1], 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..d1ed21f175 100644 --- a/tests/unit_tests/train_with_intialized_param_args_test.py +++ b/tests/unit_tests/train_with_intialized_param_args_test.py @@ -7,6 +7,7 @@ from torchmetrics import F1Score from super_gradients import Trainer +from super_gradients.common.object_names import Models from super_gradients.training import models from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader from super_gradients.training.metrics import Accuracy, Top5, ToyTestClassificationMetric @@ -22,28 +23,47 @@ def test_train_with_external_criterion(self): trainer = Trainer("external_criterion_test") 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} + model = models.get(Models.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, + } trainer.train(model=model, training_params=train_params, train_loader=dataloader, valid_loader=dataloader) def test_train_with_external_optimizer(self): trainer = Trainer("external_optimizer_test") dataloader = classification_test_dataloader(batch_size=10) - model = models.get("resnet18", arch_params={"num_classes": 5}) + model = models.get(Models.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): @@ -51,17 +71,24 @@ def test_train_with_external_scheduler(self): dataloader = classification_test_dataloader(batch_size=10) lr = 0.3 - model = models.get("resnet18", arch_params={"num_classes": 5}) + model = models.get(Models.RESNET18, arch_params={"num_classes": 5}) optimizer = SGD(params=model.parameters(), lr=lr) 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) @@ -69,15 +96,21 @@ def test_train_with_external_scheduler_class(self): trainer = Trainer("external_scheduler_test") dataloader = classification_test_dataloader(batch_size=10) - model = models.get("resnet18", arch_params={"num_classes": 5}) + model = models.get(Models.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): @@ -85,18 +118,24 @@ def test_train_with_reduce_on_plateau(self): dataloader = classification_test_dataloader(batch_size=10) lr = 0.3 - model = models.get("resnet18", arch_params={"num_classes": 5}) + model = models.get(Models.RESNET18, arch_params={"num_classes": 5}) optimizer = SGD(params=model.parameters(), lr=lr) 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) @@ -104,37 +143,55 @@ def test_train_with_external_metric(self): trainer = Trainer("external_metric_test") 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} + model = models.get(Models.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, + } 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} + model = models.get(Models.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, + } 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/vit_unit_test.py b/tests/unit_tests/vit_unit_test.py index 928385b61b..9914f7b469 100644 --- a/tests/unit_tests/vit_unit_test.py +++ b/tests/unit_tests/vit_unit_test.py @@ -1,5 +1,6 @@ import unittest +from super_gradients.common.object_names import Models from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader from super_gradients import Trainer from super_gradients.training.metrics import Accuracy, Top5 @@ -31,7 +32,7 @@ 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) + model = models.get(Models.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() ) From 20492c161f6b31481f3d68fc3590853617f11a39 Mon Sep 17 00:00:00 2001 From: shayaharon Date: Tue, 10 Jan 2023 13:06:57 +0200 Subject: [PATCH 3/7] more models import fixes --- .../quantization/resnet_qat_example.py | 3 +- .../pretrained_models_test.py | 248 ++++++++++-------- tests/unit_tests/kd_ema_test.py | 13 +- tests/unit_tests/kd_trainer_test.py | 29 +- .../unit_tests/pretrained_models_unit_test.py | 5 +- .../unit_tests/quantization_utility_tests.py | 6 +- tests/unit_tests/test_without_train_test.py | 3 +- 7 files changed, 171 insertions(+), 136 deletions(-) diff --git a/src/super_gradients/examples/quantization/resnet_qat_example.py b/src/super_gradients/examples/quantization/resnet_qat_example.py index 9cdd277cc5..a8447b4f75 100644 --- a/src/super_gradients/examples/quantization/resnet_qat_example.py +++ b/src/super_gradients/examples/quantization/resnet_qat_example.py @@ -4,6 +4,7 @@ import super_gradients from super_gradients import Trainer +from super_gradients.common.object_names import Models from super_gradients.modules.quantization.resnet_bottleneck import QuantBottleneck as sg_QuantizedBottleneck from super_gradients.training import MultiGPUMode from super_gradients.training import models as sg_models @@ -55,7 +56,7 @@ def selective_quantize(model: nn.Module): def sg_vanilla_resnet50(): - return sg_models.get("resnet50", pretrained_weights="imagenet", num_classes=1000) + return sg_models.get(Models.RESNET50, pretrained_weights="imagenet", num_classes=1000) def sg_naive_qdq_resnet50(): diff --git a/tests/integration_tests/pretrained_models_test.py b/tests/integration_tests/pretrained_models_test.py index 7a0e468437..d489cd8ca6 100644 --- a/tests/integration_tests/pretrained_models_test.py +++ b/tests/integration_tests/pretrained_models_test.py @@ -1,6 +1,7 @@ #!/usr/bin/env python3 import unittest +from super_gradients.common.object_names import Models from super_gradients.training import Trainer from super_gradients.training.dataloaders import imagenet_val, imagenet_vit_base_val from super_gradients.training.dataloaders.dataloaders import ( @@ -35,19 +36,19 @@ def setUp(self) -> None: self.imagenet_pretrained_arch_params = { "resnet": {}, "regnet": {}, - "repvgg_a0": {"build_residual_branches": True}, - "efficientnet_b0": {}, + Models.REPVGG_A0: {"build_residual_branches": True}, + Models.EFFICIENTNET_B0: {}, "mobilenet": {}, - "vit_base": {"image_size": (224, 224), "patch_size": (16, 16)}, + Models.VIT_BASE: {"image_size": (224, 224), "patch_size": (16, 16)}, } self.imagenet_pretrained_trainsfer_learning_arch_params = { "resnet": {}, "regnet": {}, - "repvgg_a0": {"build_residual_branches": True}, - "efficientnet_b0": {}, + Models.REPVGG_A0: {"build_residual_branches": True}, + Models.EFFICIENTNET_B0: {}, "mobilenet": {}, - "vit_base": {"image_size": (224, 224), "patch_size": (16, 16)}, + Models.VIT_BASE: {"image_size": (224, 224), "patch_size": (16, 16)}, } self.imagenet_pretrained_ckpt_params = {"pretrained_weights": "imagenet"} @@ -55,21 +56,21 @@ def setUp(self) -> None: self.imagenet21k_pretrained_ckpt_params = {"pretrained_weights": "imagenet21k"} self.imagenet_pretrained_accuracies = { - "resnet50": 0.8191, - "resnet34": 0.7413, - "resnet18": 0.706, - "repvgg_a0": 0.7205, - "regnetY800": 0.7707, - "regnetY600": 0.7618, - "regnetY400": 0.7474, - "regnetY200": 0.7088, - "efficientnet_b0": 0.7762, - "mobilenet_v3_large": 0.7452, - "mobilenet_v3_small": 0.6745, - "mobilenet_v2": 0.7308, - "vit_base": 0.8415, - "vit_large": 0.8564, - "beit_base_patch16_224": 0.85, + Models.RESNET50: 0.8191, + Models.RESNET34: 0.7413, + Models.RESNET18: 0.706, + Models.REPVGG_A0: 0.7205, + Models.REGNETY800: 0.7707, + Models.REGNETY600: 0.7618, + Models.REGNETY400: 0.7474, + Models.REGNETY200: 0.7088, + Models.EFFICIENTNET_B0: 0.7762, + Models.MOBILENET_V3_LARGE: 0.7452, + Models.MOBILENET_V3_SMALL: 0.6745, + Models.MOBILENET_V2: 0.7308, + Models.VIT_BASE: 0.8415, + Models.VIT_LARGE: 0.8564, + Models.BEIT_BASE_PATCH16_224: 0.85, } self.imagenet_dataset = imagenet_val(dataloader_params={"batch_size": 128}) @@ -104,11 +105,11 @@ def setUp(self) -> None: self.coco_pretrained_maps = { "ssd_lite_mobilenet_v2": 0.2041, "coco_ssd_mobilenet_v1": 0.243, - "yolox_s": 0.4047, - "yolox_m": 0.4640, - "yolox_l": 0.4925, - "yolox_n": 0.2677, - "yolox_t": 0.3718, + Models.YOLOX_S: 0.4047, + Models.YOLOX_M: 0.4640, + Models.YOLOX_L: 0.4925, + Models.YOLOX_N: 0.2677, + Models.YOLOX_T: 0.3718, } self.transfer_detection_dataset = detection_test_dataloader() @@ -151,10 +152,10 @@ def setUp(self) -> None: self.coco_segmentation_subclass_pretrained_mious = {"shelfnet34_lw": 0.651} self.coco_segmentation_dataset = coco_segmentation_val() - self.cityscapes_pretrained_models = ["ddrnet_23", "ddrnet_23_slim", "stdc1_seg50", "regseg48"] + self.cityscapes_pretrained_models = ["ddrnet_23", "ddrnet_23_slim", Models.STDC1_SEG50, Models.REGSEG48] self.cityscapes_pretrained_arch_params = { "ddrnet_23": {"aux_head": True}, - "regseg48": {}, + Models.REGSEG48: {}, "stdc": {"use_aux_heads": True, "aux_head": True}, "pplite_seg": {"use_aux_heads": True}, } @@ -163,11 +164,11 @@ def setUp(self) -> None: self.cityscapes_pretrained_mious = { "ddrnet_23": 0.8026, "ddrnet_23_slim": 0.7801, - "stdc1_seg50": 0.7511, - "stdc1_seg75": 0.7687, - "stdc2_seg50": 0.7644, - "stdc2_seg75": 0.7893, - "regseg48": 0.7815, + Models.STDC1_SEG50: 0.7511, + Models.STDC1_SEG75: 0.7687, + Models.STDC2_SEG50: 0.7644, + Models.STDC2_SEG75: 0.7893, + Models.REGSEG48: 0.7815, "pp_lite_t_seg50": 0.7492, "pp_lite_t_seg75": 0.7756, "pp_lite_b_seg50": 0.7648, @@ -229,13 +230,13 @@ def setUp(self) -> None: def test_pretrained_resnet50_imagenet(self): trainer = Trainer("imagenet_pretrained_resnet50") - model = models.get("resnet50", arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.RESNET50, arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["resnet50"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.RESNET50], delta=0.001) def test_transfer_learning_resnet50_imagenet(self): trainer = Trainer("imagenet_pretrained_resnet50_transfer_learning") - model = models.get("resnet50", arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) + model = models.get(Models.RESNET50, arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -246,13 +247,13 @@ def test_transfer_learning_resnet50_imagenet(self): def test_pretrained_resnet34_imagenet(self): trainer = Trainer("imagenet_pretrained_resnet34") - model = models.get("resnet34", arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.RESNET34, arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["resnet34"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.RESNET34], delta=0.001) def test_transfer_learning_resnet34_imagenet(self): trainer = Trainer("imagenet_pretrained_resnet34_transfer_learning") - model = models.get("resnet34", arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) + model = models.get(Models.RESNET34, arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -263,13 +264,13 @@ def test_transfer_learning_resnet34_imagenet(self): def test_pretrained_resnet18_imagenet(self): trainer = Trainer("imagenet_pretrained_resnet18") - model = models.get("resnet18", arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.RESNET18, arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["resnet18"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.RESNET18], delta=0.001) def test_transfer_learning_resnet18_imagenet(self): trainer = Trainer("imagenet_pretrained_resnet18_transfer_learning") - model = models.get("resnet18", arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) + model = models.get(Models.RESNET18, arch_params=self.imagenet_pretrained_arch_params["resnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -280,13 +281,13 @@ def test_transfer_learning_resnet18_imagenet(self): def test_pretrained_regnetY800_imagenet(self): trainer = Trainer("imagenet_pretrained_regnetY800") - model = models.get("regnetY800", arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.REGNETY800, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY800"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.REGNETY800], delta=0.001) def test_transfer_learning_regnetY800_imagenet(self): trainer = Trainer("imagenet_pretrained_regnetY800_transfer_learning") - model = models.get("regnetY800", arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) + model = models.get(Models.REGNETY800, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -297,13 +298,13 @@ def test_transfer_learning_regnetY800_imagenet(self): def test_pretrained_regnetY600_imagenet(self): trainer = Trainer("imagenet_pretrained_regnetY600") - model = models.get("regnetY600", arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.REGNETY600, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY600"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.REGNETY600], delta=0.001) def test_transfer_learning_regnetY600_imagenet(self): trainer = Trainer("imagenet_pretrained_regnetY600_transfer_learning") - model = models.get("regnetY600", arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) + model = models.get(Models.REGNETY600, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -314,13 +315,13 @@ def test_transfer_learning_regnetY600_imagenet(self): def test_pretrained_regnetY400_imagenet(self): trainer = Trainer("imagenet_pretrained_regnetY400") - model = models.get("regnetY400", arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.REGNETY400, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY400"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.REGNETY400], delta=0.001) def test_transfer_learning_regnetY400_imagenet(self): trainer = Trainer("imagenet_pretrained_regnetY400_transfer_learning") - model = models.get("regnetY400", arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) + model = models.get(Models.REGNETY400, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -331,13 +332,13 @@ def test_transfer_learning_regnetY400_imagenet(self): def test_pretrained_regnetY200_imagenet(self): trainer = Trainer("imagenet_pretrained_regnetY200") - model = models.get("regnetY200", arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.REGNETY200, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY200"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.REGNETY200], delta=0.001) def test_transfer_learning_regnetY200_imagenet(self): trainer = Trainer("imagenet_pretrained_regnetY200_transfer_learning") - model = models.get("regnetY200", arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) + model = models.get(Models.REGNETY200, arch_params=self.imagenet_pretrained_arch_params["regnet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -348,13 +349,15 @@ def test_transfer_learning_regnetY200_imagenet(self): def test_pretrained_repvgg_a0_imagenet(self): trainer = Trainer("imagenet_pretrained_repvgg_a0") - model = models.get("repvgg_a0", arch_params=self.imagenet_pretrained_arch_params["repvgg_a0"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.REPVGG_A0, arch_params=self.imagenet_pretrained_arch_params[Models.REPVGG_A0], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["repvgg_a0"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.REPVGG_A0], delta=0.001) def test_transfer_learning_repvgg_a0_imagenet(self): trainer = Trainer("imagenet_pretrained_repvgg_a0_transfer_learning") - model = models.get("repvgg_a0", arch_params=self.imagenet_pretrained_arch_params["repvgg_a0"], **self.imagenet_pretrained_ckpt_params, num_classes=5) + model = models.get( + Models.REPVGG_A0, arch_params=self.imagenet_pretrained_arch_params[Models.REPVGG_A0], **self.imagenet_pretrained_ckpt_params, num_classes=5 + ) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -364,7 +367,7 @@ def test_transfer_learning_repvgg_a0_imagenet(self): def test_pretrained_regseg48_cityscapes(self): trainer = Trainer("cityscapes_pretrained_regseg48") - model = models.get("regseg48", arch_params=self.cityscapes_pretrained_arch_params["regseg48"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.REGSEG48, arch_params=self.cityscapes_pretrained_arch_params[Models.REGSEG48], **self.cityscapes_pretrained_ckpt_params) res = ( trainer.test( model=model, test_loader=self.cityscapes_dataset, test_metrics_list=[IoU(num_classes=20, ignore_index=19)], metrics_progress_verbose=True @@ -372,11 +375,11 @@ def test_pretrained_regseg48_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["regseg48"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.REGSEG48], delta=0.001) def test_transfer_learning_regseg48_cityscapes(self): trainer = Trainer("regseg48_cityscapes_transfer_learning") - model = models.get("regseg48", arch_params=self.cityscapes_pretrained_arch_params["regseg48"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.REGSEG48, arch_params=self.cityscapes_pretrained_arch_params[Models.REGSEG48], **self.cityscapes_pretrained_ckpt_params) trainer.train( model=model, train_loader=self.transfer_segmentation_dataset, @@ -441,15 +444,20 @@ def test_pretrained_coco_segmentation_subclass_pretrained_shelfnet34_lw(self): def test_pretrained_efficientnet_b0_imagenet(self): trainer = Trainer("imagenet_pretrained_efficientnet_b0") - model = models.get("efficientnet_b0", arch_params=self.imagenet_pretrained_arch_params["efficientnet_b0"], **self.imagenet_pretrained_ckpt_params) + model = models.get( + Models.EFFICIENTNET_B0, arch_params=self.imagenet_pretrained_arch_params[Models.EFFICIENTNET_B0], **self.imagenet_pretrained_ckpt_params + ) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["efficientnet_b0"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.EFFICIENTNET_B0], delta=0.001) def test_transfer_learning_efficientnet_b0_imagenet(self): trainer = Trainer("imagenet_pretrained_efficientnet_b0_transfer_learning") model = models.get( - "efficientnet_b0", arch_params=self.imagenet_pretrained_arch_params["efficientnet_b0"], **self.imagenet_pretrained_ckpt_params, num_classes=5 + Models.EFFICIENTNET_B0, + arch_params=self.imagenet_pretrained_arch_params[Models.EFFICIENTNET_B0], + **self.imagenet_pretrained_ckpt_params, + num_classes=5, ) trainer.train( model=model, @@ -484,7 +492,7 @@ def test_transfer_learning_ssd_lite_mobilenet_v2_coco(self): def test_pretrained_ssd_mobilenet_v1_coco(self): trainer = Trainer("coco_ssd_mobilenet_v1") - model = models.get("ssd_mobilenet_v1", arch_params=self.coco_pretrained_arch_params["coco_ssd_mobilenet_v1"], **self.coco_pretrained_ckpt_params) + model = models.get(Models.SSD_MOBILENET_V1, arch_params=self.coco_pretrained_arch_params["coco_ssd_mobilenet_v1"], **self.coco_pretrained_ckpt_params) ssd_post_prediction_callback = SSDPostPredictCallback() res = trainer.test( model=model, @@ -495,60 +503,60 @@ def test_pretrained_ssd_mobilenet_v1_coco(self): self.assertAlmostEqual(res, self.coco_pretrained_maps["coco_ssd_mobilenet_v1"], delta=0.001) def test_pretrained_yolox_s_coco(self): - trainer = Trainer("yolox_s") + trainer = Trainer(Models.YOLOX_S) - model = models.get("yolox_s", **self.coco_pretrained_ckpt_params) + model = models.get(Models.YOLOX_S, **self.coco_pretrained_ckpt_params) res = trainer.test( model=model, test_loader=self.coco_dataset["yolox"], test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=80, normalize_targets=True)], )[2] - self.assertAlmostEqual(res, self.coco_pretrained_maps["yolox_s"], delta=0.001) + self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.YOLOX_S], delta=0.001) def test_pretrained_yolox_m_coco(self): - trainer = Trainer("yolox_m") - model = models.get("yolox_m", **self.coco_pretrained_ckpt_params) + trainer = Trainer(Models.YOLOX_M) + model = models.get(Models.YOLOX_M, **self.coco_pretrained_ckpt_params) res = trainer.test( model=model, test_loader=self.coco_dataset["yolox"], test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=80, normalize_targets=True)], )[2] - self.assertAlmostEqual(res, self.coco_pretrained_maps["yolox_m"], delta=0.001) + self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.YOLOX_M], delta=0.001) def test_pretrained_yolox_l_coco(self): - trainer = Trainer("yolox_l") - model = models.get("yolox_l", **self.coco_pretrained_ckpt_params) + trainer = Trainer(Models.YOLOX_L) + model = models.get(Models.YOLOX_L, **self.coco_pretrained_ckpt_params) res = trainer.test( model=model, test_loader=self.coco_dataset["yolox"], test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=80, normalize_targets=True)], )[2] - self.assertAlmostEqual(res, self.coco_pretrained_maps["yolox_l"], delta=0.001) + self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.YOLOX_L], delta=0.001) def test_pretrained_yolox_n_coco(self): - trainer = Trainer("yolox_n") + trainer = Trainer(Models.YOLOX_N) - model = models.get("yolox_n", **self.coco_pretrained_ckpt_params) + model = models.get(Models.YOLOX_N, **self.coco_pretrained_ckpt_params) res = trainer.test( model=model, test_loader=self.coco_dataset["yolox"], test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=80, normalize_targets=True)], )[2] - self.assertAlmostEqual(res, self.coco_pretrained_maps["yolox_n"], delta=0.001) + self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.YOLOX_N], delta=0.001) def test_pretrained_yolox_t_coco(self): - trainer = Trainer("yolox_t") - model = models.get("yolox_t", **self.coco_pretrained_ckpt_params) + trainer = Trainer(Models.YOLOX_T) + model = models.get(Models.YOLOX_T, **self.coco_pretrained_ckpt_params) res = trainer.test( model=model, test_loader=self.coco_dataset["yolox"], test_metrics_list=[DetectionMetrics(post_prediction_callback=YoloPostPredictionCallback(), num_cls=80, normalize_targets=True)], )[2] - self.assertAlmostEqual(res, self.coco_pretrained_maps["yolox_t"], delta=0.001) + self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.YOLOX_T], delta=0.001) def test_transfer_learning_yolox_n_coco(self): trainer = Trainer("test_transfer_learning_yolox_n_coco") - model = models.get("yolox_n", **self.coco_pretrained_ckpt_params, num_classes=5) + model = models.get(Models.YOLOX_N, **self.coco_pretrained_ckpt_params, num_classes=5) trainer.train( model=model, training_params=self.transfer_detection_train_params_yolox, @@ -560,7 +568,7 @@ def test_transfer_learning_mobilenet_v3_large_imagenet(self): trainer = Trainer("imagenet_pretrained_mobilenet_v3_large_transfer_learning") model = models.get( - "mobilenet_v3_large", arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params, num_classes=5 + Models.MOBILENET_V3_LARGE, arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params, num_classes=5 ) trainer.train( model=model, @@ -572,15 +580,15 @@ def test_transfer_learning_mobilenet_v3_large_imagenet(self): def test_pretrained_mobilenet_v3_large_imagenet(self): trainer = Trainer("imagenet_mobilenet_v3_large") - model = models.get("mobilenet_v3_large", arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.MOBILENET_V3_LARGE, arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["mobilenet_v3_large"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.MOBILENET_V3_LARGE], delta=0.001) def test_transfer_learning_mobilenet_v3_small_imagenet(self): trainer = Trainer("imagenet_pretrained_mobilenet_v3_small_transfer_learning") model = models.get( - "mobilenet_v3_small", arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params, num_classes=5 + Models.MOBILENET_V3_SMALL, arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params, num_classes=5 ) trainer.train( model=model, @@ -592,14 +600,16 @@ def test_transfer_learning_mobilenet_v3_small_imagenet(self): def test_pretrained_mobilenet_v3_small_imagenet(self): trainer = Trainer("imagenet_mobilenet_v3_small") - model = models.get("mobilenet_v3_small", arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.MOBILENET_V3_SMALL, arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["mobilenet_v3_small"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.MOBILENET_V3_SMALL], delta=0.001) def test_transfer_learning_mobilenet_v2_imagenet(self): trainer = Trainer("imagenet_pretrained_mobilenet_v2_transfer_learning") - model = models.get("mobilenet_v2", arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params, num_classes=5) + model = models.get( + Models.MOBILENET_V2, arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params, num_classes=5 + ) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -610,13 +620,14 @@ def test_transfer_learning_mobilenet_v2_imagenet(self): def test_pretrained_mobilenet_v2_imagenet(self): trainer = Trainer("imagenet_mobilenet_v2") - model = models.get("mobilenet_v2", arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.MOBILENET_V2, arch_params=self.imagenet_pretrained_arch_params["mobilenet"], **self.imagenet_pretrained_ckpt_params) res = trainer.test(model=model, test_loader=self.imagenet_dataset, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0].cpu().item() - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["mobilenet_v2"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.MOBILENET_V2], delta=0.001) def test_pretrained_stdc1_seg50_cityscapes(self): trainer = Trainer("cityscapes_pretrained_stdc1_seg50") - model = models.get("stdc1_seg50", arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params) + + model = models.get(Models.STDC1_SEG50, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params) res = ( trainer.test( model=model, @@ -627,11 +638,13 @@ def test_pretrained_stdc1_seg50_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["stdc1_seg50"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.STDC1_SEG50], delta=0.001) def test_transfer_learning_stdc1_seg50_cityscapes(self): trainer = Trainer("cityscapes_pretrained_stdc1_seg50_transfer_learning") - model = models.get("stdc1_seg50", arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params, num_classes=5) + model = models.get( + Models.STDC1_SEG50, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params, num_classes=5 + ) trainer.train( model=model, training_params=self.stdc_transfer_segmentation_train_params, @@ -641,7 +654,7 @@ def test_transfer_learning_stdc1_seg50_cityscapes(self): def test_pretrained_stdc1_seg75_cityscapes(self): trainer = Trainer("cityscapes_pretrained_stdc1_seg75") - model = models.get("stdc1_seg75", arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.STDC1_SEG75, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params) res = ( trainer.test( model=model, @@ -652,11 +665,13 @@ def test_pretrained_stdc1_seg75_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["stdc1_seg75"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.STDC1_SEG75], delta=0.001) def test_transfer_learning_stdc1_seg75_cityscapes(self): trainer = Trainer("cityscapes_pretrained_stdc1_seg75_transfer_learning") - model = models.get("stdc1_seg75", arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params, num_classes=5) + model = models.get( + Models.STDC1_SEG75, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params, num_classes=5 + ) trainer.train( model=model, training_params=self.stdc_transfer_segmentation_train_params, @@ -666,7 +681,7 @@ def test_transfer_learning_stdc1_seg75_cityscapes(self): def test_pretrained_stdc2_seg50_cityscapes(self): trainer = Trainer("cityscapes_pretrained_stdc2_seg50") - model = models.get("stdc2_seg50", arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.STDC2_SEG50, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params) res = ( trainer.test( model=model, @@ -677,11 +692,13 @@ def test_pretrained_stdc2_seg50_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["stdc2_seg50"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.STDC2_SEG50], delta=0.001) def test_transfer_learning_stdc2_seg50_cityscapes(self): trainer = Trainer("cityscapes_pretrained_stdc2_seg50_transfer_learning") - model = models.get("stdc2_seg50", arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params, num_classes=5) + model = models.get( + Models.STDC2_SEG50, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params, num_classes=5 + ) trainer.train( model=model, training_params=self.stdc_transfer_segmentation_train_params, @@ -691,7 +708,7 @@ def test_transfer_learning_stdc2_seg50_cityscapes(self): def test_pretrained_stdc2_seg75_cityscapes(self): trainer = Trainer("cityscapes_pretrained_stdc2_seg75") - model = models.get("stdc2_seg75", arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.STDC2_SEG75, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params) res = ( trainer.test( model=model, @@ -702,11 +719,13 @@ def test_pretrained_stdc2_seg75_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["stdc2_seg75"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.STDC2_SEG75], delta=0.001) def test_transfer_learning_stdc2_seg75_cityscapes(self): trainer = Trainer("cityscapes_pretrained_stdc2_seg75_transfer_learning") - model = models.get("stdc2_seg75", arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params, num_classes=5) + model = models.get( + Models.STDC2_SEG75, arch_params=self.cityscapes_pretrained_arch_params["stdc"], **self.cityscapes_pretrained_ckpt_params, num_classes=5 + ) trainer.train( model=model, training_params=self.stdc_transfer_segmentation_train_params, @@ -717,7 +736,9 @@ def test_transfer_learning_stdc2_seg75_cityscapes(self): def test_transfer_learning_vit_base_imagenet21k(self): trainer = Trainer("imagenet21k_pretrained_vit_base") - model = models.get("vit_base", arch_params=self.imagenet_pretrained_arch_params["vit_base"], **self.imagenet21k_pretrained_ckpt_params, num_classes=5) + model = models.get( + Models.VIT_BASE, arch_params=self.imagenet_pretrained_arch_params[Models.VIT_BASE], **self.imagenet21k_pretrained_ckpt_params, num_classes=5 + ) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -728,7 +749,9 @@ def test_transfer_learning_vit_base_imagenet21k(self): def test_transfer_learning_vit_large_imagenet21k(self): trainer = Trainer("imagenet21k_pretrained_vit_large") - model = models.get("vit_large", arch_params=self.imagenet_pretrained_arch_params["vit_base"], **self.imagenet21k_pretrained_ckpt_params, num_classes=5) + model = models.get( + Models.VIT_LARGE, arch_params=self.imagenet_pretrained_arch_params[Models.VIT_BASE], **self.imagenet21k_pretrained_ckpt_params, num_classes=5 + ) trainer.train( model=model, training_params=self.transfer_classification_train_params, @@ -738,39 +761,44 @@ def test_transfer_learning_vit_large_imagenet21k(self): def test_pretrained_vit_base_imagenet(self): trainer = Trainer("imagenet_pretrained_vit_base") - model = models.get("vit_base", arch_params=self.imagenet_pretrained_arch_params["vit_base"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.VIT_BASE, arch_params=self.imagenet_pretrained_arch_params[Models.VIT_BASE], **self.imagenet_pretrained_ckpt_params) res = ( trainer.test(model=model, test_loader=self.imagenet_dataset_05_mean_std, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0] .cpu() .item() ) - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["vit_base"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.VIT_BASE], delta=0.001) def test_pretrained_vit_large_imagenet(self): trainer = Trainer("imagenet_pretrained_vit_large") - model = models.get("vit_large", arch_params=self.imagenet_pretrained_arch_params["vit_base"], **self.imagenet_pretrained_ckpt_params) + model = models.get(Models.VIT_LARGE, arch_params=self.imagenet_pretrained_arch_params[Models.VIT_BASE], **self.imagenet_pretrained_ckpt_params) res = ( trainer.test(model=model, test_loader=self.imagenet_dataset_05_mean_std, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0] .cpu() .item() ) - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["vit_large"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.VIT_LARGE], delta=0.001) def test_pretrained_beit_base_imagenet(self): trainer = Trainer("imagenet_pretrained_beit_base") - model = models.get("beit_base_patch16_224", arch_params=self.imagenet_pretrained_arch_params["vit_base"], **self.imagenet_pretrained_ckpt_params) + model = models.get( + Models.BEIT_BASE_PATCH16_224, arch_params=self.imagenet_pretrained_arch_params[Models.VIT_BASE], **self.imagenet_pretrained_ckpt_params + ) res = ( trainer.test(model=model, test_loader=self.imagenet_dataset_05_mean_std, test_metrics_list=[Accuracy()], metrics_progress_verbose=True)[0] .cpu() .item() ) - self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["beit_base_patch16_224"], delta=0.001) + self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies[Models.BEIT_BASE_PATCH16_224], delta=0.001) def test_transfer_learning_beit_base_imagenet(self): trainer = Trainer("test_transfer_learning_beit_base_imagenet") model = models.get( - "beit_base_patch16_224", arch_params=self.imagenet_pretrained_arch_params["vit_base"], **self.imagenet_pretrained_ckpt_params, num_classes=5 + Models.BEIT_BASE_PATCH16_224, + arch_params=self.imagenet_pretrained_arch_params[Models.VIT_BASE], + **self.imagenet_pretrained_ckpt_params, + num_classes=5, ) trainer.train( model=model, diff --git a/tests/unit_tests/kd_ema_test.py b/tests/unit_tests/kd_ema_test.py index 7cca0c6edd..09980cd4bb 100644 --- a/tests/unit_tests/kd_ema_test.py +++ b/tests/unit_tests/kd_ema_test.py @@ -8,6 +8,7 @@ from super_gradients.training.utils.utils import check_models_have_same_weights from super_gradients.training.metrics import Accuracy from super_gradients.training.losses.kd_losses import KDLogitsLoss +from super_gradients.common.object_names import Models class KDEMATest(unittest.TestCase): @@ -39,8 +40,8 @@ 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") - student = models.get("resnet18", arch_params={"num_classes": 1000}) - teacher = models.get("resnet50", arch_params={"num_classes": 1000}, pretrained_weights="imagenet") + student = models.get(Models.RESNET18, arch_params={"num_classes": 1000}) + teacher = models.get(Models.RESNET50, arch_params={"num_classes": 1000}, pretrained_weights="imagenet") kd_model.train( training_params=self.kd_train_params, @@ -59,8 +60,8 @@ 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") - student = models.get("resnet18", arch_params={"num_classes": 1000}) - teacher = models.get("resnet50", arch_params={"num_classes": 1000}, pretrained_weights="imagenet") + student = models.get(Models.RESNET18, arch_params={"num_classes": 1000}) + teacher = models.get(Models.RESNET50, arch_params={"num_classes": 1000}, pretrained_weights="imagenet") kd_model.train( training_params=self.kd_train_params, @@ -74,8 +75,8 @@ def test_kd_ckpt_reload_net(self): # Load the trained KD trainer kd_model = KDTrainer("test_kd_ema_ckpt_reload") - student = models.get("resnet18", arch_params={"num_classes": 1000}) - teacher = models.get("resnet50", arch_params={"num_classes": 1000}, pretrained_weights="imagenet") + student = models.get(Models.RESNET18, arch_params={"num_classes": 1000}) + teacher = models.get(Models.RESNET50, arch_params={"num_classes": 1000}, pretrained_weights="imagenet") train_params["resume"] = True kd_model.train( diff --git a/tests/unit_tests/kd_trainer_test.py b/tests/unit_tests/kd_trainer_test.py index 33d0910f66..06fdf91000 100644 --- a/tests/unit_tests/kd_trainer_test.py +++ b/tests/unit_tests/kd_trainer_test.py @@ -2,6 +2,7 @@ import unittest from copy import deepcopy +from super_gradients.common.object_names import Models from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader from super_gradients.training.kd_trainer.kd_trainer import KDTrainer import torch @@ -57,8 +58,8 @@ def setUp(cls): } 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") + student = models.get(Models.RESNET18, arch_params={"num_classes": 1000}) + teacher = models.get(Models.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={}) @@ -87,8 +88,8 @@ def test_train_kd_module_external_models(self): def test_train_model_with_input_adapter(self): kd_trainer = KDTrainer("train_kd_module_with_with_input_adapter") - student = models.get("resnet18", arch_params={"num_classes": 5}) - teacher = models.get("resnet50", arch_params={"num_classes": 5}, pretrained_weights="imagenet") + student = models.get(Models.RESNET18, arch_params={"num_classes": 5}) + teacher = models.get(Models.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] @@ -108,8 +109,8 @@ def test_train_model_with_input_adapter(self): def test_load_ckpt_best_for_student(self): kd_trainer = KDTrainer("test_load_ckpt_best") - student = models.get("resnet18", arch_params={"num_classes": 5}) - teacher = models.get("resnet50", arch_params={"num_classes": 5}, pretrained_weights="imagenet") + student = models.get(Models.RESNET18, arch_params={"num_classes": 5}) + teacher = models.get(Models.RESNET50, arch_params={"num_classes": 5}, pretrained_weights="imagenet") train_params = self.kd_train_params.copy() train_params["max_epochs"] = 1 kd_trainer.train( @@ -121,14 +122,14 @@ def test_load_ckpt_best_for_student(self): ) 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(Models.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)) def test_load_ckpt_best_for_student_with_ema(self): kd_trainer = KDTrainer("test_load_ckpt_best") - student = models.get("resnet18", arch_params={"num_classes": 5}) - teacher = models.get("resnet50", arch_params={"num_classes": 5}, pretrained_weights="imagenet") + student = models.get(Models.RESNET18, arch_params={"num_classes": 5}) + teacher = models.get(Models.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 @@ -141,14 +142,14 @@ def test_load_ckpt_best_for_student_with_ema(self): ) 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(Models.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)) def test_resume_kd_training(self): kd_trainer = KDTrainer("test_resume_training_start") - student = models.get("resnet18", arch_params={"num_classes": 5}) - teacher = models.get("resnet50", arch_params={"num_classes": 5}, pretrained_weights="imagenet") + student = models.get(Models.RESNET18, arch_params={"num_classes": 5}) + teacher = models.get(Models.RESNET50, arch_params={"num_classes": 5}, pretrained_weights="imagenet") train_params = self.kd_train_params.copy() train_params["max_epochs"] = 1 kd_trainer.train( @@ -161,8 +162,8 @@ def test_resume_kd_training(self): latest_net = deepcopy(kd_trainer.net) kd_trainer = KDTrainer("test_resume_training_start") - student = models.get("resnet18", arch_params={"num_classes": 5}) - teacher = models.get("resnet50", arch_params={"num_classes": 5}, pretrained_weights="imagenet") + student = models.get(Models.RESNET18, arch_params={"num_classes": 5}) + teacher = models.get(Models.RESNET50, arch_params={"num_classes": 5}, pretrained_weights="imagenet") train_params["max_epochs"] = 2 train_params["resume"] = True diff --git a/tests/unit_tests/pretrained_models_unit_test.py b/tests/unit_tests/pretrained_models_unit_test.py index e6d75e19b1..85b03fd31a 100644 --- a/tests/unit_tests/pretrained_models_unit_test.py +++ b/tests/unit_tests/pretrained_models_unit_test.py @@ -1,5 +1,6 @@ import unittest import super_gradients +from super_gradients.common.object_names import Models from super_gradients.training import models from super_gradients.training import Trainer from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader @@ -11,11 +12,11 @@ class PretrainedModelsUnitTest(unittest.TestCase): def setUp(self) -> None: super_gradients.init_trainer() - self.imagenet_pretrained_models = ["resnet50", "repvgg_a0", "regnetY800"] + self.imagenet_pretrained_models = [Models.RESNET50, "repvgg_a0", "regnetY800"] def test_pretrained_resnet50_imagenet(self): trainer = Trainer("imagenet_pretrained_resnet50_unit_test") - model = models.get("resnet50", pretrained_weights="imagenet") + model = models.get(Models.RESNET50, pretrained_weights="imagenet") trainer.test(model=model, test_loader=classification_test_dataloader(), test_metrics_list=[Accuracy()], metrics_progress_verbose=True) def test_pretrained_regnetY800_imagenet(self): diff --git a/tests/unit_tests/quantization_utility_tests.py b/tests/unit_tests/quantization_utility_tests.py index a11050b22f..58ab83f1ae 100644 --- a/tests/unit_tests/quantization_utility_tests.py +++ b/tests/unit_tests/quantization_utility_tests.py @@ -3,6 +3,8 @@ import torchvision from torch import nn +from super_gradients.common.object_names import Models + try: import super_gradients from pytorch_quantization import nn as quant_nn @@ -787,7 +789,7 @@ def test_sg_resnet_sg_vanilla_quantization_matches_pytorch_quantization(self): if Bottleneck in sq.mapping_instructions: sq.mapping_instructions.pop(Bottleneck) - resnet_sg: nn.Module = super_gradients.training.models.get("resnet50", pretrained_weights="imagenet", num_classes=1000) + resnet_sg: nn.Module = super_gradients.training.models.get(Models.RESNET50, pretrained_weights="imagenet", num_classes=1000) sq.quantize_module(resnet_sg, preserve_state_dict=True) # PYTORCH-QUANTIZATION @@ -803,7 +805,7 @@ def test_sg_resnet_sg_vanilla_quantization_matches_pytorch_quantization(self): quant_nn.QuantAdaptiveAvgPool2d.set_default_quant_desc_input(QuantDescriptor(calib_method="histogram")) quant_modules.initialize() - resnet_pyquant: nn.Module = super_gradients.training.models.get("resnet50", pretrained_weights="imagenet", num_classes=1000) + resnet_pyquant: nn.Module = super_gradients.training.models.get(Models.RESNET50, pretrained_weights="imagenet", num_classes=1000) quant_modules.deactivate() diff --git a/tests/unit_tests/test_without_train_test.py b/tests/unit_tests/test_without_train_test.py index 2bb48b4d63..8e6457e296 100644 --- a/tests/unit_tests/test_without_train_test.py +++ b/tests/unit_tests/test_without_train_test.py @@ -8,6 +8,7 @@ from super_gradients.training.metrics.detection_metrics import DetectionMetrics from super_gradients.training.metrics.segmentation_metrics import PixelAccuracy, IoU from super_gradients.training.models.detection_models.yolo_base import YoloPostPredictionCallback +from super_gradients.common.object_names import Models class TestWithoutTrainTest(unittest.TestCase): @@ -26,7 +27,7 @@ def tearDownClass(cls) -> None: @staticmethod def get_classification_trainer(name=""): trainer = Trainer(name) - model = models.get("resnet18", num_classes=5) + model = models.get(Models.RESNET18, num_classes=5) return trainer, model @staticmethod From 77f5c8e2501990ed2f41f1337901b9e966b67790 Mon Sep 17 00:00:00 2001 From: shayaharon Date: Tue, 10 Jan 2023 15:07:13 +0200 Subject: [PATCH 4/7] lint and black --- README.md | 9 ++- docs/_sources/welcome.md.txt | 10 ++- ...er_Vision_Models_Pretrained_Checkpoints.md | 4 +- .../ema_train_integration_test.py | 3 +- .../pretrained_models_test.py | 70 ++++++++++--------- tests/unit_tests/detection_utils_test.py | 3 +- tests/unit_tests/export_onnx_test.py | 3 +- .../unit_tests/pretrained_models_unit_test.py | 4 +- tests/unit_tests/test_without_train_test.py | 4 +- tests/unit_tests/vit_unit_test.py | 3 +- 10 files changed, 68 insertions(+), 45 deletions(-) diff --git a/README.md b/README.md index 75bc2aa5a1..d76352c64b 100644 --- a/README.md +++ b/README.md @@ -44,7 +44,9 @@ ________________________________________________________________________________ ```python # Load model with pretrained weights from super_gradients.training import models -model = models.get("yolox_s", pretrained_weights="coco") +from super_gradients.common.object_names import Models + +model = models.get(Models.YOLOX_S, pretrained_weights="coco") ``` #### All Computer Vision Models - Pretrained Checkpoints can be found in the [Model Zoo](http://bit.ly/3EGfKD4) @@ -81,7 +83,10 @@ More example on how and why to use recipes can be found in [Recipes](#recipes) All SuperGradients models’ are production ready in the sense that they are compatible with deployment tools such as TensorRT (Nvidia) and OpenVINO (Intel) and can be easily taken into production. With a few lines of code you can easily integrate the models into your codebase. ```python # Load model with pretrained weights -model = models.get("yolox_s", pretrained_weights="coco") +from super_gradients.training import models +from super_gradients.common.object_names import Models + +model = models.get(Models.YOLOX_S, pretrained_weights="coco") # Prepare model for conversion # Input size is in format of [Batch x Channels x Width x Height] where 640 is the standart COCO dataset dimensions diff --git a/docs/_sources/welcome.md.txt b/docs/_sources/welcome.md.txt index 3c4e0b22ec..52c8e6433c 100644 --- a/docs/_sources/welcome.md.txt +++ b/docs/_sources/welcome.md.txt @@ -46,7 +46,10 @@ ________________________________________________________________________________ ### Ready to deploy pre-trained SOTA models ```python # Load model with pretrained weights -model = models.get("yolox_s", pretrained_weights="coco") +from super_gradients.common.object_names import Models +from super_gradients.training import models + +model = models.get(Models.YOLOX_S, pretrained_weights="coco") ``` #### Classification @@ -86,7 +89,10 @@ More example on how and why to use recipes can be found in [Recipes](#recipes) All SuperGradients models’ are production ready in the sense that they are compatible with deployment tools such as TensorRT (Nvidia) and OpenVINO (Intel) and can be easily taken into production. With a few lines of code you can easily integrate the models into your codebase. ```python # Load model with pretrained weights -model = models.get("yolox_s", pretrained_weights="coco") +from super_gradients.training import models +from super_gradients.common.object_names import Models + +model = models.get(Models.YOLOX_S, pretrained_weights="coco") # Prepare model for conversion # Input size is in format of [Batch x Channels x Width x Height] where 640 is the standart COCO dataset dimensions diff --git a/src/super_gradients/training/Computer_Vision_Models_Pretrained_Checkpoints.md b/src/super_gradients/training/Computer_Vision_Models_Pretrained_Checkpoints.md index 9c2f08ae52..1027e7a089 100644 --- a/src/super_gradients/training/Computer_Vision_Models_Pretrained_Checkpoints.md +++ b/src/super_gradients/training/Computer_Vision_Models_Pretrained_Checkpoints.md @@ -4,7 +4,9 @@ You can load any of our pretrained model in 2 lines of code: ```python from super_gradients.training import models -model = models.get("yolox_s", pretrained_weights="coco") +from super_gradients.common.object_names import Models + +model = models.get(Models.YOLOX_S, pretrained_weights="coco") ``` All the available models are listed in the column `Model name`. diff --git a/tests/integration_tests/ema_train_integration_test.py b/tests/integration_tests/ema_train_integration_test.py index 0b93fabb80..2c5ed67340 100644 --- a/tests/integration_tests/ema_train_integration_test.py +++ b/tests/integration_tests/ema_train_integration_test.py @@ -1,3 +1,4 @@ +from super_gradients.common.object_names import Models from super_gradients.training import models from super_gradients.training import Trainer from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader @@ -22,7 +23,7 @@ def __call__(self, *args, **kwargs): class EMAIntegrationTest(unittest.TestCase): def _init_model(self) -> None: self.trainer = Trainer("resnet18_cifar_ema_test") - self.model = models.get("resnet18_cifar", arch_params={"num_classes": 5}) + self.model = models.get(Models.RESNET18_CIFAR, arch_params={"num_classes": 5}) @classmethod def tearDownClass(cls) -> None: diff --git a/tests/integration_tests/pretrained_models_test.py b/tests/integration_tests/pretrained_models_test.py index d489cd8ca6..a9ea1efd5b 100644 --- a/tests/integration_tests/pretrained_models_test.py +++ b/tests/integration_tests/pretrained_models_test.py @@ -92,7 +92,7 @@ def setUp(self) -> None: "greater_metric_to_watch_is_better": True, } self.coco_pretrained_ckpt_params = {"pretrained_weights": "coco"} - self.coco_pretrained_arch_params = {"ssd_lite_mobilenet_v2": {"num_classes": 80}, "coco_ssd_mobilenet_v1": {"num_classes": 80}} + self.coco_pretrained_arch_params = {Models.SSD_LITE_MOBILENET_V2: {"num_classes": 80}, Models.SSD_MOBILENET_V1: {"num_classes": 80}} self.coco_pretrained_ckpt_params = {"pretrained_weights": "coco"} self.coco_dataset = { @@ -103,8 +103,8 @@ def setUp(self) -> None: } self.coco_pretrained_maps = { - "ssd_lite_mobilenet_v2": 0.2041, - "coco_ssd_mobilenet_v1": 0.243, + Models.SSD_LITE_MOBILENET_V2: 0.2041, + Models.SSD_MOBILENET_V1: 0.243, Models.YOLOX_S: 0.4047, Models.YOLOX_M: 0.4640, Models.YOLOX_L: 0.4925, @@ -152,9 +152,9 @@ def setUp(self) -> None: self.coco_segmentation_subclass_pretrained_mious = {"shelfnet34_lw": 0.651} self.coco_segmentation_dataset = coco_segmentation_val() - self.cityscapes_pretrained_models = ["ddrnet_23", "ddrnet_23_slim", Models.STDC1_SEG50, Models.REGSEG48] + self.cityscapes_pretrained_models = [Models.DDRNET_23, Models.DDRNET_23_SLIM, Models.STDC1_SEG50, Models.REGSEG48] self.cityscapes_pretrained_arch_params = { - "ddrnet_23": {"aux_head": True}, + Models.DDRNET_23: {"aux_head": True}, Models.REGSEG48: {}, "stdc": {"use_aux_heads": True, "aux_head": True}, "pplite_seg": {"use_aux_heads": True}, @@ -162,17 +162,17 @@ def setUp(self) -> None: self.cityscapes_pretrained_ckpt_params = {"pretrained_weights": "cityscapes"} self.cityscapes_pretrained_mious = { - "ddrnet_23": 0.8026, - "ddrnet_23_slim": 0.7801, + Models.DDRNET_23: 0.8026, + Models.DDRNET_23_SLIM: 0.7801, Models.STDC1_SEG50: 0.7511, Models.STDC1_SEG75: 0.7687, Models.STDC2_SEG50: 0.7644, Models.STDC2_SEG75: 0.7893, Models.REGSEG48: 0.7815, - "pp_lite_t_seg50": 0.7492, - "pp_lite_t_seg75": 0.7756, - "pp_lite_b_seg50": 0.7648, - "pp_lite_b_seg75": 0.7852, + Models.PP_LITE_T_SEG50: 0.7492, + Models.PP_LITE_T_SEG75: 0.7756, + Models.PP_LITE_B_SEG50: 0.7648, + Models.PP_LITE_B_SEG75: 0.7852, } self.cityscapes_dataset = cityscapes_val() @@ -389,7 +389,7 @@ def test_transfer_learning_regseg48_cityscapes(self): def test_pretrained_ddrnet23_cityscapes(self): trainer = Trainer("cityscapes_pretrained_ddrnet23") - model = models.get("ddrnet_23", arch_params=self.cityscapes_pretrained_arch_params["ddrnet_23"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.DDRNET_23, arch_params=self.cityscapes_pretrained_arch_params[Models.DDRNET_23], **self.cityscapes_pretrained_ckpt_params) res = ( trainer.test( model=model, test_loader=self.cityscapes_dataset, test_metrics_list=[IoU(num_classes=20, ignore_index=19)], metrics_progress_verbose=True @@ -397,11 +397,13 @@ def test_pretrained_ddrnet23_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["ddrnet_23"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.DDRNET_23], delta=0.001) def test_pretrained_ddrnet23_slim_cityscapes(self): trainer = Trainer("cityscapes_pretrained_ddrnet23_slim") - model = models.get("ddrnet_23_slim", arch_params=self.cityscapes_pretrained_arch_params["ddrnet_23"], **self.cityscapes_pretrained_ckpt_params) + model = models.get( + Models.DDRNET_23_SLIM, arch_params=self.cityscapes_pretrained_arch_params[Models.DDRNET_23], **self.cityscapes_pretrained_ckpt_params + ) res = ( trainer.test( model=model, test_loader=self.cityscapes_dataset, test_metrics_list=[IoU(num_classes=20, ignore_index=19)], metrics_progress_verbose=True @@ -409,11 +411,11 @@ def test_pretrained_ddrnet23_slim_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["ddrnet_23_slim"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.DDRNET_23_SLIM], delta=0.001) def test_transfer_learning_ddrnet23_cityscapes(self): trainer = Trainer("cityscapes_pretrained_ddrnet23_transfer_learning") - model = models.get("ddrnet_23", arch_params=self.cityscapes_pretrained_arch_params["ddrnet_23"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.DDRNET_23, arch_params=self.cityscapes_pretrained_arch_params[Models.DDRNET_23], **self.cityscapes_pretrained_ckpt_params) trainer.train( model=model, training_params=self.ddrnet_transfer_segmentation_train_params, @@ -423,7 +425,9 @@ def test_transfer_learning_ddrnet23_cityscapes(self): def test_transfer_learning_ddrnet23_slim_cityscapes(self): trainer = Trainer("cityscapes_pretrained_ddrnet23_slim_transfer_learning") - model = models.get("ddrnet_23_slim", arch_params=self.cityscapes_pretrained_arch_params["ddrnet_23"], **self.cityscapes_pretrained_ckpt_params) + model = models.get( + Models.DDRNET_23_SLIM, arch_params=self.cityscapes_pretrained_arch_params[Models.DDRNET_23], **self.cityscapes_pretrained_ckpt_params + ) trainer.train( model=model, training_params=self.ddrnet_transfer_segmentation_train_params, @@ -468,7 +472,9 @@ def test_transfer_learning_efficientnet_b0_imagenet(self): def test_pretrained_ssd_lite_mobilenet_v2_coco(self): trainer = Trainer("coco_ssd_lite_mobilenet_v2") - model = models.get("ssd_lite_mobilenet_v2", arch_params=self.coco_pretrained_arch_params["ssd_lite_mobilenet_v2"], **self.coco_pretrained_ckpt_params) + model = models.get( + Models.SSD_LITE_MOBILENET_V2, arch_params=self.coco_pretrained_arch_params[Models.SSD_LITE_MOBILENET_V2], **self.coco_pretrained_ckpt_params + ) ssd_post_prediction_callback = SSDPostPredictCallback() res = trainer.test( model=model, @@ -476,13 +482,13 @@ def test_pretrained_ssd_lite_mobilenet_v2_coco(self): test_metrics_list=[DetectionMetrics(post_prediction_callback=ssd_post_prediction_callback, num_cls=80)], metrics_progress_verbose=True, )[2] - self.assertAlmostEqual(res, self.coco_pretrained_maps["ssd_lite_mobilenet_v2"], delta=0.001) + self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.SSD_LITE_MOBILENET_V2], delta=0.001) def test_transfer_learning_ssd_lite_mobilenet_v2_coco(self): trainer = Trainer("coco_ssd_lite_mobilenet_v2_transfer_learning") - transfer_arch_params = self.coco_pretrained_arch_params["ssd_lite_mobilenet_v2"].copy() + transfer_arch_params = self.coco_pretrained_arch_params[Models.SSD_LITE_MOBILENET_V2].copy() transfer_arch_params["num_classes"] = 5 - model = models.get("ssd_lite_mobilenet_v2", arch_params=transfer_arch_params, **self.coco_pretrained_ckpt_params) + model = models.get(Models.SSD_LITE_MOBILENET_V2, arch_params=transfer_arch_params, **self.coco_pretrained_ckpt_params) trainer.train( model=model, training_params=self.transfer_detection_train_params_ssd, @@ -491,8 +497,8 @@ def test_transfer_learning_ssd_lite_mobilenet_v2_coco(self): ) def test_pretrained_ssd_mobilenet_v1_coco(self): - trainer = Trainer("coco_ssd_mobilenet_v1") - model = models.get(Models.SSD_MOBILENET_V1, arch_params=self.coco_pretrained_arch_params["coco_ssd_mobilenet_v1"], **self.coco_pretrained_ckpt_params) + trainer = Trainer(Models.SSD_MOBILENET_V1) + model = models.get(Models.SSD_MOBILENET_V1, arch_params=self.coco_pretrained_arch_params[Models.SSD_MOBILENET_V1], **self.coco_pretrained_ckpt_params) ssd_post_prediction_callback = SSDPostPredictCallback() res = trainer.test( model=model, @@ -500,7 +506,7 @@ def test_pretrained_ssd_mobilenet_v1_coco(self): test_metrics_list=[DetectionMetrics(post_prediction_callback=ssd_post_prediction_callback, num_cls=80)], metrics_progress_verbose=True, )[2] - self.assertAlmostEqual(res, self.coco_pretrained_maps["coco_ssd_mobilenet_v1"], delta=0.001) + self.assertAlmostEqual(res, self.coco_pretrained_maps[Models.SSD_MOBILENET_V1], delta=0.001) def test_pretrained_yolox_s_coco(self): trainer = Trainer(Models.YOLOX_S) @@ -809,7 +815,7 @@ def test_transfer_learning_beit_base_imagenet(self): def test_pretrained_pplite_t_seg50_cityscapes(self): trainer = Trainer("cityscapes_pretrained_pplite_t_seg50") - model = models.get("pp_lite_t_seg50", arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.PP_LITE_T_SEG50, arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params) res = ( trainer.test( @@ -821,11 +827,11 @@ def test_pretrained_pplite_t_seg50_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["pp_lite_t_seg50"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.PP_LITE_T_SEG50], delta=0.001) def test_pretrained_pplite_t_seg75_cityscapes(self): trainer = Trainer("cityscapes_pretrained_pplite_t_seg75") - model = models.get("pp_lite_t_seg75", arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.PP_LITE_T_SEG75, arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params) res = ( trainer.test( @@ -837,11 +843,11 @@ def test_pretrained_pplite_t_seg75_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["pp_lite_t_seg75"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.PP_LITE_T_SEG75], delta=0.001) def test_pretrained_pplite_b_seg50_cityscapes(self): trainer = Trainer("cityscapes_pretrained_pplite_b_seg50") - model = models.get("pp_lite_b_seg50", arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.PP_LITE_B_SEG50, arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params) res = ( trainer.test( @@ -853,11 +859,11 @@ def test_pretrained_pplite_b_seg50_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["pp_lite_b_seg50"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.PP_LITE_B_SEG50], delta=0.001) def test_pretrained_pplite_b_seg75_cityscapes(self): trainer = Trainer("cityscapes_pretrained_pplite_b_seg75") - model = models.get("pp_lite_b_seg75", arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params) + model = models.get(Models.PP_LITE_B_SEG75, arch_params=self.cityscapes_pretrained_arch_params["pplite_seg"], **self.cityscapes_pretrained_ckpt_params) res = ( trainer.test( @@ -869,7 +875,7 @@ def test_pretrained_pplite_b_seg75_cityscapes(self): .cpu() .item() ) - self.assertAlmostEqual(res, self.cityscapes_pretrained_mious["pp_lite_b_seg75"], delta=0.001) + self.assertAlmostEqual(res, self.cityscapes_pretrained_mious[Models.PP_LITE_B_SEG75], delta=0.001) def tearDown(self) -> None: if os.path.exists("~/.cache/torch/hub/"): diff --git a/tests/unit_tests/detection_utils_test.py b/tests/unit_tests/detection_utils_test.py index cd3d253a9a..5c51c7de8f 100644 --- a/tests/unit_tests/detection_utils_test.py +++ b/tests/unit_tests/detection_utils_test.py @@ -4,6 +4,7 @@ import numpy as np import torch.cuda +from super_gradients.common.object_names import Models from super_gradients.training import Trainer, utils as core_utils, models from super_gradients.training.dataloaders.dataloaders import coco2017_val from super_gradients.training.datasets.datasets_conf import COCO_DETECTION_CLASSES_LIST @@ -16,7 +17,7 @@ class TestDetectionUtils(unittest.TestCase): def setUp(self): self.device = "cuda" if torch.cuda.is_available() else "cpu" - self.model = models.get("yolox_n", pretrained_weights="coco").to(self.device) + self.model = models.get(Models.YOLOX_N, pretrained_weights="coco").to(self.device) self.model.eval() @unittest.skipIf(not is_data_available(), "run only when /data is available") diff --git a/tests/unit_tests/export_onnx_test.py b/tests/unit_tests/export_onnx_test.py index a51bd0e093..33e4b53019 100644 --- a/tests/unit_tests/export_onnx_test.py +++ b/tests/unit_tests/export_onnx_test.py @@ -1,6 +1,7 @@ import tempfile import unittest +from super_gradients.common.object_names import Models from super_gradients.training import models from torchvision.transforms import Compose, Normalize, Resize from super_gradients.training.transforms import Standardize @@ -9,7 +10,7 @@ class TestModelsONNXExport(unittest.TestCase): def test_models_onnx_export(self): - pretrained_model = models.get("resnet18", num_classes=1000, pretrained_weights="imagenet") + pretrained_model = models.get(Models.RESNET18, num_classes=1000, pretrained_weights="imagenet") preprocess = Compose([Resize(224), Standardize(), Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])]) with tempfile.TemporaryDirectory() as tmpdirname: out_path = os.path.join(tmpdirname, "resnet18.onnx") diff --git a/tests/unit_tests/pretrained_models_unit_test.py b/tests/unit_tests/pretrained_models_unit_test.py index 85b03fd31a..bb6a503c9b 100644 --- a/tests/unit_tests/pretrained_models_unit_test.py +++ b/tests/unit_tests/pretrained_models_unit_test.py @@ -21,12 +21,12 @@ def test_pretrained_resnet50_imagenet(self): def test_pretrained_regnetY800_imagenet(self): trainer = Trainer("imagenet_pretrained_regnetY800_unit_test") - model = models.get("regnetY800", pretrained_weights="imagenet") + model = models.get(Models.REGNETY800, pretrained_weights="imagenet") 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") - model = models.get("repvgg_a0", pretrained_weights="imagenet", arch_params={"build_residual_branches": True}) + model = models.get(Models.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) def tearDown(self) -> None: diff --git a/tests/unit_tests/test_without_train_test.py b/tests/unit_tests/test_without_train_test.py index 8e6457e296..8c896756e6 100644 --- a/tests/unit_tests/test_without_train_test.py +++ b/tests/unit_tests/test_without_train_test.py @@ -33,14 +33,14 @@ def get_classification_trainer(name=""): @staticmethod def get_detection_trainer(name=""): trainer = Trainer(name) - model = models.get("yolox_s", num_classes=5) + model = models.get(Models.YOLOX_S, num_classes=5) return trainer, model @staticmethod 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(Models.SHELFNET34_LW, arch_params=shelfnet_lw_arch_params) return trainer, model def test_test_without_train(self): diff --git a/tests/unit_tests/vit_unit_test.py b/tests/unit_tests/vit_unit_test.py index a599916dd1..a943671abb 100644 --- a/tests/unit_tests/vit_unit_test.py +++ b/tests/unit_tests/vit_unit_test.py @@ -1,5 +1,6 @@ import unittest +from super_gradients.common.object_names import Models from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader from super_gradients import Trainer from super_gradients.training.metrics import Accuracy, Top5 @@ -31,7 +32,7 @@ def test_train_vit(self): Validate vit_base """ trainer = Trainer("test_vit_base") - model = models.get("vit_base", arch_params=self.arch_params, num_classes=5) + model = models.get(Models.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() ) From 5512aec68e99d56b3d2868de86ccfa6807060323 Mon Sep 17 00:00:00 2001 From: shayaharon Date: Tue, 10 Jan 2023 15:08:07 +0200 Subject: [PATCH 5/7] welcome.md ref replaced --- documentation/source/welcome.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/documentation/source/welcome.md b/documentation/source/welcome.md index 68c2742b61..cd3023d456 100644 --- a/documentation/source/welcome.md +++ b/documentation/source/welcome.md @@ -131,9 +131,10 @@ Want to try our pre-trained models on your machine? Import SuperGradients, initi import super_gradients from super_gradients.training import Trainer, models, dataloaders +from super_gradients.common.object_names import Models trainer = Trainer(experiment_name="yoloxn_coco_experiment",ckpt_root_dir=) -model = models.get("yolox_n", pretrained_weights="coco", num_classes= 80) +model = models.get(Models.YOLOX_N, pretrained_weights="coco", num_classes= 80) train_loader = dataloaders.coco2017_train() valid_loader = dataloaders.coco2017_val() train_params = {...} From 00af7312eceea33601ee117d86e3f3f2a3d1e906 Mon Sep 17 00:00:00 2001 From: shayaharon Date: Tue, 10 Jan 2023 16:59:20 +0200 Subject: [PATCH 6/7] ci fix --- .circleci/config.yml | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/.circleci/config.yml b/.circleci/config.yml index aef3956b78..b9b0954b6e 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -325,9 +325,7 @@ jobs: steps: - deci-common/checkout_and_skip_build: check_version_file: true - - run: - name: verify __version__ - command: python3 tests/verify_version.py $CIRCLE_TAG + - run: name: add requirements.txt and requirements.pro.txt to source code command: | @@ -344,6 +342,13 @@ jobs: - deci-common/pip_test_package_installation_codeartifact_dev: package_name: "super-gradients" version: $CIRCLE_TAG + venv_name: dev-sg-test + + - run: + name: verify __version__ + command: | + source + python3 tests/verify_version.py $CIRCLE_TAG - deci-common/pip_test_package_installation_codeartifact_prod: package_name: "super-gradients" From b51047b3c8639b35a7e7480efe8c60b83447b24a Mon Sep 17 00:00:00 2001 From: shayaharon Date: Sun, 22 Jan 2023 17:06:24 +0200 Subject: [PATCH 7/7] regseg typo fix --- .../regseg_transfer_learning_example.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/super_gradients/examples/regseg_transfer_learning_example/regseg_transfer_learning_example.py b/src/super_gradients/examples/regseg_transfer_learning_example/regseg_transfer_learning_example.py index aed420ece0..e8bd1b167e 100644 --- a/src/super_gradients/examples/regseg_transfer_learning_example/regseg_transfer_learning_example.py +++ b/src/super_gradients/examples/regseg_transfer_learning_example/regseg_transfer_learning_example.py @@ -34,7 +34,7 @@ # THIS IS WHERE THE MAGIC HAPPENS- SINCE TRAINER'S CLASSES ATTRIBUTE WAS SET TO BE DIFFERENT FROM CITYSCAPES'S, AFTER # LOADING THE PRETRAINED REGSET, IT WILL CALL IT'S REPLACE_HEAD METHOD AND CHANGE IT'S SEGMENTATION HEAD LAYER ACCORDING # TO OUR BINARY SEGMENTATION DATASET -model = models.get(Models.RREGSEG48, pretrained_weights="cityscapes", num_classes=1) +model = models.get(Models.REGSEG48, pretrained_weights="cityscapes", num_classes=1) # DEFINE TRAINING PARAMS. SEE DOCS FOR THE FULL LIST. train_params = {