用循环神经网络预测股价

创作时间:

作者:

@小白创作中心

用循环神经网络预测股价

引用

CSDN

https://blog.csdn.net/weixin_42255757/article/details/139242867

循环神经网络（RNN）和长短期记忆网络（LSTM）是处理时间序列数据的强大工具。本文将详细介绍如何使用这两种模型预测股票价格，通过对比RNN和LSTM的预测效果，展示LSTM在处理时间序列数据时的优势。

数据获取与预处理

首先，我们需要获取股票数据。这里使用akshare接口来获取数据，选择浙大网新（股票代码：600797）从2021年1月1日到当前的股票数据。

df_hs300 = ak.stock_zh_a_hist(symbol="600797", period="daily", start_date="20210101", end_date=datetime.datetime.today().strftime("%Y%m%d"), adjust="")

获取数据后，我们只需要收盘价和日期两个字段，并对收盘价进行归一化处理：

close_list = df_hs300['收盘'].values
date_list = df_hs300['日期'].values
close_list_norm=[price/max(close_list) for price in close_list]

可视化处理后的数据：

%matplotlib inline
import matplotlib.pyplot as plt
plt.plot(close_list_norm)
plt.title('hs_300')
plt.xlabel('date')
plt.ylabel('colse price')
plt.show()

数据集定义

接下来，定义一个Dataset和DataLoader，选择前10天的收盘价来预测下一个交易日的收盘价，使用前700个数据作为训练数据集：

from torch.utils.data import Dataset, DataLoader

class StockDataset(Dataset):
    def __init__(self, data_list, time_step = 10, transform=None):
        self.data = data_list
        self.features = []
        self.targets = []
        for i in range(len(self.data)-time_step):
            feature = [x for x in self.data[i:i+time_step]]
            y = self.data[i+time_step]
            self.features.append(feature)
            self.targets.append(y)
        self.features = torch.tensor(self.features)
        self.features = self.features.reshape(-1, time_step, 1)
        
    def __len__(self):
        return len(self.features)
    
    def __getitem__(self, idx):
        return self.features[idx], self.targets[idx]

transform = transforms.Compose([transforms.ToTensor()])
dataset = StockDataset(close_list_norm[:700], time_step=10, transform=transform)
dataloader = torch.utils.data.DataLoader(dataset=dataset, batch_size=1, shuffle=True)

RNN模型定义与训练

定义一个简单的RNN模型：

time_step = 10
batch_size = 1
input_size,hidden_size,output_size=1,20,1

class Net(nn.Module):
    def __init__(self):
        super(Net,self).__init__()
        self.rnn=nn.RNN(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=1,
            batch_first=True
        )
        self.linear=nn.Linear(hidden_size,output_size)
    def forward(self,X,hidden_prev):
        out,ht=self.rnn(X,hidden_prev)
        batch_size, seq,  hidden_size = out.shape
        out = self.linear(out[:, -1, :])
        return out

训练模型：

model=Net()
model=model.cuda()
criterion=nn.MSELoss()
learning_rate,epochs=0.01,500
optimizer=torch.optim.Adam(model.parameters(),lr=learning_rate)
for epoch in range(epochs):
    losses = []
    for X,y in dataloader:
        X = X.cuda()
        y = y.cuda()
        y=y.to(torch.float32)
        X=X.to(torch.float32)
        optimizer.zero_grad()
        yy=model(X,hidden_prev)
        yy=yy.cuda()
        loss = criterion(y, yy)
        model.zero_grad()
        loss.backward()
        optimizer.step()
        losses.append(loss.item())
    epoch_loss=sum(losses)/len(losses)
    if epoch%50==0:
        print("epoch:{},loss:{:.8f}".format(epoch+1,epoch_loss))
torch.save(model, "model2.pt")

验证RNN模型

定义验证数据集：

transform = transforms.Compose([transforms.ToTensor()])
val_dataset = StockDataset(close_list_norm[700:], time_step=10, transform=transform)
val_dataloader = torch.utils.data.DataLoader(dataset=val_dataset, batch_size=1, shuffle=False)

验证模型：

model = torch.load('model2.pt')
Val_y,Val_predict=[],[]
Val_max_price=max(close_list) 
for X,y in val_dataloader:
    with torch.no_grad():
        X = X.cuda()
        y=y.to(torch.float32)
        X=X.to(torch.float32)
        predict=model(X,hidden_prev)
        y=y.cpu()
        predict=predict.cpu()
        Val_y.append(y[0][0]*Val_max_price) 
        Val_predict.append(predict[0][0]*Val_max_price)
fig=plt.figure(figsize=(8,5),dpi=80)
plt.plot(Val_y,linestyle='--',color='r')
plt.plot(Val_predict,color='g')
plt.title('stock price')
plt.xlabel('time')
plt.ylabel('price')
plt.show()

LSTM模型定义与训练

定义LSTM模型：

class Net_LSTM(nn.Module):
    def __init__(self):
        super(Net_LSTM,self).__init__()
        self.lstm=nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=1,
            batch_first=True
        )
        self.linear=nn.Linear(hidden_size,output_size)
    def forward(self,X):
        out,ht=self.lstm(X)      
        batch_size, seq,  hidden_size = out.shape
        out = self.linear(out[:, -1, :])
        return out

由于训练和验证的函数与RNN基本一致，这里不再赘述。下面是LSTM模型在验证集上的表现：