Register
Login
Resources
Docs Blog Datasets Glossary Case Studies Tutorials & Webinars
Product
Data Engine LLMs Platform Enterprise
Pricing Explore
Connect to our Discord channel

pretrained_models_test.py 6.0 KB

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

Press p or to see the previous file or, n or to see the next file

Comments

Loading...