#!/usr/bin/env python3 import torch import torch.nn as nn import pandas as pd import numpy as np from torch.utils.data import Dataset, DataLoader import os import argparse from models.DC_PatchTST import Model import time import matplotlib.pyplot as plt from sklearn.metrics import mean_squared_error, mean_absolute_error # 检查并创建结果文件夹 def ensure_dir(path): if not os.path.exists(path): os.makedirs(path) class SineWaveDataset(Dataset): """正弦波数据集""" def __init__(self, data_path, seq_len=96, pred_len=24, mode='train'): self.seq_len = seq_len self.pred_len = pred_len self.mode = mode # 加载数据 if mode == 'train': df = pd.read_csv(os.path.join(data_path, 'train.csv')) elif mode == 'val': df = pd.read_csv(os.path.join(data_path, 'val.csv')) else: # test df = pd.read_csv(os.path.join(data_path, 'test.csv')) # 提取特征列（除timestamp外） self.data = df[['channel1', 'channel2']].values.astype(np.float32) # 计算可用样本数量 self.total_len = len(self.data) self.samples_num = max(0, self.total_len - seq_len - pred_len + 1) print(f"{mode} 数据集: {self.total_len} 条记录, {self.samples_num} 个样本") def __len__(self): return self.samples_num def __getitem__(self, idx): # 输入序列 s_begin = idx s_end = s_begin + self.seq_len # 预测目标 r_begin = s_end r_end = r_begin + self.pred_len seq_x = self.data[s_begin:s_end] # (seq_len, n_vars) seq_y = self.data[r_begin:r_end] # (pred_len, n_vars) # 时间标记（简单的位置编码） seq_x_mark = np.arange(self.seq_len).reshape(-1, 1).astype(np.float32) seq_y_mark = np.arange(self.pred_len).reshape(-1, 1).astype(np.float32) return seq_x, seq_y, seq_x_mark, seq_y_mark class Config: """配置类""" def __init__(self): # 基础配置 self.task_name = 'long_term_forecast' self.model = 'DC_PatchTST' # 数据配置 self.seq_len = 96 # 输入序列长度 self.pred_len = 24 # 预测序列长度 self.label_len = 48 # 标签长度 self.enc_in = 2 # 输入特征维度（双通道） self.dec_in = 2 # 解码器输入维度 self.c_out = 2 # 输出维度 # 模型配置 self.d_model = 128 # 模型维度 self.n_heads = 8 # 注意力头数 self.e_layers = 2 # 编码器层数 self.d_layers = 1 # 解码器层数 self.d_ff = 256 # 前向网络维度 self.factor = 1 # 注意力因子 self.dropout = 0 # Dropout率 self.activation = 'gelu' # 训练配置 self.batch_size = 1024 self.learning_rate = 0.001 self.train_epochs = 50 self.patience = 5 # 其他配置 self.use_amp = False self.num_class = 0 # GPU配置 self.use_gpu = torch.cuda.is_available() self.device = torch.device('cuda' if self.use_gpu else 'cpu') def train_epoch(model, train_loader, criterion, optimizer, device, use_amp=False): """训练一个epoch""" model.train() total_loss = 0.0 batch_count = 0 for i, (batch_x, batch_y, batch_x_mark, batch_y_mark) in enumerate(train_loader): batch_x = batch_x.to(device) batch_y = batch_y.to(device) batch_x_mark = batch_x_mark.to(device) batch_y_mark = batch_y_mark.to(device) # 构造解码器输入 dec_inp = torch.zeros_like(batch_y).to(device) optimizer.zero_grad() if use_amp: with torch.cuda.amp.autocast(): outputs, aux = model(batch_x, batch_x_mark, dec_inp, batch_y_mark) loss = criterion(outputs, batch_y) # 添加DC的ratio loss if aux is not None and 'ratio_loss0' in aux and 'ratio_loss1' in aux: ratio_loss = aux['ratio_loss0'] + aux['ratio_loss1'] loss = loss + 0.0 * ratio_loss # ratio loss权重 else: outputs, aux = model(batch_x, batch_x_mark, dec_inp, batch_y_mark) loss = criterion(outputs, batch_y) # 添加DC的ratio loss if aux is not None and 'ratio_loss0' in aux and 'ratio_loss1' in aux: ratio_loss = aux['ratio_loss0'] + aux['ratio_loss1'] loss = loss + 0.0 * ratio_loss # ratio loss权重 loss.backward() optimizer.step() total_loss += loss.item() batch_count += 1 if i % 100 == 0: print(f'Batch {i}, Loss: {loss.item():.6f}') return total_loss / batch_count def validate(model, val_loader, criterion, device): """验证模型""" model.eval() total_loss = 0.0 batch_count = 0 with torch.no_grad(): for batch_x, batch_y, batch_x_mark, batch_y_mark in val_loader: batch_x = batch_x.to(device) batch_y = batch_y.to(device) batch_x_mark = batch_x_mark.to(device) batch_y_mark = batch_y_mark.to(device) dec_inp = torch.zeros_like(batch_y).to(device) outputs, _ = model(batch_x, batch_x_mark, dec_inp, batch_y_mark) loss = criterion(outputs, batch_y) total_loss += loss.item() batch_count += 1 return total_loss / batch_count def test_model(model, test_loader, device, save_path): """测试模型并可视化结果""" model.eval() predictions = [] ground_truths = [] with torch.no_grad(): for batch_x, batch_y, batch_x_mark, batch_y_mark in test_loader: batch_x = batch_x.to(device) batch_y = batch_y.to(device) batch_x_mark = batch_x_mark.to(device) batch_y_mark = batch_y_mark.to(device) dec_inp = torch.zeros_like(batch_y).to(device) outputs, _ = model(batch_x, batch_x_mark, dec_inp, batch_y_mark) predictions.append(outputs.cpu().numpy()) ground_truths.append(batch_y.cpu().numpy()) predictions = np.concatenate(predictions, axis=0) ground_truths = np.concatenate(ground_truths, axis=0) # 计算指标 mse = mean_squared_error(ground_truths.reshape(-1), predictions.reshape(-1)) mae = mean_absolute_error(ground_truths.reshape(-1), predictions.reshape(-1)) print(f"测试结果 - MSE: {mse:.6f}, MAE: {mae:.6f}") # 可视化前几个样本 plt.figure(figsize=(15, 10)) for i in range(min(4, len(predictions))): for ch in range(2): # 双通道 plt.subplot(4, 2, i*2 + ch + 1) plt.plot(ground_truths[i, :, ch], label='Ground Truth', color='blue') plt.plot(predictions[i, :, ch], label='Prediction', color='red') plt.title(f'Sample {i+1}, Channel {ch+1}') plt.legend() plt.tight_layout() plt.savefig(os.path.join(save_path, 'predictions.png')) print(f"预测结果可视化保存到: {save_path}/predictions.png") return mse, mae def main(): # 配置参数 config = Config() # 创建结果目录 results_dir = './results/dc_patchtst_sine_wave' ensure_dir(results_dir) print(f"使用设备: {config.device}") print(f"模型配置: seq_len={config.seq_len}, pred_len={config.pred_len}") # 加载数据 data_path = './data/sine_wave/' train_dataset = SineWaveDataset(data_path, config.seq_len, config.pred_len, 'train') val_dataset = SineWaveDataset(data_path, config.seq_len, config.pred_len, 'val') test_dataset = SineWaveDataset(data_path, config.seq_len, config.pred_len, 'test') train_loader = DataLoader(train_dataset, batch_size=config.batch_size, shuffle=True, num_workers=4) val_loader = DataLoader(val_dataset, batch_size=config.batch_size, shuffle=False, num_workers=4) test_loader = DataLoader(test_dataset, batch_size=config.batch_size, shuffle=False, num_workers=4) # 创建模型 model = Model(config).to(config.device) print(f"模型参数数量: {sum(p.numel() for p in model.parameters())}") # 优化器和损失函数 optimizer = torch.optim.Adam(model.parameters(), lr=config.learning_rate, weight_decay=1e-4) criterion = nn.MSELoss() # 训练循环 best_val_loss = float('inf') patience_counter = 0 train_losses = [] val_losses = [] print("开始训练...") start_time = time.time() for epoch in range(config.train_epochs): epoch_start = time.time() # 训练 train_loss = train_epoch(model, train_loader, criterion, optimizer, config.device, config.use_amp) # 验证 val_loss = validate(model, val_loader, criterion, config.device) train_losses.append(train_loss) val_losses.append(val_loss) epoch_time = time.time() - epoch_start print(f'Epoch {epoch+1:2d}/{config.train_epochs} | ' f'Train Loss: {train_loss:.6f} | ' f'Val Loss: {val_loss:.6f} | ' f'Time: {epoch_time:.2f}s') # 早停检查 if val_loss < best_val_loss: best_val_loss = val_loss patience_counter = 0 # 保存最佳模型 torch.save({ 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'config': config, 'epoch': epoch, 'val_loss': val_loss }, os.path.join(results_dir, 'best_model.pth')) print(f' -> 保存最佳模型 (val_loss: {val_loss:.6f})') else: patience_counter += 1 if patience_counter >= config.patience: print(f'早停触发! 最佳验证损失: {best_val_loss:.6f}') break total_time = time.time() - start_time print(f'\n训练完成! 总时间: {total_time/60:.2f} 分钟') # 加载最佳模型进行测试 checkpoint = torch.load(os.path.join(results_dir, 'best_model.pth')) model.load_state_dict(checkpoint['model_state_dict']) print("\n测试最佳模型...") test_mse, test_mae = test_model(model, test_loader, config.device, results_dir) # 保存训练历史 history = { 'train_losses': train_losses, 'val_losses': val_losses, 'test_mse': test_mse, 'test_mae': test_mae } # 绘制训练曲线 plt.figure(figsize=(12, 4)) plt.subplot(1, 2, 1) plt.plot(train_losses, label='Train Loss') plt.plot(val_losses, label='Validation Loss') plt.xlabel('Epoch') plt.ylabel('Loss') plt.legend() plt.title('Training History') plt.subplot(1, 2, 2) plt.bar(['MSE', 'MAE'], [test_mse, test_mae]) plt.title('Test Metrics') plt.ylabel('Error') plt.tight_layout() plt.savefig(os.path.join(results_dir, 'training_history.png')) print(f"\n结果保存在: {results_dir}") print(f"最佳模型: {results_dir}/best_model.pth") print(f"训练历史: {results_dir}/training_history.png") print(f"预测可视化: {results_dir}/predictions.png") if __name__ == "__main__": main()