trainer.py 7.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199
  1. # copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import os
  15. import json
  16. import time
  17. import tarfile
  18. from pathlib import Path
  19. import paddle
  20. from ..base import BaseTrainer, BaseTrainDeamon
  21. from ...utils.config import AttrDict
  22. from .model_list import MODELS
  23. class TSFCTrainer(BaseTrainer):
  24. """TS Forecast Model Trainer"""
  25. entities = MODELS
  26. def build_deamon(self, config: AttrDict) -> "TSFCTrainDeamon":
  27. """build deamon thread for saving training outputs timely
  28. Args:
  29. config (AttrDict): PaddleX pipeline config, which is loaded from pipeline yaml file.
  30. Returns:
  31. TSFCTrainDeamon: the training deamon thread object for saving training outputs timely.
  32. """
  33. return TSFCTrainDeamon(config)
  34. def train(self):
  35. """firstly, update and dump train config, then train model"""
  36. # XXX: using super().train() instead when the train_hook() is supported.
  37. os.makedirs(self.global_config.output, exist_ok=True)
  38. self.update_config()
  39. self.dump_config()
  40. train_result = self.pdx_model.train(**self.get_train_kwargs())
  41. assert (
  42. train_result.returncode == 0
  43. ), f"Encountered an unexpected error({train_result.returncode}) in \
  44. training!"
  45. self.make_tar_file()
  46. self.deamon.stop()
  47. def make_tar_file(self):
  48. """make tar file to package the training outputs"""
  49. tar_path = Path(self.global_config.output) / "best_accuracy.pdparams.tar"
  50. with tarfile.open(tar_path, "w") as tar:
  51. tar.add(self.global_config.output, arcname="best_accuracy.pdparams")
  52. def update_config(self):
  53. """update training config"""
  54. self.pdx_config.update_dataset(self.global_config.dataset_dir, "TSDataset")
  55. if self.train_config.input_len is not None:
  56. self.pdx_config.update_input_len(self.train_config.input_len)
  57. if self.train_config.time_col is not None:
  58. self.pdx_config.update_basic_info({"time_col": self.train_config.time_col})
  59. if self.train_config.target_cols is not None:
  60. self.pdx_config.update_basic_info(
  61. {"target_cols": self.train_config.target_cols.split(",")}
  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.predict_len is not None:
  70. self.pdx_config.update_predict_len(self.train_config.predict_len)
  71. if self.train_config.patience is not None:
  72. self.pdx_config.update_patience(self.train_config.patience)
  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(self.train_config.learning_rate)
  77. if self.train_config.epochs_iters is not None:
  78. self.pdx_config.update_epochs(self.train_config.epochs_iters)
  79. if self.global_config.output is not None:
  80. self.pdx_config.update_save_dir(self.global_config.output)
  81. def get_train_kwargs(self) -> dict:
  82. """get key-value arguments of model training function
  83. Returns:
  84. dict: the arguments of training function.
  85. """
  86. train_args = {"device": self.get_device()}
  87. if self.global_config.output is not None:
  88. train_args["save_dir"] = self.global_config.output
  89. return train_args
  90. class TSFCTrainDeamon(BaseTrainDeamon):
  91. """TSFCTrainResultDemon"""
  92. def get_watched_model(self):
  93. """get the models needed to be watched"""
  94. watched_models = []
  95. watched_models.append("best")
  96. return watched_models
  97. def update(self):
  98. """update train result json"""
  99. self.processing = True
  100. for i, result in enumerate(self.results):
  101. self.results[i] = self.update_result(result, self.train_outputs[i])
  102. self.save_json()
  103. self.processing = False
  104. def update_train_log(self, train_output):
  105. """update train log"""
  106. train_log_path = train_output / "train_ct.log"
  107. with open(train_log_path, "w") as f:
  108. seconds = time.time()
  109. f.write(
  110. "current training time: "
  111. + time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(seconds))
  112. )
  113. f.close()
  114. return train_log_path
  115. def update_result(self, result, train_output):
  116. """update every result"""
  117. config = Path(train_output).joinpath("config.yaml")
  118. if not config.exists():
  119. return result
  120. result["config"] = config
  121. result["train_log"] = self.update_train_log(train_output)
  122. result["visualdl_log"] = self.update_vdl_log(train_output)
  123. result["label_dict"] = self.update_label_dict(train_output)
  124. self.update_models(result, train_output, "best")
  125. return result
  126. def update_models(self, result, train_output, model_key):
  127. """update info of the models to be saved"""
  128. pdparams = Path(train_output).joinpath("best_accuracy.pdparams.tar")
  129. if pdparams.exists():
  130. score = self.get_score(Path(train_output).joinpath("score.json"))
  131. result["models"][model_key] = {
  132. "score": "%.3f" % score,
  133. "pdparams": pdparams,
  134. "pdema": "",
  135. "pdopt": "",
  136. "pdstates": "",
  137. "inference_config": "",
  138. "pdmodel": "",
  139. "pdiparams": pdparams,
  140. "pdiparams.info": "",
  141. }
  142. def get_score(self, score_path):
  143. """get the score by pdstates file"""
  144. if not Path(score_path).exists():
  145. return 0
  146. return json.load(open(score_path))["metric"]
  147. def get_best_ckp_prefix(self):
  148. """get the prefix of the best checkpoint file"""
  149. pass
  150. def get_epoch_id_by_pdparams_prefix(self):
  151. """get the epoch_id by pdparams file"""
  152. pass
  153. def get_ith_ckp_prefix(self):
  154. """get the prefix of the epoch_id checkpoint file"""
  155. pass
  156. def get_the_pdema_suffix(self):
  157. """get the suffix of pdema file"""
  158. pass
  159. def get_the_pdopt_suffix(self):
  160. """get the suffix of pdopt file"""
  161. pass
  162. def get_the_pdparams_suffix(self):
  163. """get the suffix of pdparams file"""
  164. pass
  165. def get_the_pdstates_suffix(self):
  166. """get the suffix of pdstates file"""
  167. pass