import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from sklearn.metrics import mean_squared_error
import matplotlib.pyplot as plt
# 数据生成函数
def generate_synthetic_time_series(num_samples, seq_length):
x = np.linspace(0, 4 * np.pi, num_samples)
y = np.sin(x) + 0.1 * np.random.randn(num_samples)
X, y_seq = [], []
for i in range(num_samples - seq_length):
X.append(y[i:i + seq_length])
y_seq.append(y[i + seq_length])
return np.array(X), np.array(y_seq)
# 模型定义
class Attention(nn.Module):
def __init__(self, hidden_dim):
super(Attention, self).__init__()
self.Wa = nn.Linear(hidden_dim, hidden_dim)
self.Ua = nn.Linear(hidden_dim, hidden_dim)
self.Va = nn.Linear(hidden_dim, 1)
def forward(self, hidden, encoder_outputs):
scores = self.Va(torch.tanh(self.Wa(hidden) + self.Ua(encoder_outputs)))
weights = torch.softmax(scores, dim=1)
context = torch.bmm(weights, encoder_outputs)
return context, weights
class LSTMWithAttention(nn.Module):
def __init__(self, input_dim, hidden_dim, output_dim):
super(LSTMWithAttention, self).__init__()
self.lstm = nn.LSTM(input_dim, hidden_dim, batch_first=True)
self.attention = Attention(hidden_dim)
self.fc = nn.Linear(hidden_dim, output_dim)
def forward(self, x):
lstm_out, (hn, cn) = self.lstm(x)
context, attn_weights = self.attention(hn[-1], lstm_out)
out = self.fc(context.squeeze(1))
return out, attn_weights
# 训练模型
def train_model(model, criterion, optimizer, X_train, y_train, num_epochs=10, batch_size=64):
model.train()
dataset = TensorDataset(torch.Tensor(X_train), torch.Tensor(y_train))
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)
for epoch in range(num_epochs):
epoch_loss = 0
for inputs, targets in dataloader:
optimizer.zero_grad()
outputs, _ = model(inputs)
loss = criterion(outputs.squeeze(), targets)
loss.backward()
optimizer.step()
epoch_loss += loss.item()
print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {epoch_loss / len(dataloader)}')
# 预测
def predict(model, X):
model.eval()
with torch.no_grad():
inputs = torch.Tensor(X)
predictions, _ = model(inputs)
return predictions.numpy()
# 评价指标
def evaluate(y_true, y_pred):
mse = mean_squared_error(y_true, y_pred)
rmse = np.sqrt(mse)
return mse, rmse
# 主程序
if __name__ == "__main__":
# 生成数据
num_samples = 1000
seq_length = 10
X, y = generate_synthetic_time_series(num_samples, seq_length)
# 划分数据集
split_index = int(0.8 * len(X))
X_train, X_test = X[:split_index], X[split_index:]
y_train, y_test = y[:split_index], y[split_index:]
# 定义模型、损失函数和优化器
input_dim = 1
hidden_dim = 64
output_dim = 1
model = LSTMWithAttention(input_dim, hidden_dim, output_dim)
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)
# 训练模型
train_model(model, criterion, optimizer, X_train, y_train)
# 预测
y_pred = predict(model, X_test)
# 评价
mse, rmse = evaluate(y_test, y_pred)
print(f'MSE: {mse}, RMSE: {rmse}')
# 可视化
plt.figure(figsize=(12, 6))
plt.plot(np.arange(seq_length, len(y_test) + seq_length), y_test, label='True Signal')
plt.plot(np.arange(seq_length, len(y_test) + seq_length), y_pred, label='Predicted Signal', linestyle='--')
plt.legend()
plt.show()
在 AI 技术的快速发展下
最新推荐文章于 2024-08-06 16:32:17 发布