batch_experiment.py

import os
import pandas as pd
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from sklearn.preprocessing import StandardScaler
import time

# ==========================================
# 0. GPU 状态诊断
# ==========================================
def check_device():


    print("=" * 50)
    print("正在检测硬件加速器环境...")
    if torch.cuda.is_available():
        device = torch.device("cuda")
        print(f"✅ 检测到 NVIDIA GPU: {torch.cuda.get_device_name(0)}")
        print(f"✅ CUDA 版本: {torch.version.cuda}")
    elif torch.backends.mps.is_available():
        device = torch.device("mps")
        print("✅ 检测到 Apple Silicon (Mac M系列) GPU")
    else:
        device = torch.device("cpu")
        print("❌ 未检测到 GPU，将使用 CPU 运行。")
    print("=" * 50)
    return device


# ==========================================
# 1. 懒加载 Dataset 类
# ==========================================
class TimeSeriesDataset(Dataset):
    def __init__(self, data, seq_len, pred_len):
        self.data = data
        self.seq_len = seq_len
        self.pred_len = pred_len

    def __len__(self):
        return len(self.data) - self.seq_len - self.pred_len + 1

    def __getitem__(self, index):
        s_begin = index
        s_end = s_begin + self.seq_len
        r_begin = s_end
        r_end = r_begin + self.pred_len
        x = torch.tensor(self.data[s_begin:s_end], dtype=torch.float32)
        y = torch.tensor(self.data[r_begin:r_end], dtype=torch.float32)
        return x, y


def load_and_preprocess(dataset_path, seq_len, pred_len):
    df = pd.read_csv(dataset_path)

    # 缺失值插补
    numeric_cols = df.select_dtypes(include=[np.number]).columns
    if len(numeric_cols) > 0:
        df[numeric_cols] = df[numeric_cols].interpolate(method='linear', limit_direction='both').bfill().ffill()

    # 智能去除日期列
    date_col = None
    for col in df.columns:
        if 'date' in col.lower() or 'time' in col.lower():
            date_col = col
            break
    if date_col and pd.api.types.is_datetime64_any_dtype(df[date_col]):
        data_values = df.drop(columns=[date_col]).values
    else:
        data_values = df.values

    num_features = data_values.shape[1]

    # Z-Score 标准化
    scaler = StandardScaler()
    data_scaled = scaler.fit_transform(data_values)

    # 7:3 时序切分
    train_size = int(len(data_scaled) * 0.7)
    train_data = data_scaled[:train_size]
    test_data = data_scaled[train_size:]

    # 返回懒加载 Dataset
    train_dataset = TimeSeriesDataset(train_data, seq_len, pred_len)
    test_dataset = TimeSeriesDataset(test_data, seq_len, pred_len)

    return train_dataset, test_dataset, num_features


# ==========================================
# 2. 模型定义区
# ==========================================

# ---------- LSTM ----------
class MTSModel(nn.Module):
    def __init__(self, input_size, hidden_size, output_size, pred_len):
        super(MTSModel, self).__init__()
        self.pred_len = pred_len
        self.output_size = output_size
        self.lstm = nn.LSTM(input_size, hidden_size, batch_first=True, num_layers=2)
        self.fc = nn.Linear(hidden_size, pred_len * output_size)

    def forward(self, x):
        lstm_out, _ = self.lstm(x)
        last_out = lstm_out[:, -1, :]
        out = self.fc(last_out)
        out = out.view(-1, self.pred_len, self.output_size)
        return out


