Feature/sg 493 modelnames instead of strings (#614)

* refs in unit tests updated

* formatting

* more models import fixes

* lint and black

* welcome.md ref replaced

* ci fix

* regseg typo fix

Co-authored-by: Ofri Masad <ofrimasad@users.noreply.github.com>

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

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

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=<CHECKPOINT_DIRECTORY>)
-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 = {...}

src/super_gradients/examples/ddrnet_imagenet/ddrnet_classification_example.py

@@ -15,67 +15,72 @@
 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
 from super_gradients.training import Trainer, models, dataloaders
 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()
 
 parser.add_argument("--reload", action="store_true")
 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)

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())

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,

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],

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():

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.REGSEG48, 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)

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`.

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={