import numpy as np from sklearn.preprocessing import StandardScaler from sklearn.model_selection import train_test_split from sklearn.metrics import r2_score, mean_absolute_error, mean_squared_error import logging import os from datetime import datetime import json from src.feature_analysis import FeatureAnalysis from src.database import get_db_connection from .logger import setup_logger import math # PyTorch 为可选依赖 try: import torch import torch.nn as nn from torch.utils.data import DataLoader from src.data_preparation import EquipmentDataset _HAS_TORCH = True except ImportError: torch = None nn = None DataLoader = None EquipmentDataset = None _HAS_TORCH = False logger = setup_logger(__name__) # 条件基类：有 PyTorch 时继承 nn.Module，否则继承 object _PYTORCH_BASE = nn.Module if _HAS_TORCH else object class CostPredictionModel(_PYTORCH_BASE): def __init__(self, input_size, equipment_type): if not _HAS_TORCH: raise ImportError("PyTorch is not installed. Install with: pip install torch") super().__init__() self.equipment_type = equipment_type if equipment_type == '火箭炮': self.net = nn.Sequential( nn.Linear(input_size, 32), nn.ReLU(), nn.BatchNorm1d(32), nn.Linear(32, 16), nn.ReLU(), nn.BatchNorm1d(16), nn.Linear(16, 8), nn.ReLU(), nn.BatchNorm1d(8), nn.Linear(8, 1) ) def init_weights(m): if isinstance(m, nn.Linear): torch.nn.init.orthogonal_(m.weight, gain=0.5) torch.nn.init.constant_(m.bias, 0.0) elif isinstance(m, nn.BatchNorm1d): torch.nn.init.constant_(m.weight, 0.5) torch.nn.init.constant_(m.bias, 0.0) self.net.apply(init_weights) else: self.manufacturer_net = nn.Sequential( nn.Linear(5, 4), nn.ReLU(), nn.BatchNorm1d(4), nn.Dropout(0.2) ) self.equipment_net = nn.Sequential( nn.Linear(input_size - 5, 64), nn.LeakyReLU(0.1), nn.BatchNorm1d(64), nn.Dropout(0.2), nn.Linear(64, 32), nn.LeakyReLU(0.1), nn.BatchNorm1d(32), nn.Dropout(0.2), nn.Linear(32, 16), nn.LeakyReLU(0.1), nn.BatchNorm1d(16), nn.Dropout(0.2) ) self.combined_net = nn.Sequential( nn.Linear(20, 32), nn.LeakyReLU(0.1), nn.BatchNorm1d(32), nn.Dropout(0.2), nn.Linear(32, 16), nn.LeakyReLU(0.1), nn.BatchNorm1d(16), nn.Dropout(0.2), nn.Linear(16, 8), nn.LeakyReLU(0.1), nn.BatchNorm1d(8), nn.Linear(8, 1) ) def forward(self, x): if self.equipment_type == '火箭炮': return self.net(x) else: manufacturer_features = x[:, -5:] equipment_features = x[:, :-5] manu_out = self.manufacturer_net(manufacturer_features) equip_out = self.equipment_net(equipment_features) combined = torch.cat([equip_out, manu_out], dim=1) return self.combined_net(combined) class ModelTrainer: def __init__(self): if not _HAS_TORCH: raise ImportError("PyTorch is not installed. Install with: pip install torch") self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') self.model = None self.feature_scaler = None self.target_scaler = None self.equipment_type = None self.feature_analyzer = FeatureAnalysis() def train_model(self, dataloader, epochs=100, learning_rate=0.001, equipment_type=None): """训练模型""" try: sample_features, _ = next(iter(dataloader)) input_size = sample_features.shape[1] # 设置确定性 torch.manual_seed(42) torch.backends.cudnn.deterministic = True torch.backends.cudnn.benchmark = False if torch.cuda.is_available(): torch.cuda.manual_seed_all(42) np.random.seed(42) self.model = CostPredictionModel(input_size, equipment_type).to(self.device) if equipment_type == '火箭炮': # 火箭炮使用更保守和稳定的训练设置 criterion = nn.SmoothL1Loss(beta=0.1) # 使用Huber损失，beta值较小 learning_rate = 0.0003 # 更小的学习率 weight_decay = 0.001 # 适中的权重衰减 # 使用AdamW优化器，更小的beta值 optimizer = torch.optim.AdamW( self.model.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=(0.8, 0.9), # 更小的动量值 eps=1e-8 ) # 使用带预热的学习率调度 num_steps = len(dataloader) * epochs warmup_steps = num_steps // 10 # 10%的预热步数 def lr_lambda(current_step): if current_step < warmup_steps: return float(current_step) / float(max(1, warmup_steps)) return 0.5 * (1.0 + math.cos( math.pi * (current_step - warmup_steps) / float(max(1, num_steps - warmup_steps)) )) scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda) else: # 巡飞弹保持原有设置 # 巡飞弹使用更激进的训练设置 criterion = nn.MSELoss() learning_rate = 0.001 # 较大的学习率 weight_decay = 0.001 # 较小的权重衰减 patience = 20 # 较短的耐心值 # 使用Adam优化器 optimizer = torch.optim.Adam( self.model.parameters(), lr=learning_rate, weight_decay=weight_decay, betas=(0.9, 0.999) ) # 使用余弦退火学习率调度 scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( optimizer, T_max=epochs, eta_min=learning_rate * 0.01 ) # 训练循环 best_loss = float('inf') patience = 30 patience_counter = 0 best_model_state = None moving_avg_loss = None alpha = 0.9 # 移动平均系数 for epoch in range(epochs): self.model.train() total_loss = 0 batch_count = 0 for batch_features, batch_targets in dataloader: batch_features = batch_features.to(self.device) batch_targets = batch_targets.to(self.device) # 前向传播 outputs = self.model(batch_features) loss = criterion(outputs, batch_targets.view(-1, 1)) # 反向传播 optimizer.zero_grad(set_to_none=True) # 更高效的梯度清零 loss.backward() # 梯度裁剪 if equipment_type == '火箭炮': torch.nn.utils.clip_grad_norm_( self.model.parameters(), max_norm=0.1 ) optimizer.step() if equipment_type == '火箭炮': scheduler.step() total_loss += loss.item() batch_count += 1 avg_loss = total_loss / batch_count # 使用移动平均计算损失 if moving_avg_loss is None: moving_avg_loss = avg_loss else: moving_avg_loss = alpha * moving_avg_loss + (1 - alpha) * avg_loss # 早停检查使用移动平均损失 if moving_avg_loss < best_loss: best_loss = moving_avg_loss patience_counter = 0 best_model_state = { 'state_dict': self.model.state_dict(), 'optimizer': optimizer.state_dict(), 'scheduler': scheduler.state_dict() if equipment_type == '火箭炮' else None } else: patience_counter += 1 if (epoch + 1) % 10 == 0: logger.info(f'Epoch [{epoch+1}/{epochs}], Loss: {moving_avg_loss:.4f}, ' f'LR: {optimizer.param_groups[0]["lr"]:.6f}') if patience_counter >= patience: logger.info(f"Early stopping triggered at epoch {epoch+1}") break # 恢复最佳模型 if best_model_state is not None: self.model.load_state_dict(best_model_state['state_dict']) optimizer.load_state_dict(best_model_state['optimizer']) if equipment_type == '火箭炮' and best_model_state['scheduler']: scheduler.load_state_dict(best_model_state['scheduler']) return True except Exception as e: logger.error(f"Error in model training: {str(e)}") raise def save_model(self, equipment_type, metrics=None): """保存模型""" try: timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") model_dir = 'models' os.makedirs(model_dir, exist_ok=True) # 转换装备类型为英文 equipment_type_en = 'rocket' if equipment_type == '火箭炮' else 'missile' # 保存模型 model_path = f'{model_dir}/{equipment_type_en}_{timestamp}.pth' torch.save({ 'model_state_dict': self.model.state_dict(), 'input_size': self.model.equipment_net[0].in_features + 5, 'manufacturer_net_state': self.model.manufacturer_net.state_dict(), 'equipment_net_state': self.model.equipment_net.state_dict(), 'combined_net_state': self.model.combined_net.state_dict() }, model_path) # 保存标准化器 scaler_path = f'{model_dir}/{equipment_type_en}_{timestamp}_scaler.pth' torch.save({ 'feature_scaler': self.feature_scaler, 'target_scaler': self.target_scaler }, scaler_path) # 获取评估指标 r2 = metrics['validation']['r2'] if metrics and metrics.get('validation') else metrics['train']['r2'] mae = metrics['validation']['mae'] if metrics and metrics.get('validation') else metrics['train']['mae'] rmse = metrics['validation']['rmse'] if metrics and metrics.get('validation') else metrics['train']['rmse'] # 更新数据库 with get_db_connection() as conn: cursor = conn.cursor() # 将所有同类型模型设置为非激活, 除了 PLS 模型 cursor.execute(""" UPDATE trained_models SET is_active = FALSE WHERE equipment_type = ? AND model_type != ? """, (equipment_type, 'pls')) # 保存新模型记录 cursor.execute(""" INSERT INTO trained_models ( model_name, model_type, equipment_type, model_path, scaler_path, training_date, is_active, created_by, r2_score, mae, rmse ) VALUES (?, ?, ?, ?, ?, datetime('now','localtime'), TRUE, ?, ?, ?, ?) """, ( f"{equipment_type}_{timestamp}", 'pytorch', equipment_type, model_path, scaler_path, 'system', r2, mae, rmse )) conn.commit() logger.info(f"Model saved successfully: {model_path}") return True except Exception as e: logger.error(f"Error saving model: {str(e)}") return False def load_model(self, equipment_type, model_type): """加载模型""" try: # 从数据库获取最新的激活模型 with get_db_connection() as conn: cursor = conn.cursor() cursor.execute(""" SELECT * FROM trained_models WHERE equipment_type = ? AND model_type = ? AND is_active = TRUE ORDER BY training_date DESC LIMIT 1 """, (equipment_type, model_type)) model_record = cursor.fetchone() if not model_record: raise ValueError(f"No trained model found for {equipment_type}") # 加载模型 if model_record['model_type'] == 'pytorch': # 加载PyTorch模型 checkpoint = torch.load(model_record['model_path'], encoding='latin1') input_size = checkpoint['input_size'] # 创建新模型实例 self.model = CostPredictionModel(input_size, equipment_type).to(self.device) self.model.load_state_dict(checkpoint['model_state_dict']) # 加载标准化器 scalers = torch.load(model_record['scaler_path'], encoding='latin1') self.feature_scaler = scalers['feature_scaler'] self.target_scaler = scalers['target_scaler'] else: # 加载sklearn模型 from joblib import load with open(model_record['model_path'], 'rb') as f: self.model = load(f) with open(model_record['scaler_path'], 'rb') as f: scalers = load(f) self.feature_scaler = scalers['feature_scaler'] self.target_scaler = scalers['target_scaler'] logger.info(f"Model loaded successfully from {model_record['model_path']}") return True except Exception as e: logger.error(f"Error loading model: {str(e)}") return False def predict(self, features): """使用模型进行预测""" try: self.model.eval() with torch.no_grad(): features_tensor = torch.FloatTensor(features).to(self.device) predictions = self.model(features_tensor) # 直接返回预测值 return predictions.cpu().numpy() except Exception as e: logger.error(f"Error in prediction: {str(e)}") raise def fit_model(self, X_train, y_train, models, X_val=None, y_val=None, equipment_type=None): """训练模型并返回评估结果""" try: logger.info(f"Starting model training for {equipment_type}") logger.info(f"Selected models: {models}") logger.info(f"Training data shape: {X_train.shape}") all_metrics = {} best_model = None best_score = float('-inf') best_model_type = None # 添加变量记录最佳模型类型 # 训练所有选择的模型 for model_type in models: logger.info(f"Training {model_type} model...") if model_type == 'pls': # PLS模型单独处理，不参与最优模型评选 from sklearn.cross_decomposition import PLSRegression # 使用较少的组件数来避免过拟合 n_components = min(3, X_train.shape[1] // 5) model = PLSRegression( n_components=n_components, scale=True, max_iter=500, tol=1e-6 ) model.fit(X_train, y_train) # 评估PLS模型 y_train_pred = model.predict(X_train).ravel() if X_val is not None: y_val_pred = model.predict(X_val).ravel() # 将预测值转换回原始尺度 y_train_pred_original = self.target_scaler.inverse_transform(y_train_pred.reshape(-1, 1)).ravel() y_train_original = self.target_scaler.inverse_transform(y_train.reshape(-1, 1)).ravel() train_metrics = { 'r2': float(r2_score(y_train_original, y_train_pred_original)), 'mae': float(mean_absolute_error(y_train_original, y_train_pred_original)), 'rmse': float(np.sqrt(mean_squared_error(y_train_original, y_train_pred_original))) } val_metrics = None if X_val is not None: y_val_pred_original = self.target_scaler.inverse_transform(y_val_pred.reshape(-1, 1)).ravel() y_val_original = self.target_scaler.inverse_transform(y_val.reshape(-1, 1)).ravel() val_metrics = { 'r2': float(r2_score(y_val_original, y_val_pred_original)), 'mae': float(mean_absolute_error(y_val_original, y_val_pred_original)), 'rmse': float(np.sqrt(mean_squared_error(y_val_original, y_val_pred_original))) } all_metrics[model_type] = { 'train': train_metrics, 'validation': val_metrics } # 保存PLS模型，但不参与最优模型评选 if equipment_type: self._save_sklearn_model(equipment_type, model_type, model, all_metrics[model_type]) continue # 跳过后续的最优模型评选 elif model_type == 'xgboost': import xgboost as xgb model = xgb.XGBRegressor( n_estimators=50, learning_rate=0.03, max_depth=3, min_child_weight=5, subsample=0.6, colsample_bytree=0.6, reg_alpha=0.5, reg_lambda=2.0, gamma=1, random_state=42 ) # 训练模型 model.fit(X_train, y_train) elif model_type == 'lightgbm': import lightgbm as lgb model = lgb.LGBMRegressor( n_estimators=50, learning_rate=0.03, max_depth=3, num_leaves=8, subsample=0.6, colsample_bytree=0.6, reg_alpha=0.5, reg_lambda=2.0, min_child_samples=10, min_split_gain=1.0, random_state=42 ) # 训练模型 model.fit(X_train, y_train) elif model_type == 'gbm': from sklearn.ensemble import GradientBoostingRegressor model = GradientBoostingRegressor( n_estimators=50, learning_rate=0.03, max_depth=3, min_samples_split=10, min_samples_leaf=5, subsample=0.6, min_impurity_decrease=0.01, random_state=42 ) # 训练模型 model.fit(X_train, y_train) elif model_type == 'rf': from sklearn.ensemble import RandomForestRegressor model = RandomForestRegressor( n_estimators=100, max_depth=4, min_samples_split=5, min_samples_leaf=3, max_features='sqrt', bootstrap=True, random_state=42 ) # 训练模型 model.fit(X_train, y_train) elif model_type == 'pytorch': # 训练PyTorch模型（如未安装则跳过） try: train_dataset = EquipmentDataset(X_train, y_train) train_loader = DataLoader(train_dataset, batch_size=32, shuffle=True) training_success = self.train_model(train_loader) if not training_success: continue except ImportError as e: logger.warning(f"PyTorch not available, skipping: {e}") continue # 评估模型性能 if model_type == 'pytorch': with torch.no_grad(): X_train_tensor = torch.FloatTensor(X_train).to(self.device) y_train_pred = self.model(X_train_tensor).cpu().numpy() # 直接获取输出 if X_val is not None: X_val_tensor = torch.FloatTensor(X_val).to(self.device) y_val_pred = self.model(X_val_tensor).cpu().numpy() # 直接获取输出 else: # 使用训练好的模型进行预测 y_train_pred = model.predict(X_train) if X_val is not None: if model_type == 'pls': y_val_pred = model.predict(X_val).ravel() # 记录PLS的一些额外信息 if hasattr(model, 'score'): train_r2 = model.score(X_train, y_train) val_r2 = model.score(X_val, y_val) logger.info(f"PLS built-in R² - Train: {train_r2:.4f}, Val: {val_r2:.4f}") # 记录每个组件解释的方差比例 if hasattr(model, 'explained_variance_ratio_'): logger.info("PLS explained variance ratios: " + ", ".join([f"{v:.4f}" for v in model.explained_variance_ratio_])) else: y_val_pred = model.predict(X_val) # 将测值转换回始尺度 y_train_pred_original = self.target_scaler.inverse_transform(y_train_pred.reshape(-1, 1)).ravel() y_train_original = self.target_scaler.inverse_transform(y_train.reshape(-1, 1)).ravel() train_metrics = { 'r2': float(r2_score(y_train_original, y_train_pred_original)), 'mae': float(mean_absolute_error(y_train_original, y_train_pred_original)), 'rmse': float(np.sqrt(mean_squared_error(y_train_original, y_train_pred_original))) } val_metrics = None if X_val is not None and y_val is not None: y_val_pred_original = self.target_scaler.inverse_transform(y_val_pred.reshape(-1, 1)).ravel() y_val_original = self.target_scaler.inverse_transform(y_val.reshape(-1, 1)).ravel() val_metrics = { 'r2': float(r2_score(y_val_original, y_val_pred_original)), 'mae': float(mean_absolute_error(y_val_original, y_val_pred_original)), 'rmse': float(np.sqrt(mean_squared_error(y_val_original, y_val_pred_original))) } all_metrics[model_type] = { 'train': train_metrics, 'validation': val_metrics } # 更新最佳模型（不包括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 best_model = { 'type': model_type, 'r2': current_score, '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: self.best_model = model self.best_model_metrics = all_metrics[model_type] # 单独保存PLS模型（不参与最佳模型评选） if model_type == 'pls' and 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 } except Exception as e: logger.error(f"Error in model fitting: {str(e)}") logger.error("Detailed traceback:", exc_info=True) raise def _calculate_feature_importance(self, X): """计算特征重要性""" try: if self.model is None: return None self.model.eval() feature_importance = np.zeros(X.shape[1]) # 使用特征扰动计算重要性 with torch.no_grad(): X_tensor = torch.FloatTensor(X).to(self.device) baseline_pred = self.model(X_tensor).cpu().numpy() # 直接获取预测值 for i in range(X.shape[1]): # 创建动后的特征 X_perturbed = X.copy() X_perturbed[:, i] = np.random.permutation(X_perturbed[:, i]) # 预测并计算影响 X_perturbed_tensor = torch.FloatTensor(X_perturbed).to(self.device) perturbed_pred = self.model(X_perturbed_tensor).cpu().numpy() # 直接获取预测值 # 特征重要性为预变化的平均绝对值 feature_importance[i] = np.mean(np.abs(baseline_pred - perturbed_pred)) # 归一化特征重要性 feature_importance = feature_importance / np.sum(feature_importance) return feature_importance except Exception as e: logger.error(f"Error calculating feature importance: {str(e)}") return None def _save_sklearn_model(self, equipment_type, model_type, model, metrics=None): """保存sklearn类型的模型""" try: timestamp = datetime.now().strftime("%Y%m%d_%H%M%S") model_dir = 'models' os.makedirs(model_dir, exist_ok=True) # 转换装备类型为英文 equipment_type_en = 'rocket' if equipment_type == '火箭炮' else 'missile' # 保存模型 model_path = f'{model_dir}/{equipment_type_en}_{model_type}_{timestamp}.joblib' from joblib import dump dump(model, model_path) # 保存标准化器 scaler_path = f'{model_dir}/{equipment_type_en}_{model_type}_{timestamp}_scaler.joblib' dump({ 'feature_scaler': self.feature_scaler, 'target_scaler': self.target_scaler }, scaler_path) # 获取评估指标 r2 = metrics['validation']['r2'] if metrics and metrics.get('validation') else metrics['train']['r2'] mae = metrics['validation']['mae'] if metrics and metrics.get('validation') else metrics['train']['mae'] rmse = metrics['validation']['rmse'] if metrics and metrics.get('validation') else metrics['train']['rmse'] # 更新数据库 with get_db_connection() as conn: cursor = conn.cursor() # 将同类型的其他模型设置为非激活 if model_type != 'pls': cursor.execute(""" UPDATE trained_models SET is_active = FALSE WHERE equipment_type = ? AND model_type != ? """, (equipment_type, 'pls')) else: cursor.execute(""" UPDATE trained_models SET is_active = FALSE WHERE equipment_type = ? AND model_type = ? """, (equipment_type, 'pls')) # 保存新模型记录 cursor.execute(""" INSERT INTO trained_models ( model_name, model_type, equipment_type, model_path, scaler_path, training_date, is_active, created_by, r2_score, mae, rmse ) VALUES (?, ?, ?, ?, ?, datetime('now','localtime'), TRUE, ?, ?, ?, ?) """, ( f"{equipment_type}_{model_type}_{timestamp}", model_type, equipment_type, model_path, scaler_path, 'system', r2, mae, rmse )) conn.commit() logger.info(f"Model {model_type} saved successfully: {model_path}") return True except Exception as e: logger.error(f"Error saving sklearn model: {str(e)}") return False