335 lines
12 KiB
Python
335 lines
12 KiB
Python
#!/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() |