在Pytorch中使用lstm预测amzon股票价格

使用 WSL

安装 python 扩展

安装依赖

Copy
sudo apt update
sudo apt upgrade
sudo apt install python3-pip
sudo apt install python3-pandas
pip3 install torch
pip3 install numpy
pip3 install scikit-learn

导入

Copy
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import torch
import torch.nn as nn
from sklearn.preprocessing import MinMaxScaler
from torch.utils.data import Dataset

加载数据

Copy
data = pd.read_csv('AMZN.csv')

取出其中的日期和收盘价

Copy
data = data[['Date', 'Close']]

选择我们使用的设备是 cpu 还是 gpu

Copy
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'

准备一个数据框，我们用前七天的收盘价预测后一天的收盘价

Copy
def prepare_dataframe_for_lstm(df, n_steps):
df = dc(df)
df.set_index('Date', inplace=True)
for i in range(1, n_steps+1):
df[f'Close(t-{i})'] = df['Close'].shift(i)
df.dropna(inplace=True)
return df
lookback = 7
shifted_df = prepare_dataframe_for_lstm(data, lookback)

转换成 numpy

Copy
shifted_df_as_np = shifted_df.to_numpy()

缩放数据 - 1 到 1 之间

Copy
scaler = MinMaxScaler(feature_range=(-1, 1))
shifted_df_as_np = scaler.fit_transform(shifted_df_as_np)

处理数据，x 是输入值，矩阵的第一列， y 是输出值，矩阵的后面 7 列。

Copy
X = shifted_df_as_np[:, 1:]
y = shifted_df_as_np[:, 0]

翻转 x 在水平轴上的镜像

Copy
X = dc(np.flip(X, axis=1))

分割，其中的百分之 95 用来训练，百分之 5 用来测试

Copy
X = dc(np.flip(X, axis=1))
split_index = int(len(X) * 0.95)
X_train = X[:split_index]
X_test = X[split_index:]
y_train = y[:split_index]
y_test = y[split_index:]

重塑矩阵得到维度

Copy
X_train = X_train.reshape((-1, lookback, 1))
X_test = X_test.reshape((-1, lookback, 1))
y_train = y_train.reshape((-1, 1))
y_test = y_test.reshape((-1, 1))

将 numpy 中的所有数据转变成 tensor

Copy
X_train = torch.tensor(X_train).float()
y_train = torch.tensor(y_train).float()
X_test = torch.tensor(X_test).float()
y_test = torch.tensor(y_test).float()

创建数据集

Copy
class TimeSeriesDataset(Dataset):
def __init__(self, X, y):
self.X = X
self.y = y
def __len__(self):
return len(self.X)
def __getitem__(self, i):
return self.X[i], self.y[i]
train_dataset = TimeSeriesDataset(X_train, y_train)
test_dataset = TimeSeriesDataset(X_test, y_test)

加载数据集

Copy
from torch.utils.data import DataLoader
batch_size = 16
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

搭建 lstm 模型

Copy
class LSTM(nn.Module):
def __init__(self, input_size, hidden_size, num_stacked_layers):
super().__init__()
self.hidden_size = hidden_size
self.num_stacked_layers = num_stacked_layers

self.lstm = nn.LSTM(input_size, hidden_size, num_stacked_layers,
batch_first=True)

self.fc = nn.Linear(hidden_size, 1)

def forward(self, x):
batch_size = x.size(0)
h0 = torch.zeros(self.num_stacked_layers, batch_size, self.hidden_size).to(device)
c0 = torch.zeros(self.num_stacked_layers, batch_size, self.hidden_size).to(device)

out, _ = self.lstm(x, (h0, c0))
out = self.fc(out[:, -1, :])
return out

model = LSTM(1, 4, 1)
model.to(device)

定义参数

Copy
learning_rate = 0.001
num_epochs = 10
loss_function = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

创建训练函数

Copy
def train_one_epoch():
model.train(True)
print(f'Epoch: {epoch + 1}')
running_loss = 0.0

for batch_index, batch in enumerate(train_loader):
x_batch, y_batch = batch[0].to(device), batch[1].to(device)

output = model(x_batch)
loss = loss_function(output, y_batch)
running_loss += loss.item()

optimizer.zero_grad()
loss.backward()
optimizer.step()

if batch_index % 100 == 99: # print every 100 batches
avg_loss_across_batches = running_loss / 100
print('Batch {0}, Loss: {1:.3f}'.format(batch_index+1,
avg_loss_across_batches))
running_loss = 0.0
print()

创建训练函数

Copy
def train_one_epoch():
model.train(True)
print(f'Epoch: {epoch + 1}')
running_loss = 0.0

for batch_index, batch in enumerate(train_loader):
x_batch, y_batch = batch[0].to(device), batch[1].to(device)

output = model(x_batch)
loss = loss_function(output, y_batch)
running_loss += loss.item()

optimizer.zero_grad()
loss.backward()
optimizer.step()

if batch_index % 100 == 99: # print every 100 batches
avg_loss_across_batches = running_loss / 100
print('Batch {0}, Loss: {1:.3f}'.format(batch_index+1,
avg_loss_across_batches))
running_loss = 0.0
print()

创建测试函数

Copy
def validate_one_epoch():
model.train(False)
running_loss = 0.0

for batch_index, batch in enumerate(test_loader):
x_batch, y_batch = batch[0].to(device), batch[1].to(device)

with torch.no_grad():
output = model(x_batch)
loss = loss_function(output, y_batch)
running_loss += loss.item()

avg_loss_across_batches = running_loss / len(test_loader)

print('Val Loss: {0:.3f}'.format(avg_loss_across_batches))
print('***************************************************')
print()

循环

Copy
for epoch in range(num_epochs):
train_one_epoch()
validate_one_epoch()

可视化

Copy
with torch.no_grad():
predicted = model(X_train.to(device)).to('cpu').numpy()
plt.plot(y_train, label='Actual Close')
plt.plot(predicted, label='Predicted Close')
plt.xlabel('Day')
plt.ylabel('Close')
plt.legend()
plt.show()