# ---------- DLinear ----------
class moving_avg(nn.Module):
    def __init__(self, kernel_size, stride):
        super(moving_avg, self).__init__()
        self.kernel_size = kernel_size
        self.avg = nn.AvgPool1d(kernel_size=kernel_size, stride=stride, padding=0)

    def forward(self, x):
        front = x[:, 0:1, :].repeat(1, (self.kernel_size - 1) // 2, 1)
        end = x[:, -1:, :].repeat(1, (self.kernel_size - 1) // 2, 1)
        x = torch.cat([front, x, end], dim=1)
        x = self.avg(x.permute(0, 2, 1))
        x = x.permute(0, 2, 1)
        return x

class series_decomp(nn.Module):
    def __init__(self, kernel_size):
        super(series_decomp, self).__init__()
        self.moving_avg = moving_avg(kernel_size, stride=1)

    def forward(self, x):
        moving_mean = self.moving_avg(x)
        res = x - moving_mean
        return res, moving_mean

class DLinear(nn.Module):
    def __init__(self, seq_len, pred_len, channels):
        super(DLinear, self).__init__()
        self.seq_len = seq_len
        self.pred_len = pred_len
        self.decompsition = series_decomp(25)
        self.Linear_Seasonal = nn.Linear(self.seq_len, self.pred_len)
        self.Linear_Trend = nn.Linear(self.seq_len, self.pred_len)

    def forward(self, x):
        seasonal_init, trend_init = self.decompsition(x)
        seasonal_init, trend_init = seasonal_init.permute(0, 2, 1), trend_init.permute(0, 2, 1)
        seasonal_output = self.Linear_Seasonal(seasonal_init)
        trend_output = self.Linear_Trend(trend_init)
        x = seasonal_output + trend_output
        return x.permute(0, 2, 1)


# ---------- FITS ----------
class FITS(nn.Module):
    def __init__(self, seq_len, pred_len, channels, cutoff_ratio=0.5):
        super(FITS, self).__init__()
        self.seq_len = seq_len
        self.pred_len = pred_len
        self.channels = channels
        self.freq_len = max(1, int(seq_len * cutoff_ratio))
        self.freq_real = nn.Linear(self.freq_len, pred_len)
        self.freq_imag = nn.Linear(self.freq_len, pred_len)

    def forward(self, x):
        B, L, C = x.shape
        x = x.permute(0, 2, 1)
        xf = torch.fft.rfft(x, dim=-1, norm='ortho')
        xf = xf[:, :, :self.freq_len]
        real = xf.real
        imag = xf.imag
        real_out = self.freq_real(real)
        imag_out = self.freq_imag(imag)
        pred_spec = torch.complex(real_out, imag_out)
        pred = torch.fft.irfft(pred_spec, n=self.pred_len, dim=-1, norm='ortho')
        pred = pred.permute(0, 2, 1)
        return pred


# ==========================================
# 3. 训练与评估
# ==========================================
def train_and_evaluate(model_name, train_dataset, test_dataset, num_features,
                       seq_len=96, pred_len=96, epochs=20, batch_size=32,
                       device=torch.device('cpu')):
    print(f"  模型: {model_name} | 预测长度: {pred_len} ", end='', flush=True)
    start_time = time.time()

    use_pin_memory = (device.type in ['cuda', 'mps'])
    num_workers = 0 if os.name == 'nt' else 4
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True,
                              pin_memory=use_pin_memory, num_workers=num_workers)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False,
                             pin_memory=use_pin_memory, num_workers=num_workers)

    if model_name == 'LSTM_Baseline':
        model = MTSModel(input_size=num_features, hidden_size=64,
                         output_size=num_features, pred_len=pred_len).to(device)
    elif model_name == 'DLinear':
        model = DLinear(seq_len=seq_len, pred_len=pred_len, channels=num_features).to(device)
    elif model_name == 'FITS':
        model = FITS(seq_len=seq_len, pred_len=pred_len, channels=num_features).to(device)
    else:
        raise ValueError(f"未知模型名称: {model_name}")

    criterion = nn.MSELoss()
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

    model.train()
    for epoch in range(epochs):
        for batch_x, batch_y in train_loader:
            batch_x, batch_y = batch_x.to(device, non_blocking=True), batch_y.to(device, non_blocking=True)
            optimizer.zero_grad()
            outputs = model(batch_x)
            loss = criterion(outputs, batch_y)
            loss.backward()
            optimizer.step()

    model.eval()
    test_loss_mse = 0.0
    test_loss_mae = 0.0
    criterion_mae = nn.L1Loss()
    with torch.no_grad():
        for batch_x, batch_y in test_loader:
            batch_x, batch_y = batch_x.to(device, non_blocking=True), batch_y.to(device, non_blocking=True)
            outputs = model(batch_x)
            test_loss_mse += criterion(outputs, batch_y).item()
            test_loss_mae += criterion_mae(outputs, batch_y).item()

    avg_mse = test_loss_mse / len(test_loader)
    avg_mae = test_loss_mae / len(test_loader)
    cost_time = time.time() - start_time
    print(f" | 耗时: {cost_time:.1f}s | MSE: {avg_mse:.4f} | MAE: {avg_mae:.4f}")
    return avg_mse, avg_mae


# ==========================================
# 4. 主函数
# ==========================================
def main():
    device = check_device()
    data_dir = "./data"
    dataset_files = ["Electricity.csv", "ETTh1.csv", "Exchange.csv", "Traffic.csv", "Weather.csv"]
    models_to_test = ["LSTM_Baseline", "DLinear", "FITS"]
    pred_lengths = [96, 192, 336, 720]
    seq_len = 96
    batch_size = 32
    epochs = 20

    results = []
    for file in dataset_files:
        file_path = os.path.join(data_dir, file)
        if not os.path.exists(file_path):
            print(f"警告：文件 {file_path} 不存在，已跳过。")
            continue
        dataset_name = file.split('.')[0]
        print(f"\n====== 数据集: {dataset_name} ======")

        for pl in pred_lengths:
            print(f"预处理预测长度 {pl}...")
            train_ds, test_ds, num_features = load_and_preprocess(file_path, seq_len, pl)

            for model_name in models_to_test:
                mse, mae = train_and_evaluate(
                    model_name=model_name,
                    train_dataset=train_ds,
                    test_dataset=test_ds,
                    num_features=num_features,
                    seq_len=seq_len,
                    pred_len=pl,
                    epochs=epochs,
                    batch_size=batch_size,
                    device=device
                )
                results.append({
                    'Model': model_name,
                    'Dataset': dataset_name,
                    'Pred_Len': pl,
                    'MSE': round(mse, 4),
                    'MAE': round(mae, 4)
                })

    results_df = pd.DataFrame(results)
    results_df.to_csv('benchmark_results_gpu.csv', index=False)

    print("\n" + "=" * 50)
    print("🎉 跑批任务全部完成！结果已保存至 benchmark_results_gpu.csv")
    print("=" * 50)


if __name__ == "__main__":
    main()