720 lines
32 KiB
Python
720 lines
32 KiB
Python
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 |