train.py 1.5 KB

1234567891011121314151617181920212223242526272829303132333435363738394041
  1. # -*- coding: utf-8 -*-
  2. """
  3. @author: yq
  4. @time: 2024/11/1
  5. @desc: 模型训练管道
  6. """
  7. from entitys import DataSplitEntity
  8. from feature.filter_strategy_base import FilterStrategyBase
  9. from init import init
  10. from model import ModelBase
  11. from monitor.report_generate import Report
  12. init()
  13. class TrainPipeline():
  14. def __init__(self, filter_strategy: FilterStrategyBase, model: ModelBase, data: DataSplitEntity):
  15. self._filter_strategy = filter_strategy
  16. self._model = model
  17. self._data = data
  18. self._model._train_config.set_save_path_func(self._filter_strategy.data_process_config.f_get_save_path)
  19. def train(self, ):
  20. # 处理数据,获取候选特征
  21. candidate_feature = self._filter_strategy.filter(self._data)
  22. # 生成训练数据
  23. data_prepared = self._filter_strategy.feature_generate(self._data, candidate_feature)
  24. # 特征信息
  25. metric_value_dict_feature = self._filter_strategy.feature_report(self._data, candidate_feature)
  26. metric_value_dict_train = self._model.train(data_prepared, *data_prepared.args, **data_prepared.kwargs)
  27. self.metric_value_dict = {**metric_value_dict_feature, **metric_value_dict_train}
  28. def generate_report(self, ):
  29. Report.generate_report(self.metric_value_dict, self._model.get_template_path(),
  30. save_path=self._filter_strategy.data_process_config.f_get_save_path("模型报告.docx"))
  31. if __name__ == "__main__":
  32. pass