trainer.py 4.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129
  1. # !/usr/bin/env python3
  2. # -*- coding: UTF-8 -*-
  3. ################################################################################
  4. #
  5. # Copyright (c) 2024 Baidu.com, Inc. All Rights Reserved
  6. #
  7. ################################################################################
  8. """
  9. Author: PaddlePaddle Authors
  10. """
  11. from pathlib import Path
  12. import paddle
  13. from ..base.trainer import BaseTrainer
  14. from ..base.train_deamon import BaseTrainDeamon
  15. from ...utils.config import AttrDict
  16. from .support_models import SUPPORT_MODELS
  17. class DetTrainer(BaseTrainer):
  18. """ Object Detection Model Trainer """
  19. support_models = SUPPORT_MODELS
  20. def build_deamon(self, config: AttrDict) -> "DetTrainDeamon":
  21. """build deamon thread for saving training outputs timely
  22. Args:
  23. config (AttrDict): PaddleX pipeline config, which is loaded from pipeline yaml file.
  24. Returns:
  25. DetTrainDeamon: the training deamon thread object for saving training outputs timely.
  26. """
  27. return DetTrainDeamon(config)
  28. def _update_dataset(self):
  29. """update dataset settings
  30. """
  31. self.pdx_config.update_dataset(self.global_config.dataset_dir,
  32. "COCODetDataset")
  33. def update_config(self):
  34. """update training config
  35. """
  36. if self.train_config.log_interval:
  37. self.pdx_config.update_log_interval(self.train_config.log_interval)
  38. if self.train_config.eval_interval:
  39. self.pdx_config.update_eval_interval(
  40. self.train_config.eval_interval)
  41. self._update_dataset()
  42. if self.train_config.num_classes is not None:
  43. self.pdx_config.update_num_class(self.train_config.num_classes)
  44. if self.train_config.pretrain_weight_path and self.train_config.pretrain_weight_path != "":
  45. self.pdx_config.update_pretrained_weights(
  46. self.train_config.pretrain_weight_path)
  47. if self.train_config.batch_size is not None:
  48. self.pdx_config.update_batch_size(self.train_config.batch_size)
  49. if self.train_config.learning_rate is not None:
  50. self.pdx_config.update_learning_rate(
  51. self.train_config.learning_rate)
  52. if self.train_config.epochs_iters is not None:
  53. self.pdx_config.update_epochs(self.train_config.epochs_iters)
  54. epochs_iters = self.train_config.epochs_iters
  55. else:
  56. epochs_iters = self.pdx_config.get_epochs_iters()
  57. if self.global_config.output is not None:
  58. self.pdx_config.update_save_dir(self.global_config.output)
  59. if "PicoDet" in self.global_config.model:
  60. assigner_epochs = max(int(epochs_iters / 10), 1)
  61. try:
  62. self.pdx_config.update_static_assigner_epochs(assigner_epochs)
  63. except Exception:
  64. print(
  65. f"The model({self.global_config.model}) don't support to update_static_assigner_epochs!"
  66. )
  67. def get_train_kwargs(self) -> dict:
  68. """get key-value arguments of model training function
  69. Returns:
  70. dict: the arguments of training function.
  71. """
  72. train_args = {"device": self.get_device()}
  73. if self.train_config.resume_path is not None and self.train_config.resume_path != "":
  74. train_args["resume_path"] = self.train_config.resume_path
  75. return train_args
  76. class DetTrainDeamon(BaseTrainDeamon):
  77. """ DetTrainResultDemon """
  78. def __init__(self, *args, **kwargs):
  79. super().__init__(*args, **kwargs)
  80. def get_the_pdparams_suffix(self):
  81. """ get the suffix of pdparams file """
  82. return "pdparams"
  83. def get_the_pdema_suffix(self):
  84. """ get the suffix of pdema file """
  85. return "pdema"
  86. def get_the_pdopt_suffix(self):
  87. """ get the suffix of pdopt file """
  88. return "pdopt"
  89. def get_the_pdstates_suffix(self):
  90. """ get the suffix of pdstates file """
  91. return "pdstates"
  92. def get_ith_ckp_prefix(self, epoch_id):
  93. """ get the prefix of the epoch_id checkpoint file """
  94. return f"{epoch_id}"
  95. def get_best_ckp_prefix(self):
  96. """ get the prefix of the best checkpoint file """
  97. return "best_model"
  98. def get_score(self, pdstates_path):
  99. """ get the score by pdstates file """
  100. if not Path(pdstates_path).exists():
  101. return 0
  102. return paddle.load(pdstates_path)["metric"]
  103. def get_epoch_id_by_pdparams_prefix(self, pdparams_prefix):
  104. """ get the epoch_id by pdparams file """
  105. return int(pdparams_prefix)