基于Transformer的电力负荷预测代码示例

1. 环境准备

首先，确保已安装以下库：

pip install torch torchvision pandas numpy scikit-learn matplotlib

2. 导入库

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt

from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split

3. 数据准备

3.1 加载数据

我们有一个电力负荷的数据集，包含日期时间和相应的负荷值。此处提供数据集的具体内容，数据集只有两个特征，一个为日期时间，另一个为相应日期的电力负荷值

data = pd.read_csv('load_data.csv', parse_dates=['Datetime'], index_col='Datetime')

3.2 数据预处理

# 提取时间特征
data['Hour'] = data.index.hour
data['DayOfWeek'] = data.index.dayofweek
data['Month'] = data.index.month

# 选择特征和目标
features = ['Hour', 'DayOfWeek', 'Month']
target = ['Load']

# 数据标准化
feature_scaler = StandardScaler()
target_scaler = StandardScaler()

data[features] = feature_scaler.fit_transform(data[features])
data[target] = target_scaler.fit_transform(data[target])

# 定义序列长度
sequence_length = 24  # 以过去24小时的数据预测未来负荷

# 构建输入和输出序列
def create_sequences(data, seq_length):
    xs = []
    ys = []
    for i in range(len(data) - seq_length):
        x = data[features].iloc[i:(i + seq_length)].values
        y = data[target].iloc[i + seq_length].values
        xs.append(x)
        ys.append(y)
    return np.array(xs), np.array(ys)

X, y = create_sequences(data, sequence_length)

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, shuffle=False)

# 转换为张量
X_train = torch.tensor(X_train, dtype=torch.float32)
y_train = torch.tensor(y_train, dtype=torch.float32)

X_test = torch.tensor(X_test, dtype=torch.float32)
y_test = torch.tensor(y_test, dtype=torch.float32)

3.3 创建数据集和数据加载器

class LoadDataset(Dataset):
    def __init__(self, X, y):
        self.X = X
        self.y = y
   
    def __len__(self):
        return len(self.X)
   
    def __getitem__(self, idx):
        return self.X[idx], self.y[idx]

batch_size = 64

train_dataset = LoadDataset(X_train, y_train)
test_dataset = LoadDataset(X_test, y_test)

train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

4. 定义Transformer模型（重头戏）

class TransformerModel(nn.Module):
    def __init__(self, feature_size=3, num_layers=2, dropout=0.1):
        super(TransformerModel, self).__init__()
        self.model_type = 'Transformer'
        self.feature_size = feature_size
        self.src_mask = None
        
        self.pos_encoder = PositionalEncoding(feature_size, dropout)
        encoder_layers = nn.TransformerEncoderLayer(
            d_model=feature_size, nhead=2, dropout=dropout)
        self.transformer_encoder = nn.TransformerEncoder(
            encoder_layers, num_layers=num_layers)
        self.decoder = nn.Linear(feature_size, 1)
        
    def forward(self, src):
        # src shape: [seq_len, batch_size, feature_size]
        src = src.permute(1, 0, 2)  # 调整维度顺序
        src = self.pos_encoder(src)
        output = self.transformer_encoder(src)
        output = output.mean(dim=0)  # 对序列维度求平均
        output = self.decoder(output)
        return output

class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout=0.1, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        
        pe = torch.zeros(max_len, d_model)
        position = torch.arange(
            0, max_len, dtype=torch.float).unsqueeze(1)
        div_term = torch.exp(torch.arange(
            0, d_model, 2).float() * (-np.log(10000.0) / d_model))
        pe[:, 0::2] = torch.sin(
            position * div_term)
        pe[:, 1::2] = torch.cos(
            position * div_term)
        pe = pe.unsqueeze(1)
        self.register_buffer('pe', pe)
   
    def forward(self, x):
        x = x + self.pe[:x.size(0), :]
        return self.dropout(x)

5. 定义训练和评估函数

def train(model, dataloader, criterion, optimizer):
    model.train()
    total_loss = 0
    for X_batch, y_batch in dataloader:
        optimizer.zero_grad()
        output = model(X_batch)
        loss = criterion(output, y_batch)
        loss.backward()
        optimizer.step()
        total_loss += loss.item() * X_batch.size(0)
    return total_loss / len(dataloader.dataset)

def evaluate(model, dataloader, criterion):
    model.eval()
    total_loss = 0
    with torch.no_grad():
        for X_batch, y_batch in dataloader:
            output = model(X_batch)
            loss = criterion(output, y_batch)
            total_loss += loss.item() * X_batch.size(0)
    return total_loss / len(dataloader.dataset)

6. 模型训练

model = TransformerModel(feature_size=3, num_layers=2, dropout=0.1)

# 定义损失函数和优化器
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

# 训练模型
num_epochs = 20
train_losses = []
test_losses = []

for epoch in range(1, num_epochs + 1):
    train_loss = train(model, train_loader, criterion, optimizer)
    test_loss = evaluate(model, test_loader, criterion)
    train_losses.append(train_loss)
    test_losses.append(test_loss)
    print(f'Epoch {epoch}, Training Loss: {train_loss:.4f}, Test Loss: {test_loss:.4f}')

7. 模型测试

# 模型预测
model.eval()
with torch.no_grad():
    y_pred = model(X_test).numpy()

# 反标准化
y_test_inv = target_scaler.inverse_transform(y_test.numpy())
y_pred_inv = target_scaler.inverse_transform(y_pred)

# 绘制实际值和预测值
plt.figure(figsize=(12,6))
plt.plot(y_test_inv, label='Actual Load')
plt.plot(y_pred_inv, label='Predicted Load')
plt.xlabel('Sample')
plt.ylabel('Load')
plt.legend()
plt.show()

注意事项

• 数据集：请使用实际的电力负荷数据集，替换代码中读取数据的部分。
• 参数调整：根据实际情况调整模型的参数，如特征数量、序列长度、模型层数等。
• 特征选择：可以加入更多的特征，如天气数据、节假日标识等，以提高模型性能。
• 模型优化：尝试调整学习率、批次大小、优化器等参数，或使用学习率调度策略