Skip to content

Commit 1f6a183

Browse files
shaydeciofrimasad
shaydeci
and
ofrimasad
authored
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 <[email protected]>
1 parent 8f7f220 commit 1f6a183

35 files changed

+403
-306
lines changed

‎README.md

+7-2
Original file line numberDiff line numberDiff line change
@@ -44,7 +44,9 @@ ________________________________________________________________________________
4444
```python
4545
# Load model with pretrained weights
4646
from super_gradients.training import models
47-
model = models.get("yolox_s", pretrained_weights="coco")
47+
from super_gradients.common.object_names import Models
48+
49+
model = models.get(Models.YOLOX_S, pretrained_weights="coco")
4850
```
4951
#### All Computer Vision Models - Pretrained Checkpoints can be found in the [Model Zoo](http://bit.ly/3EGfKD4)
5052

@@ -81,7 +83,10 @@ More example on how and why to use recipes can be found in [Recipes](#recipes)
8183
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.
8284
```python
8385
# Load model with pretrained weights
84-
model = models.get("yolox_s", pretrained_weights="coco")
86+
from super_gradients.training import models
87+
from super_gradients.common.object_names import Models
88+
89+
model = models.get(Models.YOLOX_S, pretrained_weights="coco")
8590

8691
# Prepare model for conversion
8792
# 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

+8-2
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,10 @@ ________________________________________________________________________________
4646
### Ready to deploy pre-trained SOTA models
4747
```python
4848
# Load model with pretrained weights
49-
model = models.get("yolox_s", pretrained_weights="coco")
49+
from super_gradients.common.object_names import Models
50+
from super_gradients.training import models
51+
52+
model = models.get(Models.YOLOX_S, pretrained_weights="coco")
5053
```
5154

5255
#### Classification
@@ -86,7 +89,10 @@ More example on how and why to use recipes can be found in [Recipes](#recipes)
8689
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.
8790
```python
8891
# Load model with pretrained weights
89-
model = models.get("yolox_s", pretrained_weights="coco")
92+
from super_gradients.training import models
93+
from super_gradients.common.object_names import Models
94+
95+
model = models.get(Models.YOLOX_S, pretrained_weights="coco")
9096

9197
# Prepare model for conversion
9298
# 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

+2-1
Original file line numberDiff line numberDiff line change
@@ -131,9 +131,10 @@ Want to try our pre-trained models on your machine? Import SuperGradients, initi
131131

132132
import super_gradients
133133
from super_gradients.training import Trainer, models, dataloaders
134+
from super_gradients.common.object_names import Models
134135

135136
trainer = Trainer(experiment_name="yoloxn_coco_experiment",ckpt_root_dir=<CHECKPOINT_DIRECTORY>)
136-
model = models.get("yolox_n", pretrained_weights="coco", num_classes= 80)
137+
model = models.get(Models.YOLOX_N, pretrained_weights="coco", num_classes= 80)
137138
train_loader = dataloaders.coco2017_train()
138139
valid_loader = dataloaders.coco2017_val()
139140
train_params = {...}

‎src/super_gradients/examples/ddrnet_imagenet/ddrnet_classification_example.py

+46-41
Original file line numberDiff line numberDiff line change
@@ -15,67 +15,72 @@
1515
import torch
1616

1717
from super_gradients.common import MultiGPUMode
18+
from super_gradients.common.object_names import Models
1819
from super_gradients.training.datasets.datasets_utils import RandomResizedCropAndInterpolation
1920
from torchvision.transforms import RandomHorizontalFlip, ColorJitter, ToTensor, Normalize
2021
import super_gradients
2122
from super_gradients.training import Trainer, models, dataloaders
2223
import argparse
2324
from super_gradients.training.metrics import Accuracy, Top5
2425
from super_gradients.training.datasets.data_augmentation import RandomErase
26+
2527
parser = argparse.ArgumentParser()
2628
super_gradients.init_trainer()
2729

2830
parser.add_argument("--reload", action="store_true")
2931
parser.add_argument("--max_epochs", type=int, default=100)
3032
parser.add_argument("--batch", type=int, default=3)
3133
parser.add_argument("--experiment_name", type=str, default="ddrnet_23")
32-
parser.add_argument("-s", "--slim", action="store_true", help='train the slim version of DDRNet23')
34+
parser.add_argument("-s", "--slim", action="store_true", help="train the slim version of DDRNet23")
3335

3436
args, _ = parser.parse_known_args()
3537
distributed = super_gradients.is_distributed()
3638
devices = torch.cuda.device_count() if not distributed else 1
3739

38-
train_params_ddr = {"max_epochs": args.max_epochs,
39-
"lr_mode": "step",
40-
"lr_updates": [30, 60, 90],
41-
"lr_decay_factor": 0.1,
42-
"initial_lr": 0.1 * devices,
43-
"optimizer": "SGD",
44-
"optimizer_params": {"weight_decay": 0.0001, "momentum": 0.9, "nesterov": True},
45-
"loss": "cross_entropy",
46-
"train_metrics_list": [Accuracy(), Top5()],
47-
"valid_metrics_list": [Accuracy(), Top5()],
48-
49-
"metric_to_watch": "Accuracy",
50-
"greater_metric_to_watch_is_better": True
51-
}
52-
53-
dataset_params = {"batch_size": args.batch,
54-
"color_jitter": 0.4,
55-
"random_erase_prob": 0.2,
56-
"random_erase_value": 'random',
57-
"train_interpolation": 'random',
58-
}
59-
60-
61-
train_transforms = [RandomResizedCropAndInterpolation(size=224, interpolation="random"),
62-
RandomHorizontalFlip(),
63-
ColorJitter(0.4, 0.4, 0.4),
64-
ToTensor(),
65-
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
66-
RandomErase(0.2, "random")
67-
]
68-
69-
trainer = Trainer(experiment_name=args.experiment_name,
70-
multi_gpu=MultiGPUMode.DISTRIBUTED_DATA_PARALLEL if distributed else MultiGPUMode.DATA_PARALLEL,
71-
device='cuda')
72-
73-
train_loader = dataloaders.imagenet_train(dataset_params={"transforms": train_transforms},
74-
dataloader_params={"batch_size": args.batch})
40+
train_params_ddr = {
41+
"max_epochs": args.max_epochs,
42+
"lr_mode": "step",
43+
"lr_updates": [30, 60, 90],
44+
"lr_decay_factor": 0.1,
45+
"initial_lr": 0.1 * devices,
46+
"optimizer": "SGD",
47+
"optimizer_params": {"weight_decay": 0.0001, "momentum": 0.9, "nesterov": True},
48+
"loss": "cross_entropy",
49+
"train_metrics_list": [Accuracy(), Top5()],
50+
"valid_metrics_list": [Accuracy(), Top5()],
51+
"metric_to_watch": "Accuracy",
52+
"greater_metric_to_watch_is_better": True,
53+
}
54+
55+
dataset_params = {
56+
"batch_size": args.batch,
57+
"color_jitter": 0.4,
58+
"random_erase_prob": 0.2,
59+
"random_erase_value": "random",
60+
"train_interpolation": "random",
61+
}
62+
63+
64+
train_transforms = [
65+
RandomResizedCropAndInterpolation(size=224, interpolation="random"),
66+
RandomHorizontalFlip(),
67+
ColorJitter(0.4, 0.4, 0.4),
68+
ToTensor(),
69+
Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
70+
RandomErase(0.2, "random"),
71+
]
72+
73+
trainer = Trainer(
74+
experiment_name=args.experiment_name, multi_gpu=MultiGPUMode.DISTRIBUTED_DATA_PARALLEL if distributed else MultiGPUMode.DATA_PARALLEL, device="cuda"
75+
)
76+
77+
train_loader = dataloaders.imagenet_train(dataset_params={"transforms": train_transforms}, dataloader_params={"batch_size": args.batch})
7578
valid_loader = dataloaders.imagenet_val()
7679

77-
model = models.get("ddrnet_23_slim" if args.slim else "ddrnet_23",
78-
arch_params={"aux_head": False, "classification_mode": True, 'dropout_prob': 0.3},
79-
num_classes=1000)
80+
model = models.get(
81+
Models.DDRNET_23_SLIM if args.slim else Models.DDRNET_23,
82+
arch_params={"aux_head": False, "classification_mode": True, "dropout_prob": 0.3},
83+
num_classes=1000,
84+
)
8085

8186
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

+20-9
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22
# Reaches ~94.9 Accuracy after 250 Epochs
33
import super_gradients
44
from super_gradients import Trainer
5+
from super_gradients.common.object_names import Models
56
from super_gradients.training import models, dataloaders
67
from super_gradients.training.metrics.classification_metrics import Accuracy, Top5
78
from super_gradients.training.utils.early_stopping import EarlyStop
@@ -13,18 +14,28 @@
1314
early_stop_acc = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="Accuracy", mode="max", patience=3, verbose=True)
1415
early_stop_val_loss = EarlyStop(Phase.VALIDATION_EPOCH_END, monitor="LabelSmoothingCrossEntropyLoss", mode="min", patience=3, verbose=True)
1516

