trainer.py 4.9 KB

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