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

#578 Feature/sg 516 support head replacement for local pretrained weights unknown dataset

Merged
Ghost merged 1 commits into Deci-AI:master from deci-ai:feature/SG-516_support_head_replacement_for_local_pretrained_weights_unknown_dataset
1 changed files with 40 additions and 0 deletions
  1. 40
    0
      tests/unit_tests/local_ckpt_head_replacement_test.py
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
  1. import unittest
  2. from super_gradients.training import Trainer, models
  3. from super_gradients.training.dataloaders.dataloaders import classification_test_dataloader
  4. from super_gradients.training.metrics import Accuracy, Top5
  5. from super_gradients.training.utils.utils import check_models_have_same_weights
  6. import os
  7. class LocalCkptHeadReplacementTest(unittest.TestCase):
  8. def test_local_ckpt_head_replacement(self):
  9. train_params = {
  10. "max_epochs": 1,
  11. "lr_updates": [1],
  12. "lr_decay_factor": 0.1,
  13. "lr_mode": "step",
  14. "lr_warmup_epochs": 0,
  15. "initial_lr": 0.1,
  16. "loss": "cross_entropy",
  17. "optimizer": "SGD",
  18. "criterion_params": {},
  19. "optimizer_params": {"weight_decay": 1e-4, "momentum": 0.9},
  20. "train_metrics_list": [Accuracy(), Top5()],
  21. "valid_metrics_list": [Accuracy(), Top5()],
  22. "metric_to_watch": "Accuracy",
  23. "greater_metric_to_watch_is_better": True,
  24. }
  25. # Define Model
  26. net = models.get("resnet18", num_classes=5)
  27. trainer = Trainer("test_resume_training")
  28. trainer.train(model=net, training_params=train_params, train_loader=classification_test_dataloader(), valid_loader=classification_test_dataloader())
  29. ckpt_path = os.path.join(trainer.checkpoints_dir_path, "ckpt_latest.pth")
  30. net2 = models.get("resnet18", num_classes=10, checkpoint_num_classes=5, checkpoint_path=ckpt_path)
  31. self.assertFalse(check_models_have_same_weights(net, net2))
  32. net.linear = None
  33. net2.linear = None
  34. self.assertTrue(check_models_have_same_weights(net, net2))
Discard