trainer.py 3.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112
  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. from pathlib import Path
  13. import paddle
  14. from ..base import BaseTrainer, BaseTrainDeamon
  15. from ...utils.config import AttrDict
  16. from .model_list import MODELS
  17. class TableRecTrainer(BaseTrainer):
  18. """ Table Recognition Model Trainer """
  19. entities = MODELS
  20. def build_deamon(self, config: AttrDict) -> "TableRecTrainDeamon":
  21. """build deamon thread for saving training outputs timely
  22. Args:
  23. config (AttrDict): PaddleX pipeline config, which is loaded from pipeline yaml file.
  24. Returns:
  25. TableRecTrainDeamon: the training deamon thread object for saving training outputs timely.
  26. """
  27. return TableRecTrainDeamon(config)
  28. def update_config(self):
  29. """update training config
  30. """
  31. if self.train_config.log_interval:
  32. self.pdx_config.update_log_interval(self.train_config.log_interval)
  33. if self.train_config.eval_interval:
  34. self.pdx_config._update_eval_interval_by_epoch(
  35. self.train_config.eval_interval)
  36. if self.train_config.save_interval:
  37. self.pdx_config.update_save_interval(
  38. self.train_config.save_interval)
  39. self.pdx_config.update_dataset(self.global_config.dataset_dir,
  40. "PubTabTableRecDataset")
  41. if self.train_config.pretrain_weight_path:
  42. self.pdx_config.update_pretrained_weights(
  43. self.train_config.pretrain_weight_path)
  44. if self.train_config.batch_size is not None:
  45. self.pdx_config.update_batch_size(self.train_config.batch_size)
  46. if self.train_config.learning_rate is not None:
  47. self.pdx_config.update_learning_rate(
  48. self.train_config.learning_rate)
  49. if self.train_config.epochs_iters is not None:
  50. self.pdx_config._update_epochs(self.train_config.epochs_iters)
  51. if self.train_config.resume_path is not None and self.train_config.resume_path != "":
  52. self.pdx_config._update_checkpoints(self.train_config.resume_path)
  53. if self.global_config.output is not None:
  54. self.pdx_config._update_output_dir(self.global_config.output)
  55. def get_train_kwargs(self) -> dict:
  56. """get key-value arguments of model training function
  57. Returns:
  58. dict: the arguments of training function.
  59. """
  60. return {"device": self.get_device()}
  61. class TableRecTrainDeamon(BaseTrainDeamon):
  62. """ TableRecTrainDeamon """
  63. def __init__(self, *args, **kwargs):
  64. super().__init__(*args, **kwargs)
  65. def get_the_pdparams_suffix(self):
  66. """ get the suffix of pdparams file """
  67. return "pdparams"
  68. def get_the_pdema_suffix(self):
  69. """ get the suffix of pdema file """
  70. return "pdema"
  71. def get_the_pdopt_suffix(self):
  72. """ get the suffix of pdopt file """
  73. return "pdopt"
  74. def get_the_pdstates_suffix(self):
  75. """ get the suffix of pdstates file """
  76. return "states"
  77. def get_ith_ckp_prefix(self, epoch_id):
  78. """ get the prefix of the epoch_id checkpoint file """
  79. return f"iter_epoch_{epoch_id}"
  80. def get_best_ckp_prefix(self):
  81. """ get the prefix of the best checkpoint file """
  82. return "best_accuracy"
  83. def get_score(self, pdstates_path):
  84. """ get the score by pdstates file """
  85. if not Path(pdstates_path).exists():
  86. return 0
  87. return paddle.load(pdstates_path)['best_model_dict']['acc']
  88. def get_epoch_id_by_pdparams_prefix(self, pdparams_prefix):
  89. """ get the epoch_id by pdparams file """
  90. return int(pdparams_prefix.split(".")[0].split("_")[-1])