tian/CostPrediction
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: 训练后保存所有模型而非仅最优; 预测页可切换任意模型

Browse Source
This commit is contained in:
tian 2026-06-10 10:53:02 +08:00
parent 8be35bfebf
commit 849168bfa2
1 changed files with 8 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

30
src/model_trainer.py
View File

@ -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
}