# -*- coding: utf-8 -*- """ CNN-BiLSTM-Attention 模型训练器 """ import torch import torch.optim as optim import numpy as np import time import copy from models.model_registry import register_trainer from utils.model_manager import model_manager from analysis.metrics import evaluate_model from utils.data_utils import create_dataset from sklearn.preprocessing import MinMaxScaler from utils.visualization import plot_loss_curve # 导入绘图函数 # 导入新创建的模型 from models.cnn_bilstm_attention import CnnBiLstmAttention def train_with_cnn_bilstm_attention(product_id, model_identifier, product_df, store_id, training_mode, aggregation_method, epochs, sequence_length, forecast_horizon, model_dir, **kwargs): """ 使用 CNN-BiLSTM-Attention 模型进行训练，并实现早停和最佳模型保存。 """ print(f"🚀 CNN-BiLSTM-Attention 训练器启动: model_identifier='{model_identifier}'") start_time = time.time() # --- 1. 数据准备 --- features = ['sales', 'weekday', 'month', 'is_holiday', 'is_weekend', 'is_promotion', 'temperature'] product_name = product_df['product_name'].iloc[0] if 'product_name' in product_df.columns else model_identifier X = product_df[features].values y = product_df[['sales']].values scaler_X = MinMaxScaler(feature_range=(0, 1)) scaler_y = MinMaxScaler(feature_range=(0, 1)) X_scaled = scaler_X.fit_transform(X) y_scaled = scaler_y.fit_transform(y) train_size = int(len(X_scaled) * 0.8) X_train_raw, X_test_raw = X_scaled[:train_size], X_scaled[train_size:] y_train_raw, y_test_raw = y_scaled[:train_size], y_scaled[train_size:] trainX, trainY = create_dataset(X_train_raw, y_train_raw, sequence_length, forecast_horizon) testX, testY = create_dataset(X_test_raw, y_test_raw, sequence_length, forecast_horizon) trainX = torch.from_numpy(trainX).float() trainY = torch.from_numpy(trainY).float() testX = torch.from_numpy(testX).float() testY = torch.from_numpy(testY).float() # --- 2. 实例化模型和优化器 --- input_dim = trainX.shape[2] model = CnnBiLstmAttention( input_dim=input_dim, output_dim=forecast_horizon, sequence_length=sequence_length ) optimizer = optim.Adam(model.parameters(), lr=kwargs.get('learning_rate', 0.001)) criterion = torch.nn.MSELoss() # --- 3. 训练循环与早停 --- print("开始训练 CNN-BiLSTM-Attention 模型 (含早停)...") # 版本锁定：在训练开始前确定本次训练的版本号 current_version = model_manager.peek_next_version( model_type='cnn_bilstm_attention', product_id=product_id, store_id=store_id, training_mode=training_mode, aggregation_method=aggregation_method ) print(f"🔒 本次训练版本锁定为: {current_version}") loss_history = {'train': [], 'val': []} best_val_loss = float('inf') best_model_state = None best_model_path = None # 用于存储最佳模型的路径 patience = kwargs.get('patience', 15) patience_counter = 0 for epoch in range(epochs): model.train() optimizer.zero_grad() outputs = model(trainX) train_loss = criterion(outputs, trainY.squeeze(-1)) train_loss.backward() optimizer.step() # 验证 model.eval() with torch.no_grad(): val_outputs = model(testX) val_loss = criterion(val_outputs, testY.squeeze(-1)) loss_history['train'].append(train_loss.item()) loss_history['val'].append(val_loss.item()) if (epoch + 1) % 10 == 0: print(f'Epoch [{epoch+1}/{epochs}], Train Loss: {train_loss.item():.4f}, Val Loss: {val_loss.item():.4f}') # 早停逻辑 if val_loss.item() < best_val_loss: best_val_loss = val_loss.item() best_model_state = copy.deepcopy(model.state_dict()) patience_counter = 0 print(f"✨ 新的最佳模型! Epoch: {epoch+1}, Val Loss: {best_val_loss:.4f}") # 立即保存最佳模型 best_model_data = { 'model_state_dict': best_model_state, 'scaler_X': scaler_X, 'scaler_y': scaler_y, 'config': { 'model_type': 'cnn_bilstm_attention', 'input_dim': input_dim, 'output_dim': forecast_horizon, 'sequence_length': sequence_length, 'features': features }, 'epoch': epoch + 1 } best_model_path, _ = model_manager.save_model( model_data=best_model_data, product_id=product_id, model_type='cnn_bilstm_attention', store_id=store_id, training_mode=training_mode, aggregation_method=aggregation_method, product_name=product_name, version=f"{current_version}_best" ) else: patience_counter += 1 if patience_counter >= patience: print(f"🚫 早停触发! 在 epoch {epoch+1} 停止。") break training_time = time.time() - start_time print(f"模型训练完成，耗时: {training_time:.2f}秒") # --- 4. 使用最佳模型进行评估 --- if best_model_state: model.load_state_dict(best_model_state) print("最佳模型已加载用于最终评估。") model.eval() with torch.no_grad(): test_pred_scaled = model(testX) test_pred_unscaled = scaler_y.inverse_transform(test_pred_scaled.numpy()) test_true_unscaled = scaler_y.inverse_transform(testY.squeeze(-1).numpy()) metrics = evaluate_model(test_true_unscaled.flatten(), test_pred_unscaled.flatten()) metrics['training_time'] = training_time metrics['best_val_loss'] = best_val_loss metrics['stopped_epoch'] = epoch + 1 print("\n最佳模型评估指标:") print(f"MSE: {metrics['mse']:.4f}, RMSE: {metrics['rmse']:.4f}, MAE: {metrics['mae']:.4f}, R²: {metrics['r2']:.4f}, MAPE: {metrics['mape']:.2f}%") # --- 5. 保存工件 --- # 准备 scope 和 identifier 以生成标准化的文件名 scope = training_mode if scope == 'product': identifier = product_id elif scope == 'store': identifier = store_id elif scope == 'global': identifier = aggregation_method else: identifier = product_name # 后备方案 # 绘制带有版本号的损失曲线图 loss_curve_path = plot_loss_curve( train_losses=loss_history['train'], val_losses=loss_history['val'], model_type='cnn_bilstm_attention', scope=scope, identifier=identifier, version=current_version, # 使用锁定的版本 model_dir=model_dir ) print(f"📈 带版本号的损失曲线已保存: {loss_curve_path}") # 准备要保存的最终模型数据 model_data = { 'model_state_dict': best_model_state, # 保存最佳模型的状态 'scaler_X': scaler_X, 'scaler_y': scaler_y, 'config': { 'model_type': 'cnn_bilstm_attention', 'input_dim': input_dim, 'output_dim': forecast_horizon, 'sequence_length': sequence_length, 'features': features }, 'metrics': metrics, 'loss_history': loss_history, 'loss_curve_path': loss_curve_path # 直接包含路径 } # 使用模型管理器保存最终模型 final_model_path, final_version = model_manager.save_model( model_data=model_data, product_id=product_id, model_type='cnn_bilstm_attention', store_id=store_id, training_mode=training_mode, aggregation_method=aggregation_method, product_name=product_name, version=current_version # 使用锁定的版本号 ) print(f"✅ CNN-BiLSTM-Attention 最终模型已保存，版本: {final_version}") # 组装返回的工件 artifacts = { "versioned_model": final_model_path, "loss_curve_plot": loss_curve_path, "best_model": best_model_path, "version": final_version } return model, metrics, artifacts # --- 关键步骤: 将训练器注册到系统中 --- register_trainer('cnn_bilstm_attention', train_with_cnn_bilstm_attention)