422 lines
18 KiB
Python
422 lines
18 KiB
Python
import numpy as np
|
|
import pandas as pd
|
|
from sklearn.preprocessing import StandardScaler
|
|
from sklearn.model_selection import cross_val_score
|
|
from sklearn.ensemble import RandomForestRegressor, GradientBoostingRegressor
|
|
from sklearn.impute import SimpleImputer
|
|
import xgboost as xgb
|
|
import lightgbm as lgb
|
|
import logging
|
|
import joblib
|
|
import os
|
|
from src.feature_analysis import FeatureAnalysis
|
|
from datetime import datetime
|
|
import json
|
|
from src.database import get_db_connection
|
|
from src.data_preparation import DataPreparation
|
|
|
|
class ModelTrainer:
|
|
def __init__(self):
|
|
self.models = {
|
|
'xgboost': self._create_xgboost_model(),
|
|
'lightgbm': self._create_lightgbm_model(),
|
|
'gbdt': self._create_gbdt_model(),
|
|
'rf': self._create_rf_model()
|
|
}
|
|
self.best_model = None
|
|
self.imputer = SimpleImputer(strategy='mean')
|
|
self.feature_scaler = None
|
|
self.target_scaler = None
|
|
|
|
def fit_model(self, X_train, y_train, model_names, X_val=None, y_val=None, equipment_type=None):
|
|
"""
|
|
训练模型并返回评估结果
|
|
"""
|
|
try:
|
|
# 记录数据范围
|
|
logging.info(f"Training data range - X: min={X_train.min()}, max={X_train.max()}")
|
|
logging.info(f"Training data range - y: min={y_train.min()}, max={y_train.max()}")
|
|
|
|
results = {}
|
|
best_score = -float('inf')
|
|
best_model_info = None
|
|
|
|
for model_name in model_names:
|
|
if model_name not in self.models:
|
|
logging.warning(f"Unknown model: {model_name}")
|
|
continue
|
|
|
|
logging.info(f"Training {model_name}...")
|
|
model = self.models[model_name]
|
|
|
|
# 训练模型
|
|
model.fit(X_train, y_train)
|
|
|
|
# 计算评估指标
|
|
metrics = self._calculate_metrics(
|
|
model,
|
|
X_train, y_train,
|
|
X_val, y_val
|
|
)
|
|
|
|
# 更新最佳模型
|
|
if metrics['validation']['r2'] > best_score:
|
|
best_score = metrics['validation']['r2']
|
|
self.best_model = model
|
|
best_model_info = {
|
|
'type': model_name,
|
|
'r2': float(metrics['validation']['r2']),
|
|
'mae': float(metrics['validation']['mae']) if metrics['validation']['mae'] is not None else None,
|
|
'rmse': float(metrics['validation']['rmse']) if metrics['validation']['rmse'] is not None else None
|
|
}
|
|
|
|
# 转换 numpy 数据类型为 Python 原生类型
|
|
results[model_name] = {
|
|
'train': {
|
|
'r2': float(metrics['train']['r2']),
|
|
'mae': float(metrics['train']['mae']),
|
|
'rmse': float(metrics['train']['rmse'])
|
|
},
|
|
'validation': {
|
|
'r2': float(metrics['validation']['r2']),
|
|
'mae': float(metrics['validation']['mae']) if metrics['validation']['mae'] is not None else None,
|
|
'rmse': float(metrics['validation']['rmse']) if metrics['validation']['rmse'] is not None else None
|
|
}
|
|
}
|
|
|
|
# 保存最佳模型
|
|
if equipment_type and best_model_info:
|
|
self._save_best_model(equipment_type, best_model_info, X_train)
|
|
|
|
# 转换特征重要性为列表
|
|
feature_importance = None
|
|
if self.best_model and hasattr(self.best_model, 'feature_importances_'):
|
|
feature_importance = [float(x) for x in self.best_model.feature_importances_]
|
|
|
|
return {
|
|
'metrics': results,
|
|
'best_model': best_model_info,
|
|
'feature_importance': feature_importance
|
|
}
|
|
|
|
except Exception as e:
|
|
logging.error(f"Error in model training: {str(e)}")
|
|
raise
|
|
|
|
def _calculate_metrics(self, model, X_train, y_train, X_val=None, y_val=None):
|
|
"""
|
|
计算模型评估指标
|
|
"""
|
|
from sklearn.metrics import r2_score, mean_absolute_error, mean_squared_error
|
|
|
|
# 训练集评估
|
|
train_pred = model.predict(X_train)
|
|
|
|
# 如果使用了标准化，需要转换回原始范围
|
|
if hasattr(self, 'target_scaler'):
|
|
train_pred = self.target_scaler.inverse_transform(train_pred.reshape(-1, 1)).ravel()
|
|
y_train_orig = self.target_scaler.inverse_transform(y_train.reshape(-1, 1)).ravel()
|
|
else:
|
|
y_train_orig = y_train
|
|
|
|
# 记录预测范围
|
|
logging.info(f"Train predictions range: min={train_pred.min()}, max={train_pred.max()}")
|
|
logging.info(f"Train actual range: min={y_train_orig.min()}, max={y_train_orig.max()}")
|
|
|
|
train_metrics = {
|
|
'r2': r2_score(y_train_orig, train_pred),
|
|
'mae': mean_absolute_error(y_train_orig, train_pred),
|
|
'rmse': np.sqrt(mean_squared_error(y_train_orig, train_pred))
|
|
}
|
|
|
|
# 验证集评估
|
|
if X_val is not None and y_val is not None:
|
|
val_pred = model.predict(X_val)
|
|
|
|
# 如果使用了标准化，需要转换回原始范围
|
|
if hasattr(self, 'target_scaler'):
|
|
val_pred = self.target_scaler.inverse_transform(val_pred.reshape(-1, 1)).ravel()
|
|
y_val_orig = self.target_scaler.inverse_transform(y_val.reshape(-1, 1)).ravel()
|
|
else:
|
|
y_val_orig = y_val
|
|
|
|
# 记录预测范围
|
|
logging.info(f"Validation predictions range: min={val_pred.min()}, max={val_pred.max()}")
|
|
logging.info(f"Validation actual range: min={y_val_orig.min()}, max={y_val_orig.max()}")
|
|
|
|
val_metrics = {
|
|
'r2': r2_score(y_val_orig, val_pred),
|
|
'mae': mean_absolute_error(y_val_orig, val_pred),
|
|
'rmse': np.sqrt(mean_squared_error(y_val_orig, val_pred))
|
|
}
|
|
else:
|
|
# 使用交叉验证
|
|
cv_scores = cross_val_score(model, X_train, y_train, cv=5)
|
|
val_metrics = {
|
|
'r2': cv_scores.mean(),
|
|
'mae': None,
|
|
'rmse': None
|
|
}
|
|
|
|
return {
|
|
'train': train_metrics,
|
|
'validation': val_metrics
|
|
}
|
|
|
|
def _create_xgboost_model(self):
|
|
"""
|
|
创建 XGBoost 模型，增强正则化
|
|
"""
|
|
return xgb.XGBRegressor(
|
|
n_estimators=50, # 减少树的数量
|
|
learning_rate=0.05, # 减小学习率
|
|
max_depth=3, # 减小树的深度
|
|
min_child_weight=3, # 增加最小子节点权重
|
|
subsample=0.7, # 减小样本采样比例
|
|
colsample_bytree=0.7, # 减小特征采样比例
|
|
reg_alpha=0.1, # L1 正则化
|
|
reg_lambda=1, # L2 正则化
|
|
random_state=42
|
|
)
|
|
|
|
def _create_lightgbm_model(self):
|
|
"""
|
|
创建 LightGBM 模型，增强正则化
|
|
"""
|
|
return lgb.LGBMRegressor(
|
|
n_estimators=50,
|
|
learning_rate=0.05,
|
|
max_depth=3,
|
|
num_leaves=7,
|
|
min_data_in_leaf=3,
|
|
min_sum_hessian_in_leaf=1e-3,
|
|
subsample=0.7,
|
|
colsample_bytree=0.7,
|
|
reg_alpha=0.1,
|
|
reg_lambda=1,
|
|
random_state=42,
|
|
verbose=-1
|
|
)
|
|
|
|
def _create_gbdt_model(self):
|
|
"""
|
|
创建 GBDT 模型，增强正则化以减轻过拟合
|
|
"""
|
|
return GradientBoostingRegressor(
|
|
n_estimators=20, # 减少树的数量
|
|
learning_rate=0.01, # 减小学习率
|
|
max_depth=2, # 减小树的深度
|
|
min_samples_split=4, # 增加分裂所需的最小样本数
|
|
min_samples_leaf=3, # 增加叶子节点最小样本数
|
|
subsample=0.5, # 减小样本采样比例
|
|
random_state=42,
|
|
validation_fraction=0.2 # 使用部分训练数据作为验证集
|
|
)
|
|
|
|
def _create_rf_model(self):
|
|
"""
|
|
创建随机森林模型，针对小样本数据调整参数
|
|
"""
|
|
return RandomForestRegressor(
|
|
n_estimators=100, # 增加树的数量
|
|
max_depth=4, # 限制树的深度
|
|
min_samples_split=2, # 减小分需的最小样本数
|
|
min_samples_leaf=1, # 减小叶子节点最小样本数
|
|
max_features='sqrt', # 特征采样
|
|
bootstrap=True, # 使用 bootstrap 采样
|
|
oob_score=True, # 计算袋外分数
|
|
random_state=42
|
|
)
|
|
|
|
def _save_best_model(self, equipment_type, best_model_info, X_train):
|
|
"""
|
|
保存最佳模型
|
|
"""
|
|
try:
|
|
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
|
model_dir = 'models'
|
|
os.makedirs(model_dir, exist_ok=True)
|
|
|
|
# 保存模型文件
|
|
model_path = f'{model_dir}/{equipment_type}_{timestamp}'
|
|
if isinstance(self.best_model, xgb.XGBRegressor):
|
|
self.best_model.save_model(f'{model_path}.json')
|
|
model_format = 'json'
|
|
else:
|
|
joblib.dump(self.best_model, f'{model_path}.joblib')
|
|
model_format = 'joblib'
|
|
|
|
# 验证标准化器
|
|
if not isinstance(self.feature_scaler, StandardScaler):
|
|
raise ValueError("Invalid feature scaler")
|
|
if not isinstance(self.target_scaler, StandardScaler):
|
|
raise ValueError("Invalid target scaler")
|
|
|
|
# 保存标准化器
|
|
scaler_path = f'{model_dir}/{equipment_type}_{timestamp}_scaler.joblib'
|
|
joblib.dump({
|
|
'feature_scaler': self.feature_scaler,
|
|
'target_scaler': self.target_scaler
|
|
}, scaler_path)
|
|
|
|
logging.info(f"Saved model to {model_path}.{model_format}")
|
|
logging.info(f"Saved scalers to {scaler_path}")
|
|
|
|
# 更新数据库中的模型记录
|
|
with get_db_connection() as conn:
|
|
cursor = conn.cursor()
|
|
|
|
# 将之前的激活模型设置为非激活
|
|
cursor.execute("""
|
|
UPDATE trained_models
|
|
SET is_active = FALSE
|
|
WHERE equipment_type = %s
|
|
""", (equipment_type,))
|
|
|
|
# 插入新的模型记录
|
|
cursor.execute("""
|
|
INSERT INTO trained_models (
|
|
model_name, model_type, equipment_type, model_path,
|
|
scaler_path, r2_score, mae, rmse, feature_importance,
|
|
training_data_size, training_date, is_active, created_by
|
|
) VALUES (%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, NOW(), TRUE, 'system')
|
|
""", (
|
|
f'{best_model_info["type"]}_{timestamp}',
|
|
best_model_info["type"],
|
|
equipment_type,
|
|
f'{model_path}.{model_format}',
|
|
scaler_path,
|
|
best_model_info["r2"],
|
|
best_model_info["mae"],
|
|
best_model_info["rmse"],
|
|
json.dumps(self.feature_importance) if hasattr(self, 'feature_importance') else None,
|
|
len(X_train)
|
|
))
|
|
|
|
conn.commit()
|
|
|
|
logging.info(f"Best model saved: {model_path}")
|
|
return True
|
|
|
|
except Exception as e:
|
|
logging.error(f"Error saving best model: {str(e)}")
|
|
return False
|
|
|
|
def load_model(self, equipment_type):
|
|
"""
|
|
加载已训练的模型
|
|
"""
|
|
try:
|
|
logging.info(f"Loading model for {equipment_type}")
|
|
|
|
# 从数据库获最新的激活模型
|
|
with get_db_connection() as conn:
|
|
cursor = conn.cursor(dictionary=True)
|
|
cursor.execute("""
|
|
SELECT * FROM trained_models
|
|
WHERE equipment_type = %s AND is_active = TRUE
|
|
ORDER BY training_date DESC LIMIT 1
|
|
""", (equipment_type,))
|
|
|
|
model_record = cursor.fetchone()
|
|
if not model_record:
|
|
raise ValueError(f"No active model found for {equipment_type}")
|
|
|
|
logging.info(f"Found model: {model_record['model_name']}")
|
|
logging.info(f"Model path: {model_record['model_path']}")
|
|
logging.info(f"Scaler path: {model_record['scaler_path']}")
|
|
|
|
# 检查文件是否存在
|
|
if not os.path.exists(model_record['model_path']):
|
|
raise FileNotFoundError(f"Model file not found: {model_record['model_path']}")
|
|
if not os.path.exists(model_record['scaler_path']):
|
|
raise FileNotFoundError(f"Scaler file not found: {model_record['scaler_path']}")
|
|
|
|
# 加载模型文件
|
|
if model_record['model_type'] == 'xgboost':
|
|
self.best_model = xgb.XGBRegressor()
|
|
self.best_model.load_model(model_record['model_path'])
|
|
else:
|
|
self.best_model = joblib.load(model_record['model_path'])
|
|
|
|
# 加载标准化器
|
|
try:
|
|
scalers = joblib.load(model_record['scaler_path'])
|
|
logging.info(f"Loaded scalers: {scalers.keys()}")
|
|
|
|
if 'feature_scaler' not in scalers or 'target_scaler' not in scalers:
|
|
raise ValueError("Missing scalers in saved file")
|
|
|
|
self.feature_scaler = scalers['feature_scaler']
|
|
self.target_scaler = scalers['target_scaler']
|
|
|
|
# 验证标准化器
|
|
if not hasattr(self.feature_scaler, 'transform') or not hasattr(self.target_scaler, 'transform'):
|
|
raise ValueError("Invalid scaler objects")
|
|
|
|
logging.info("Model and scalers loaded successfully")
|
|
logging.info(f"Feature scaler type: {type(self.feature_scaler)}")
|
|
logging.info(f"Target scaler type: {type(self.target_scaler)}")
|
|
|
|
except Exception as e:
|
|
logging.error(f"Error loading scalers: {str(e)}")
|
|
logging.error(f"Scaler file content: {scalers if 'scalers' in locals() else 'Not loaded'}")
|
|
raise ValueError(f"Failed to load scalers: {str(e)}")
|
|
|
|
return True
|
|
|
|
except Exception as e:
|
|
logging.error(f"Error loading model: {str(e)}")
|
|
logging.error("Detailed traceback:", exc_info=True)
|
|
return False
|
|
|
|
def predict(self, features):
|
|
"""
|
|
使用加载的模型进行预测
|
|
"""
|
|
try:
|
|
if not self.best_model:
|
|
raise ValueError("No model loaded")
|
|
|
|
if not self.feature_scaler:
|
|
raise ValueError("Feature scaler not loaded")
|
|
|
|
if not self.target_scaler:
|
|
raise ValueError("Target scaler not loaded")
|
|
|
|
logging.info("Starting prediction")
|
|
logging.info(f"Input features shape: {features.shape}")
|
|
logging.info(f"Input features: \n{features}")
|
|
|
|
# 处理缺失值
|
|
features_filled = np.array(features, dtype=float)
|
|
features_filled[np.isnan(features_filled)] = 0
|
|
features_filled = np.nan_to_num(features_filled, 0)
|
|
|
|
logging.info(f"Filled features: \n{features_filled}")
|
|
|
|
# 标准化特征
|
|
X = self.feature_scaler.transform(features_filled)
|
|
logging.info(f"Transformed features shape: {X.shape}")
|
|
logging.info(f"Transformed features: \n{X}")
|
|
|
|
# 预测
|
|
y_pred_scaled = self.best_model.predict(X)
|
|
logging.info(f"Scaled prediction shape: {y_pred_scaled.shape}")
|
|
logging.info(f"Scaled prediction: {y_pred_scaled}")
|
|
|
|
# 反标准化
|
|
y_pred = self.target_scaler.inverse_transform(y_pred_scaled.reshape(-1, 1))
|
|
logging.info(f"Final prediction shape: {y_pred.shape}")
|
|
logging.info(f"Final prediction: {y_pred}")
|
|
|
|
# 记录标准化器的参数
|
|
logging.info("Target scaler params:")
|
|
logging.info(f"Mean: {self.target_scaler.mean_}")
|
|
logging.info(f"Scale: {self.target_scaler.scale_}")
|
|
|
|
return y_pred.ravel()
|
|
|
|
except Exception as e:
|
|
logging.error(f"Error in prediction: {str(e)}")
|
|
raise |