trainer.py 6.5 KB

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