16-
train_params = {"max_epochs": 250, "lr_updates": [100, 150, 200], "lr_decay_factor": 0.1, "lr_mode": "step",
17-
"lr_warmup_epochs": 0, "initial_lr": 0.1, "loss": "cross_entropy", "optimizer": "SGD",
18-
"criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9},
19-
"train_metrics_list": [Accuracy(), Top5()], "valid_metrics_list": [Accuracy(), Top5()],
20-
"metric_to_watch": "Accuracy",
21-
"greater_metric_to_watch_is_better": True, "phase_callbacks": [early_stop_acc, early_stop_val_loss]}
17+
train_params = {
18+
"max_epochs": 250,
19+
"lr_updates": [100, 150, 200],
20+
"lr_decay_factor": 0.1,
21+
"lr_mode": "step",
22+
"lr_warmup_epochs": 0,
23+
"initial_lr": 0.1,
24+
"loss": "cross_entropy",
25+
"optimizer": "SGD",
26+
"criterion_params": {},
27+
"optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9},
28+
"train_metrics_list": [Accuracy(), Top5()],
29+
"valid_metrics_list": [Accuracy(), Top5()],
30+
"metric_to_watch": "Accuracy",
31+
"greater_metric_to_watch_is_better": True,
32+
"phase_callbacks": [early_stop_acc, early_stop_val_loss],
33+
}
2234

