trainer.py 6.7 KB

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