From 849168bfa2f984dce36e1124220d10d8dfb8226f Mon Sep 17 00:00:00 2001 From: tian <11429339@qq.com> Date: Wed, 10 Jun 2026 10:53:02 +0800 Subject: [PATCH] =?UTF-8?q?fix:=20=E8=AE=AD=E7=BB=83=E5=90=8E=E4=BF=9D?= =?UTF-8?q?=E5=AD=98=E6=89=80=E6=9C=89=E6=A8=A1=E5=9E=8B=E8=80=8C=E9=9D=9E?= =?UTF-8?q?=E4=BB=85=E6=9C=80=E4=BC=98;=20=E9=A2=84=E6=B5=8B=E9=A1=B5?= =?UTF-8?q?=E5=8F=AF=E5=88=87=E6=8D=A2=E4=BB=BB=E6=84=8F=E6=A8=A1=E5=9E=8B?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/model_trainer.py | 30 ++++++++---------------------- 1 file changed, 8 insertions(+), 22 deletions(-) diff --git a/src/model_trainer.py b/src/model_trainer.py index 7d7dfc5..3fef1e8 100644 --- a/src/model_trainer.py +++ b/src/model_trainer.py @@ -393,6 +393,8 @@ class ModelTrainer: all_metrics = {} best_model = None + best_score = -float('inf') + self.best_model = None best_score = float('-inf') best_model_type = None # 添加变量记录最佳模型类型 @@ -582,7 +584,7 @@ class ModelTrainer: 'validation': val_metrics } - # 更新最佳模型（不包括PLS） + # 跟踪最佳模型（PLS不参与） current_score = val_metrics['r2'] if val_metrics else train_metrics['r2'] if model_type != 'pls' and current_score > best_score: best_score = current_score @@ -592,33 +594,17 @@ class ModelTrainer: 'mae': val_metrics['mae'] if val_metrics else train_metrics['mae'], 'rmse': val_metrics['rmse'] if val_metrics else train_metrics['rmse'] } - best_model_type = model_type # 记录最佳模型类型 - - # 保存最佳模型实例（但不立即写入数据库） - if model_type == 'pytorch': - self.best_pytorch_model = self.model.state_dict() # 保存模型状态 - self.best_pytorch_metrics = all_metrics[model_type] # 保存指标 - else: + if model_type != 'pytorch': self.best_model = model - self.best_model_metrics = all_metrics[model_type] - # 单独保存PLS模型（不参与最佳模型评选） - if model_type == 'pls' and equipment_type: + # 保存每个模型 + if model_type == 'pytorch': + self.save_model(equipment_type, all_metrics[model_type]) + elif equipment_type: self._save_sklearn_model(equipment_type, model_type, model, all_metrics[model_type]) - # 在所有模型训练完成后，只保存最佳模型 - if best_model_type and equipment_type: - if best_model_type == 'pytorch': - # 恢复最佳PyTorch模型状态并保存 - self.model.load_state_dict(self.best_pytorch_model) - self.save_model(equipment_type, self.best_pytorch_metrics) - else: - # 保存最佳sklearn模型 - self._save_sklearn_model(equipment_type, best_model_type, self.best_model, self.best_model_metrics) - return { 'metrics': all_metrics, - 'feature_importance': None, 'best_model': best_model }