1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
|
- import unittest
- import super_gradients
- from super_gradients.training import MultiGPUMode
- from super_gradients.training import SgModel
- from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ImageNetDatasetInterface, \
- ClassificationTestDatasetInterface
- from super_gradients.training.metrics import Accuracy
- import os
- import shutil
- class PretrainedModelsTest(unittest.TestCase):
- def setUp(self) -> None:
- super_gradients.init_trainer()
- self.imagenet_pretrained_models = ["resnet50", "repvgg_a0", "regnetY800"]
- self.imagenet_pretrained_arch_params = {"resnet50": {"pretrained_weights": "imagenet"},
- "regnetY800": {"pretrained_weights": "imagenet"},
- "repvgg_a0": {"pretrained_weights": "imagenet",
- "build_residual_branches": True}}
- self.imagenet_pretrained_accuracies = {"resnet50": 0.763,
- "repvgg_a0": 0.7205,
- "regnetY800": 0.7605}
- self.imagenet_dataset = ImageNetDatasetInterface(data_dir="/data/Imagenet", dataset_params={"batch_size": 128})
- self.transfer_classification_dataset = ClassificationTestDatasetInterface(image_size=224)
- self.transfer_classification_train_params = {"max_epochs": 3,
- "lr_updates": [1],
- "lr_decay_factor": 0.1,
- "initial_lr": 0.6,
- "loss": "cross_entropy",
- "lr_mode": "step",
- "optimizer_params": {"weight_decay": 0.000,
- "momentum": 0.9},
- "train_metrics_list": [Accuracy()],
- "valid_metrics_list": [Accuracy()],
- "loss_logging_items_names": ["Loss"],
- "metric_to_watch": "Accuracy",
- "greater_metric_to_watch_is_better": True}
- def test_pretrained_resnet50_imagenet(self):
- trainer = SgModel('imagenet_pretrained_resnet50', model_checkpoints_location='local',
- multi_gpu=MultiGPUMode.OFF)
- trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
- trainer.build_model("resnet50", arch_params=self.imagenet_pretrained_arch_params["resnet50"])
- res = trainer.test(test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
- metrics_progress_verbose=True)[0].cpu().item()
- self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["resnet50"])
- def test_transfer_learning_resnet50_imagenet(self):
- trainer = SgModel('imagenet_pretrained_resnet50_transfer_learning', model_checkpoints_location='local',
- multi_gpu=MultiGPUMode.OFF)
- trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
- trainer.build_model("resnet50", arch_params=self.imagenet_pretrained_arch_params["resnet50"])
- trainer.train(training_params=self.transfer_classification_train_params)
- def test_pretrained_regnetY800_imagenet(self):
- trainer = SgModel('imagenet_pretrained_regnetY800', model_checkpoints_location='local',
- multi_gpu=MultiGPUMode.OFF)
- trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
- trainer.build_model("regnetY800", arch_params=self.imagenet_pretrained_arch_params["regnetY800"])
- res = trainer.test(test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
- metrics_progress_verbose=True)[0].cpu().item()
- self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["regnetY800"])
- def test_transfer_learning_regnetY800_imagenet(self):
- trainer = SgModel('imagenet_pretrained_regnetY800_transfer_learning', model_checkpoints_location='local',
- multi_gpu=MultiGPUMode.OFF)
- trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
- trainer.build_model("regnetY800", arch_params=self.imagenet_pretrained_arch_params["regnetY800"])
- trainer.train(training_params=self.transfer_classification_train_params)
- def test_pretrained_repvgg_a0_imagenet(self):
- trainer = SgModel('imagenet_pretrained_repvgg_a0', model_checkpoints_location='local',
- multi_gpu=MultiGPUMode.OFF)
- trainer.connect_dataset_interface(self.imagenet_dataset, data_loader_num_workers=8)
- trainer.build_model("repvgg_a0", arch_params=self.imagenet_pretrained_arch_params["repvgg_a0"])
- res = trainer.test(test_loader=self.imagenet_dataset.val_loader, test_metrics_list=[Accuracy()],
- metrics_progress_verbose=True)[0].cpu().item()
- self.assertAlmostEqual(res, self.imagenet_pretrained_accuracies["repvgg_a0"])
- def test_transfer_learning_repvgg_a0_imagenet(self):
- trainer = SgModel('imagenet_pretrained_repvgg_a0_transfer_learning', model_checkpoints_location='local',
- multi_gpu=MultiGPUMode.OFF)
- trainer.connect_dataset_interface(self.transfer_classification_dataset, data_loader_num_workers=8)
- trainer.build_model("repvgg_a0", arch_params=self.imagenet_pretrained_arch_params["repvgg_a0"])
- trainer.train(training_params=self.transfer_classification_train_params)
- def tearDown(self) -> None:
- if os.path.exists('~/.cache/torch/hub/'):
- shutil.rmtree('~/.cache/torch/hub/')
- if __name__ == '__main__':
- unittest.main()
|