2335
# Define Model
2436
trainer = Trainer("Callback_Example")
2537

2638
# Build Model
27-
model = models.get("resnet18_cifar", num_classes=10)
39+
model = models.get(Models.RESNET18_CIFAR, num_classes=10)
2840

29-
trainer.train(model=model, training_params=train_params,
30-
train_loader=dataloaders.cifar10_train(), valid_loader=dataloaders.cifar10_val())
41+
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

+2-1
Original file line numberDiff line numberDiff line change
@@ -1,10 +1,11 @@
1+
from super_gradients.common.object_names import Models
12
from super_gradients.training import Trainer, models
23
from super_gradients.training.metrics.classification_metrics import Accuracy, Top5
34
from super_gradients.training.dataloaders.dataloaders import cifar10_train, cifar10_val
45

56

67
trainer = Trainer(experiment_name="demo-clearml-logger")
7-
model = models.get("resnet18", num_classes=10)
8+
model = models.get(Models.RESNET18, num_classes=10)
89

910
training_params = {
1011
"max_epochs": 20,

‎src/super_gradients/examples/loggers_examples/deci_platform_logger_example.py

+3-1
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,6 @@
11
import os
2+
3+
from super_gradients.common.object_names import Models
24
from super_gradients.training import Trainer, models
35
from super_gradients.training.metrics.classification_metrics import Accuracy, Top5
46
from super_gradients.training.dataloaders.dataloaders import cifar10_train, cifar10_val
@@ -7,7 +9,7 @@
79

810

911
trainer = Trainer(experiment_name="demo-deci-platform-logger")
10-
model = models.get("resnet18", num_classes=10)
12+
model = models.get(Models.RESNET18, num_classes=10)
1113
training_params = {
1214
"max_epochs": 20,
1315
"lr_updates": [5, 10, 15],

‎src/super_gradients/examples/quantization/resnet_qat_example.py

+2-1
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44

55
import super_gradients
66
from super_gradients import Trainer
7+
from super_gradients.common.object_names import Models
78
from super_gradients.modules.quantization.resnet_bottleneck import QuantBottleneck as sg_QuantizedBottleneck
89
from super_gradients.training import MultiGPUMode
910
from super_gradients.training import models as sg_models
@@ -55,7 +56,7 @@ def selective_quantize(model: nn.Module):
5556

5657

5758
def sg_vanilla_resnet50():
58-
return sg_models.get("resnet50", pretrained_weights="imagenet", num_classes=1000)
59+
return sg_models.get(Models.RESNET50, pretrained_weights="imagenet", num_classes=1000)
5960

6061

6162
def sg_naive_qdq_resnet50():

‎src/super_gradients/examples/regseg_transfer_learning_example/regseg_transfer_learning_example.py

+41-38
Original file line numberDiff line numberDiff line change
@@ -1,20 +1,31 @@
1+
from super_gradients.common.object_names import Models
12
from super_gradients.training import models, dataloaders
23

34
from super_gradients.training.sg_trainer import Trainer
45
from super_gradients.training.metrics import BinaryIOU
5-
from super_gradients.training.transforms.transforms import SegResize, SegRandomFlip, SegRandomRescale, SegCropImageAndMask, \
6-
SegPadShortToCropSize, SegColorJitter
6+
from super_gradients.training.transforms.transforms import (
7+
SegResize,
8+
SegRandomFlip,
9+
SegRandomRescale,
10+
SegCropImageAndMask,
11+
SegPadShortToCropSize,
12+
SegColorJitter,
13+
)
714
from super_gradients.training.utils.callbacks import BinarySegmentationVisualizationCallback, Phase
815

916
# DEFINE DATA TRANSFORMATIONS
1017

1118
dl_train = dataloaders.supervisely_persons_train(
12-
dataset_params={"transforms": [SegColorJitter(brightness=0.5, contrast=0.5, saturation=0.5),
13-
SegRandomFlip(),
14-
SegRandomRescale(scales=[0.25, 1.]),
15-
SegPadShortToCropSize([320, 480]),
16-
SegCropImageAndMask(crop_size=[320, 480],
17-
mode="random")]})
19+
dataset_params={
20+
"transforms": [
21+
SegColorJitter(brightness=0.5, contrast=0.5, saturation=0.5),
22+
SegRandomFlip(),
23+
SegRandomRescale(scales=[0.25, 1.0]),
24+
SegPadShortToCropSize([320, 480]),
25+
SegCropImageAndMask(crop_size=[320, 480], mode="random"),
26+
]
27+
}
28+
)
1829

1930
dl_val = dataloaders.supervisely_persons_val(dataset_params={"transforms": [SegResize(h=480, w=320)]})
2031

@@ -23,35 +34,27 @@
2334
# THIS IS WHERE THE MAGIC HAPPENS- SINCE TRAINER'S CLASSES ATTRIBUTE WAS SET TO BE DIFFERENT FROM CITYSCAPES'S, AFTER
2435
# LOADING THE PRETRAINED REGSET, IT WILL CALL IT'S REPLACE_HEAD METHOD AND CHANGE IT'S SEGMENTATION HEAD LAYER ACCORDING
2536
# TO OUR BINARY SEGMENTATION DATASET
26-
model = models.get("regseg48", pretrained_weights="cityscapes", num_classes=1)
37+
model = models.get(Models.REGSEG48, pretrained_weights="cityscapes", num_classes=1)
2738

2839
# DEFINE TRAINING PARAMS. SEE DOCS FOR THE FULL LIST.
29-
train_params = {"max_epochs": 50,
30-
"lr_mode": "cosine",
31-
"initial_lr": 0.0064, # for batch_size=16
32-
"optimizer_params": {"momentum": 0.843,
33-
"weight_decay": 0.00036,
34-
"nesterov": True},
35-
36-
"cosine_final_lr_ratio": 0.1,
37-
"multiply_head_lr": 10,
38-
"optimizer": "SGD",
39-
"loss": "bce_dice_loss",
40-
"ema": True,
41-
"zero_weight_decay_on_bias_and_bn": True,
42-
"average_best_models": True,
43-
"mixed_precision": False,
44-
"metric_to_watch": "mean_IOU",
45-
"greater_metric_to_watch_is_better": True,
46-
"train_metrics_list": [BinaryIOU()],
47-
"valid_metrics_list": [BinaryIOU()],
48-
49-
"phase_callbacks": [BinarySegmentationVisualizationCallback(phase=Phase.VALIDATION_BATCH_END,
50-
freq=1,
51-
last_img_idx_in_batch=4)],
52-
}
53-
54-
trainer.train(model=model,
55-
training_params=train_params,
56-
train_loader=dl_train,
57-
valid_loader=dl_val)
40+
train_params = {
41+
"max_epochs": 50,
42+
"lr_mode": "cosine",
43+
"initial_lr": 0.0064, # for batch_size=16
44+
"optimizer_params": {"momentum": 0.843, "weight_decay": 0.00036, "nesterov": True},
45+
"cosine_final_lr_ratio": 0.1,
46+
"multiply_head_lr": 10,
47+
"optimizer": "SGD",
48+
"loss": "bce_dice_loss",
49+
"ema": True,
50+
"zero_weight_decay_on_bias_and_bn": True,
51+
"average_best_models": True,
52+
"mixed_precision": False,
53+
"metric_to_watch": "mean_IOU",
54+
"greater_metric_to_watch_is_better": True,
55+
"train_metrics_list": [BinaryIOU()],
56+
"valid_metrics_list": [BinaryIOU()],
57+
"phase_callbacks": [BinarySegmentationVisualizationCallback(phase=Phase.VALIDATION_BATCH_END, freq=1, last_img_idx_in_batch=4)],
58+
}
59+
60+
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

+3-1
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,9 @@
44
You can load any of our pretrained model in 2 lines of code:
55
```python
66
from super_gradients.training import models
7-
model = models.get("yolox_s", pretrained_weights="coco")
7+
from super_gradients.common.object_names import Models
8+
9+
model = models.get(Models.YOLOX_S, pretrained_weights="coco")
810
```
911

1012
All the available models are listed in the column `Model name`.

‎tests/end_to_end_tests/cifar_trainer_test.py

+3-2
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
import unittest
22

3+
from super_gradients.common.object_names import Models
34
from super_gradients.training import models
45

56
import super_gradients
@@ -18,7 +19,7 @@ def test_train_cifar10_dataloader(self):
1819
super_gradients.init_trainer()
1920
trainer = Trainer("test")
2021
cifar10_train_dl, cifar10_val_dl = cifar10_train(), cifar10_val()
21-
model = models.get("resnet18_cifar", arch_params={"num_classes": 10})
22+
model = models.get(Models.RESNET18_CIFAR, arch_params={"num_classes": 10})
2223
trainer.train(
2324
model=model,
2425
training_params={
@@ -37,7 +38,7 @@ def test_train_cifar100_dataloader(self):
3738
super_gradients.init_trainer()
3839
trainer = Trainer("test")
3940
cifar100_train_dl, cifar100_val_dl = cifar100_train(), cifar100_val()
40-
model = models.get("resnet18_cifar", arch_params={"num_classes": 100})
41+
model = models.get(Models.RESNET18_CIFAR, arch_params={"num_classes": 100})
4142
trainer.train(
4243
model=model,
4344
training_params={

0 commit comments

Comments
 (0)