trainer.py 6.5 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193
  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 os
  12. import json
  13. import time
  14. import tarfile
  15. from pathlib import Path
  16. import paddle
  17. from ..base import BaseTrainer, BaseTrainDeamon
  18. from ...utils.config import AttrDict
  19. from .model_list import MODELS
  20. class TSCLSTrainer(BaseTrainer):
  21. """ TS Classification Model Trainer """
  22. entities = MODELS
  23. def build_deamon(self, config: AttrDict) -> "TSCLSTrainDeamon":
  24. """build deamon thread for saving training outputs timely
  25. Args:
  26. config (AttrDict): PaddleX pipeline config, which is loaded from pipeline yaml file.
  27. Returns:
  28. TSCLSTrainDeamon: the training deamon thread object for saving training outputs timely.
  29. """
  30. return TSCLSTrainDeamon(config)
  31. def train(self):
  32. """firstly, update and dump train config, then train model
  33. """
  34. rtn = super().train()
  35. self.make_tar_file()
  36. return rtn
  37. def make_tar_file(self):
  38. """make tar file to package the training outputs
  39. """
  40. tar_path = Path(
  41. self.global_config.output) / "best_accuracy.pdparams.tar"
  42. with tarfile.open(tar_path, 'w') as tar:
  43. tar.add(self.global_config.output, arcname='best_accuracy.pdparams')
  44. def update_config(self):
  45. """update training config
  46. """
  47. self.pdx_config.update_dataset(self.global_config.dataset_dir,
  48. "TSCLSDataset")
  49. if self.train_config.time_col is not None:
  50. self.pdx_config.update_basic_info({
  51. 'time_col': self.train_config.time_col
  52. })
  53. if self.train_config.target_cols is not None:
  54. self.pdx_config.update_basic_info({
  55. 'target_cols': self.train_config.target_cols.split(',')
  56. })
  57. if self.train_config.group_id is not None:
  58. self.pdx_config.update_basic_info({
  59. 'group_id': self.train_config.group_id
  60. })
  61. if self.train_config.static_cov_cols is not None:
  62. self.pdx_config.update_basic_info({
  63. 'static_cov_cols': self.train_config.static_cov_cols
  64. })
  65. if self.train_config.freq is not None:
  66. try:
  67. self.train_config.freq = int(self.train_config.freq)
  68. except ValueError:
  69. pass
  70. self.pdx_config.update_basic_info({'freq': self.train_config.freq})
  71. if self.train_config.batch_size is not None:
  72. self.pdx_config.update_batch_size(self.train_config.batch_size)
  73. if self.train_config.learning_rate is not None:
  74. self.pdx_config.update_learning_rate(
  75. self.train_config.learning_rate)
  76. if self.train_config.epochs_iters is not None:
  77. self.pdx_config.update_epochs(self.train_config.epochs_iters)
  78. if self.global_config.output is not None:
  79. self.pdx_config.update_save_dir(self.global_config.output)
  80. def get_train_kwargs(self) -> dict:
  81. """get key-value arguments of model training function
  82. Returns:
  83. dict: the arguments of training function.
  84. """
  85. train_args = {"device": self.get_device()}
  86. if self.global_config.output is not None:
  87. train_args["save_dir"] = self.global_config.output
  88. return train_args
  89. class TSCLSTrainDeamon(BaseTrainDeamon):
  90. """ TSCLSTrainResultDemon """
  91. def get_watched_model(self):
  92. """ get the models needed to be watched """
  93. watched_models = []
  94. watched_models.append("best")
  95. return watched_models
  96. def update(self):
  97. """ update train result json """
  98. self.processing = True
  99. for i, result in enumerate(self.results):
  100. self.results[i] = self.update_result(result, self.train_outputs[i])
  101. self.save_json()
  102. self.processing = False
  103. def update_train_log(self, train_output):
  104. """ update train log """
  105. train_log_path = train_output / "train_ct.log"
  106. with open(train_log_path, 'w') as f:
  107. seconds = time.time()
  108. f.write('current training time: ' + time.strftime(
  109. "%Y-%m-%d %H:%M:%S", time.localtime(seconds)))
  110. f.close()
  111. return train_log_path
  112. def update_result(self, result, train_output):
  113. """ update every result """
  114. config = Path(train_output).joinpath("config.yaml")
  115. if not config.exists():
  116. return result
  117. result["config"] = config
  118. result["train_log"] = self.update_train_log(train_output)
  119. result["visualdl_log"] = self.update_vdl_log(train_output)
  120. result["label_dict"] = self.update_label_dict(train_output)
  121. self.update_models(result, train_output, "best")
  122. return result
  123. def update_models(self, result, train_output, model_key):
  124. """ update info of the models to be saved """
  125. pdparams = Path(train_output).joinpath("best_accuracy.pdparams.tar")
  126. if pdparams.exists():
  127. score = self.get_score(Path(train_output).joinpath("score.json"))
  128. result["models"][model_key] = {
  129. "score": "%.3f" % score,
  130. "pdparams": pdparams,
  131. "pdema": "",
  132. "pdopt": "",
  133. "pdstates": "",
  134. "inference_config": "",
  135. "pdmodel": "",
  136. "pdiparams": pdparams,
  137. "pdiparams.info": ""
  138. }
  139. def get_score(self, score_path):
  140. """ get the score by pdstates file """
  141. if not Path(score_path).exists():
  142. return 0
  143. return json.load(open(score_path))["metric"]
  144. def get_best_ckp_prefix(self):
  145. """ get the prefix of the best checkpoint file """
  146. pass
  147. def get_epoch_id_by_pdparams_prefix(self):
  148. """ get the epoch_id by pdparams file """
  149. pass
  150. def get_ith_ckp_prefix(self):
  151. """ get the prefix of the epoch_id checkpoint file """
  152. pass
  153. def get_the_pdema_suffix(self):
  154. """ get the suffix of pdema file """
  155. pass
  156. def get_the_pdopt_suffix(self):
  157. """ get the suffix of pdopt file """
  158. pass
  159. def get_the_pdparams_suffix(self):
  160. """ get the suffix of pdparams file """
  161. pass
  162. def get_the_pdstates_suffix(self):
  163. """ get the suffix of pdstates file """
  164. pass