LSTM实战之预测股票

news2024/10/2 10:32:30

📈 用PyTorch搭建LSTM模型，轻松预测股票价格！🚀

Hey小伙伴们，今天给大家带来一个超级实用的项目教程——如何用PyTorch和LSTM模型来预测股票价格！🌟

🔍 项目背景

我们都知道股市是个风云变幻的地方，而预测股价则是很多投资者梦寐以求的能力。今天，我们就来尝试一下用机器学习的方法来预测股价，让数据说话！

📑 准备工作

首先，我们要准备好开发环境，确保安装了以下Python库：

numpy: 数组处理
pandas: 数据处理
matplotlib: 数据可视化
scikit-learn: 数据预处理
torch: 构建LSTM模型

💻 实战演练

1️⃣ 导入库 & 加载数据

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.preprocessing import MinMaxScaler
import torch
import torch.nn as nn
import torch.optim as optim

# 加载数据
df = pd.read_csv('stock_data.csv')
# 只保留收盘价
data = df.filter(['Close'])
# 将数据转换为numpy数组
dataset = data.values
# 归一化数据
scaler = MinMaxScaler(feature_range=(0, 1))
scaled_data = scaler.fit_transform(dataset)

2️⃣ 创建数据集

# 训练集和测试集划分
training_data_len = int(np.ceil(len(dataset) * .8))

# 创建训练数据集
def create_dataset(data, time_step=1):
    X_train, y_train = [], []
    for i in range(len(data)-time_step-1):
        X_train.append(data[i:(i+time_step), 0])
        y_train.append(data[i + time_step, 0])
    return np.array(X_train), np.array(y_train)

time_step = 60
X_train, y_train = create_dataset(scaled_data[:training_data_len], time_step)
X_test, y_test = create_dataset(scaled_data[training_data_len-time_step:], time_step)

# 调整数据形状以适应LSTM
X_train = np.reshape(X_train, (X_train.shape[0], X_train.shape[1], 1))
X_test = np.reshape(X_test, (X_test.shape[0], X_test.shape[1], 1))

# 转换为PyTorch张量
X_train = torch.from_numpy(X_train).float()
y_train = torch.from_numpy(y_train).float()
X_test = torch.from_numpy(X_test).float()
y_test = torch.from_numpy(y_test).float()

3️⃣ 构建LSTM模型

class LSTMModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, layer_dim, output_dim):
        super(LSTMModel, self).__init__()
        self.hidden_dim = hidden_dim
        self.layer_dim = layer_dim
        self.lstm = nn.LSTM(input_dim, hidden_dim, layer_dim, batch_first=True)
        self.fc = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        h0 = torch.zeros(self.layer_dim, x.size(0), self.hidden_dim).requires_grad_()
        c0 = torch.zeros(self.layer_dim, x.size(0), self.hidden_dim).requires_grad_()
        
        out, (hn, cn) = self.lstm(x, (h0.detach(), c0.detach()))
        out = self.fc(out[:, -1, :]) 
        return out

input_dim = 1
hidden_dim = 50
layer_dim = 1
output_dim = 1

model = LSTMModel(input_dim, hidden_dim, layer_dim, output_dim)

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

4️⃣ 训练模型

num_epochs = 100

for epoch in range(num_epochs):
    outputs = model(X_train)
    optimizer.zero_grad()
    
    # 获取损失
    loss = criterion(outputs, y_train)
    
    # 反向传播和优化
    loss.backward()
    optimizer.step()
    
    if (epoch+1) % 10 == 0:
        print(f"Epoch: {epoch+1}, Loss: {loss.item()}")

5️⃣ 预测和评估

# 获取模型预测值
train_predictions = model(X_train).detach().numpy()
test_predictions = model(X_test).detach().numpy()

# 反归一化预测值
train_predictions = scaler.inverse_transform(train_predictions)
test_predictions = scaler.inverse_transform(test_predictions)

# 计算均方根误差（RMSE）
rmse = np.sqrt(np.mean(((test_predictions - y_test.numpy()) ** 2)))
print("Root Mean Squared Error:", rmse)

# 可视化结果
train = data[:training_data_len+1]
valid = data[training_data_len+1:]
valid['Predictions'] = test_predictions

plt.figure(figsize=(16,8))
plt.title('Model')
plt.xlabel('Date', fontsize=18)
plt.ylabel('Close Price', fontsize=18)
plt.plot(train['Close'])
plt.plot(valid[['Close', 'Predictions']])
plt.legend(['Train', 'Val', 'Predictions'], loc='upper right')
plt.show()