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

kd_ema_test.py 4.6 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
  1. import unittest
  2. from super_gradients.training import models
  3. from super_gradients.training import Trainer
  4. from super_gradients.training.kd_trainer import KDTrainer
  5. import torch
  6. from super_gradients.training.utils.utils import check_models_have_same_weights
  7. from super_gradients.training.datasets.dataset_interfaces.dataset_interface import ClassificationTestDatasetInterface
  8. from super_gradients.training.metrics import Accuracy
  9. from super_gradients.training.losses.kd_losses import KDLogitsLoss
  10. class KDEMATest(unittest.TestCase):
  11. @classmethod
  12. def setUp(cls):
  13. cls.sg_trained_teacher = Trainer("sg_trained_teacher", device='cpu')
  14. cls.dataset_params = {"batch_size": 5}
  15. cls.dataset = ClassificationTestDatasetInterface(dataset_params=cls.dataset_params)
  16. cls.kd_train_params = {"max_epochs": 3, "lr_updates": [1], "lr_decay_factor": 0.1, "lr_mode": "step",
  17. "lr_warmup_epochs": 0, "initial_lr": 0.1,
  18. "loss": KDLogitsLoss(torch.nn.CrossEntropyLoss()),
  19. "optimizer": "SGD",
  20. "criterion_params": {}, "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9},
  21. "train_metrics_list": [Accuracy()], "valid_metrics_list": [Accuracy()],
  22. "metric_to_watch": "Accuracy",
  23. 'loss_logging_items_names': ["Loss", "Task Loss", "Distillation Loss"],
  24. "greater_metric_to_watch_is_better": True, "average_best_models": False,
  25. "ema": True}
  26. def test_teacher_ema_not_duplicated(self):
  27. """Check that the teacher EMA is a reference to the teacher net (not a copy)."""
  28. kd_model = KDTrainer("test_teacher_ema_not_duplicated", device='cpu')
  29. kd_model.connect_dataset_interface(self.dataset)
  30. student = models.get('resnet18', arch_params={'num_classes': 1000})
  31. teacher = models.get('resnet50', arch_params={'num_classes': 1000},
  32. pretrained_weights="imagenet")
  33. kd_model.train(training_params=self.kd_train_params, student=student, teacher=teacher)
  34. self.assertTrue(kd_model.ema_model.ema.module.teacher is kd_model.net.module.teacher)
  35. self.assertTrue(kd_model.ema_model.ema.module.student is not kd_model.net.module.student)
  36. def test_kd_ckpt_reload_net(self):
  37. """Check that the KD trainer load correctly from checkpoint when "load_ema_as_net=False"."""
  38. # Create a KD trainer and train it
  39. train_params = self.kd_train_params.copy()
  40. kd_model = KDTrainer("test_kd_ema_ckpt_reload", device='cpu')
  41. kd_model.connect_dataset_interface(self.dataset)
  42. student = models.get('resnet18', arch_params={'num_classes': 1000})
  43. teacher = models.get('resnet50', arch_params={'num_classes': 1000},
  44. pretrained_weights="imagenet")
  45. kd_model.train(training_params=self.kd_train_params, student=student, teacher=teacher)
  46. ema_model = kd_model.ema_model.ema
  47. net = kd_model.net
  48. # Load the trained KD trainer
  49. kd_model = KDTrainer("test_kd_ema_ckpt_reload", device='cpu')
  50. kd_model.connect_dataset_interface(self.dataset)
  51. student = models.get('resnet18', arch_params={'num_classes': 1000})
  52. teacher = models.get('resnet50', arch_params={'num_classes': 1000},
  53. pretrained_weights="imagenet")
  54. train_params["resume"] = True
  55. kd_model.train(training_params=train_params, student=student, teacher=teacher)
  56. reloaded_ema_model = kd_model.ema_model.ema
  57. reloaded_net = kd_model.net
  58. # trained ema == loaded ema (Should always be true as long as "ema=True" in train_params)
  59. self.assertTrue(check_models_have_same_weights(ema_model, reloaded_ema_model))
  60. # loaded net == trained net (since load_ema_as_net = False)
  61. self.assertTrue(check_models_have_same_weights(reloaded_net, net))
  62. # loaded net != trained ema (since load_ema_as_net = False)
  63. self.assertTrue(not check_models_have_same_weights(reloaded_net, ema_model))
  64. # loaded student ema == loaded student net (since load_ema_as_net = False)
  65. self.assertTrue(not check_models_have_same_weights(reloaded_ema_model.module.student, reloaded_net.module.student))
  66. # loaded teacher ema == loaded teacher net (teacher always loads ema)
  67. self.assertTrue(check_models_have_same_weights(reloaded_ema_model.module.teacher, reloaded_net.module.teacher))
  68. if __name__ == '__main__':
  69. unittest.main()
Tip!

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

Comments

Loading...