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 import BaseTrainer, BaseTrainDeamon
  14. from ...utils.config import AttrDict
  15. from ...utils import logging
  16. from .model_list import MODELS
  17. class DetTrainer(BaseTrainer):
  18. """ Object Detection Model Trainer """
  19. entities = 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. logging.info(
  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)