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

extreme_batch_cb_test.py 7.2 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
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
  1. import unittest
  2. from super_gradients import Trainer
  3. from super_gradients.common.object_names import Models
  4. from super_gradients.training import models
  5. from super_gradients.training.dataloaders.dataloaders import segmentation_test_dataloader, detection_test_dataloader
  6. from super_gradients.training.losses import PPYoloELoss
  7. from super_gradients.training.losses.ddrnet_loss import DDRNetLoss
  8. from super_gradients.training.metrics import IoU, DetectionMetrics_050
  9. from super_gradients.training.models.detection_models.pp_yolo_e import PPYoloEPostPredictionCallback
  10. from super_gradients.training.utils.callbacks.callbacks import ExtremeBatchSegVisualizationCallback, ExtremeBatchDetectionVisualizationCallback
  11. # Helper method to set up Trainer and model with common parameters
  12. def setup_trainer_and_model_seg(experiment_name: str):
  13. trainer = Trainer(experiment_name)
  14. model = models.get(Models.DDRNET_23, arch_params={"use_aux_heads": True}, pretrained_weights="cityscapes")
  15. return trainer, model
  16. def setup_trainer_and_model_detection(experiment_name: str):
  17. trainer = Trainer(experiment_name)
  18. model = models.get(Models.YOLO_NAS_S, num_classes=1)
  19. return trainer, model
  20. class DummyIOU(IoU):
  21. """
  22. Metric for testing the segmentation callback works with compound metrics
  23. """
  24. def compute(self):
  25. diou = super(DummyIOU, self).compute()
  26. return {"diou": diou, "diou_minus": -1 * diou}
  27. class ExtremeBatchSanityTest(unittest.TestCase):
  28. @classmethod
  29. def setUpClass(cls):
  30. cls.seg_training_params = {
  31. "max_epochs": 3,
  32. "initial_lr": 1e-2,
  33. "loss": DDRNetLoss(),
  34. "lr_mode": "PolyLRScheduler",
  35. "ema": True,
  36. "optimizer": "SGD",
  37. "mixed_precision": False,
  38. "optimizer_params": {"weight_decay": 5e-4, "momentum": 0.9},
  39. "load_opt_params": False,
  40. "train_metrics_list": [IoU(5)],
  41. "valid_metrics_list": [IoU(5)],
  42. "metric_to_watch": "IoU",
  43. "greater_metric_to_watch_is_better": True,
  44. }
  45. cls.od_training_params = {
  46. "max_epochs": 3,
  47. "initial_lr": 1e-2,
  48. "loss": PPYoloELoss(num_classes=1, use_static_assigner=False, reg_max=16),
  49. "lr_mode": "PolyLRScheduler",
  50. "ema": True,
  51. "optimizer": "SGD",
  52. "mixed_precision": False,
  53. "optimizer_params": {"weight_decay": 5e-4, "momentum": 0.9},
  54. "load_opt_params": False,
  55. "valid_metrics_list": [
  56. DetectionMetrics_050(
  57. normalize_targets=True,
  58. post_prediction_callback=PPYoloEPostPredictionCallback(score_threshold=0.03, nms_top_k=1000, max_predictions=300, nms_threshold=0.65),
  59. num_cls=1,
  60. )
  61. ],
  62. "train_metrics_list": [],
  63. "metric_to_watch": "mAP@0.50",
  64. "greater_metric_to_watch_is_better": True,
  65. }
  66. def test_detection_extreme_batch_with_metric_sanity(self):
  67. trainer, model = setup_trainer_and_model_detection("test_detection_extreme_batch_with_metric_sanity")
  68. self.od_training_params["phase_callbacks"] = [
  69. ExtremeBatchDetectionVisualizationCallback(
  70. classes=["1"],
  71. metric=DetectionMetrics_050(
  72. normalize_targets=True,
  73. post_prediction_callback=PPYoloEPostPredictionCallback(score_threshold=0.03, nms_top_k=1000, max_predictions=300, nms_threshold=0.65),
  74. num_cls=1,
  75. ),
  76. metric_component_name="mAP@0.50",
  77. post_prediction_callback=PPYoloEPostPredictionCallback(score_threshold=0.03, nms_top_k=1000, max_predictions=300, nms_threshold=0.65),
  78. )
  79. ]
  80. trainer.train(model=model, training_params=self.od_training_params, train_loader=detection_test_dataloader(), valid_loader=detection_test_dataloader())
  81. def test_detection_extreme_batch_with_loss_sanity(self):
  82. trainer, model = setup_trainer_and_model_detection("test_detection_extreme_batch_with_loss_sanity")
  83. self.od_training_params["phase_callbacks"] = [
  84. ExtremeBatchDetectionVisualizationCallback(
  85. classes=["1"],
  86. loss_to_monitor="PPYoloELoss/loss_cls",
  87. post_prediction_callback=PPYoloEPostPredictionCallback(score_threshold=0.03, nms_top_k=1000, max_predictions=300, nms_threshold=0.65),
  88. )
  89. ]
  90. trainer.train(model=model, training_params=self.od_training_params, train_loader=detection_test_dataloader(), valid_loader=detection_test_dataloader())
  91. def test_segmentation_extreme_batch_with_metric_sanity(self):
  92. trainer, model = setup_trainer_and_model_seg("test_segmentation_extreme_batch_with_metric_sanity")
  93. self.seg_training_params["phase_callbacks"] = [ExtremeBatchSegVisualizationCallback(IoU(5))]
  94. trainer.train(
  95. model=model, training_params=self.seg_training_params, train_loader=segmentation_test_dataloader(), valid_loader=segmentation_test_dataloader()
  96. )
  97. def test_segmentation_extreme_batch_with_compound_metric_sanity(self):
  98. trainer, model = setup_trainer_and_model_seg("test_segmentation_extreme_batch_with_compound_metric_sanity")
  99. self.seg_training_params["phase_callbacks"] = [ExtremeBatchSegVisualizationCallback(DummyIOU(5), metric_component_name="diou_minus")]
  100. trainer.train(
  101. model=model, training_params=self.seg_training_params, train_loader=segmentation_test_dataloader(), valid_loader=segmentation_test_dataloader()
  102. )
  103. def test_segmentation_extreme_batch_with_loss_sanity(self):
  104. trainer, model = setup_trainer_and_model_seg("test_segmentation_extreme_batch_with_loss_sanity")
  105. self.seg_training_params["phase_callbacks"] = [ExtremeBatchSegVisualizationCallback(loss_to_monitor="DDRNetLoss/aux_loss1")]
  106. trainer.train(
  107. model=model, training_params=self.seg_training_params, train_loader=segmentation_test_dataloader(), valid_loader=segmentation_test_dataloader()
  108. )
  109. def test_segmentation_extreme_batch_train_only(self):
  110. trainer, model = setup_trainer_and_model_seg("test_segmentation_extreme_batch_train_only")
  111. self.seg_training_params["phase_callbacks"] = [
  112. ExtremeBatchSegVisualizationCallback(loss_to_monitor="DDRNetLoss/aux_loss1", enable_on_train_loader=True, enable_on_valid_loader=False)
  113. ]
  114. trainer.train(
  115. model=model, training_params=self.seg_training_params, train_loader=segmentation_test_dataloader(), valid_loader=segmentation_test_dataloader()
  116. )
  117. def test_segmentation_extreme_batch_train_and_valid(self):
  118. trainer, model = setup_trainer_and_model_seg("test_segmentation_extreme_batch_train_and_valid")
  119. self.seg_training_params["phase_callbacks"] = [
  120. ExtremeBatchSegVisualizationCallback(loss_to_monitor="DDRNetLoss/aux_loss1", enable_on_train_loader=True, enable_on_valid_loader=True)
  121. ]
  122. trainer.train(
  123. model=model, training_params=self.seg_training_params, train_loader=segmentation_test_dataloader(), valid_loader=segmentation_test_dataloader()
  124. )
  125. if __name__ == "__main__":
  126. unittest.main()
Tip!

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

Comments

Loading...