trainer.py 5.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140
  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. import json
  12. import shutil
  13. import paddle
  14. from pathlib import Path
  15. from ..base.trainer import BaseTrainer
  16. from ..base.train_deamon import BaseTrainDeamon
  17. from .support_models import SUPPORT_MODELS
  18. from ...utils.config import AttrDict
  19. class ClsTrainer(BaseTrainer):
  20. """ Image Classification Model Trainer """
  21. support_models = SUPPORT_MODELS
  22. def dump_label_dict(self, src_label_dict_path: str):
  23. """dump label dict config
  24. Args:
  25. src_label_dict_path (str): path to label dict file to be saved.
  26. """
  27. dst_label_dict_path = Path(self.global_config.output).joinpath(
  28. "label_dict.txt")
  29. shutil.copyfile(src_label_dict_path, dst_label_dict_path)
  30. def build_deamon(self, config: AttrDict) -> "ClsTrainDeamon":
  31. """build deamon thread for saving training outputs timely
  32. Args:
  33. config (AttrDict): PaddleX pipeline config, which is loaded from pipeline yaml file.
  34. Returns:
  35. ClsTrainDeamon: the training deamon thread object for saving training outputs timely.
  36. """
  37. return ClsTrainDeamon(config)
  38. def update_config(self):
  39. """update training config
  40. """
  41. if self.train_config.log_interval:
  42. self.pdx_config.update_log_interval(self.train_config.log_interval)
  43. if self.train_config.eval_interval:
  44. self.pdx_config.update_eval_interval(
  45. self.train_config.eval_interval)
  46. if self.train_config.save_interval:
  47. self.pdx_config.update_save_interval(
  48. self.train_config.save_interval)
  49. self.pdx_config.update_dataset(self.global_config.dataset_dir,
  50. "ClsDataset")
  51. if self.train_config.num_classes is not None:
  52. self.pdx_config.update_num_classes(self.train_config.num_classes)
  53. if self.train_config.pretrain_weight_path and self.train_config.pretrain_weight_path != "":
  54. self.pdx_config.update_pretrained_weights(
  55. self.train_config.pretrain_weight_path)
  56. label_dict_path = Path(self.global_config.dataset_dir).joinpath(
  57. "label.txt")
  58. if label_dict_path.exists():
  59. self.dump_label_dict(label_dict_path)
  60. if self.train_config.batch_size is not None:
  61. self.pdx_config.update_batch_size(self.train_config.batch_size)
  62. if self.train_config.learning_rate is not None:
  63. self.pdx_config.update_learning_rate(
  64. self.train_config.learning_rate)
  65. if self.train_config.epochs_iters is not None:
  66. self.pdx_config._update_epochs(self.train_config.epochs_iters)
  67. if self.train_config.warmup_steps is not None:
  68. self.pdx_config.update_warmup_epochs(self.train_config.warmup_steps)
  69. if self.global_config.output is not None:
  70. self.pdx_config._update_output_dir(self.global_config.output)
  71. def get_train_kwargs(self) -> dict:
  72. """get key-value arguments of model training function
  73. Returns:
  74. dict: the arguments of training function.
  75. """
  76. train_args = {"device": self.get_device()}
  77. if self.train_config.resume_path is not None and self.train_config.resume_path != "":
  78. train_args["resume_path"] = self.train_config.resume_path
  79. return train_args
  80. class ClsTrainDeamon(BaseTrainDeamon):
  81. """ ClsTrainResultDemon """
  82. def __init__(self, *args, **kwargs):
  83. super().__init__(*args, **kwargs)
  84. def get_the_pdparams_suffix(self):
  85. """ get the suffix of pdparams file """
  86. return "pdparams"
  87. def get_the_pdema_suffix(self):
  88. """ get the suffix of pdema file """
  89. return "pdema"
  90. def get_the_pdopt_suffix(self):
  91. """ get the suffix of pdopt file """
  92. return "pdopt"
  93. def get_the_pdstates_suffix(self):
  94. """ get the suffix of pdstates file """
  95. return "pdstates"
  96. def get_ith_ckp_prefix(self, epoch_id):
  97. """ get the prefix of the epoch_id checkpoint file """
  98. return f"epoch_{epoch_id}"
  99. def get_best_ckp_prefix(self):
  100. """ get the prefix of the best checkpoint file """
  101. return "best_model"
  102. def get_score(self, pdstates_path):
  103. """ get the score by pdstates file """
  104. if not Path(pdstates_path).exists():
  105. return 0
  106. return paddle.load(pdstates_path)["metric"]
  107. def get_epoch_id_by_pdparams_prefix(self, pdparams_prefix):
  108. """ get the epoch_id by pdparams file """
  109. return int(pdparams_prefix.split("_")[-1])
  110. def update_label_dict(self, train_output):
  111. """ update label dict """
  112. dict_path = train_output.joinpath("label_dict.txt")
  113. if not dict_path.exists():
  114. return ""
  115. return dict_path