#导包
import numpy as np
import torch
from torch.autograd import Variable
from sklearn.preprocessing import MinMaxScaler
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt#声明随机种子
torch.manual_seed(777)#读取数据
data=np.loadtxt('../data-02-stock_daily.csv',delimiter=',')#顺序颠倒
data=data[::-1]#归一化
min=MinMaxScaler()
data=min.fit_transform(data)x=data
y=data[:,-1:]# 7天为一组
length=7
num_class=5data_x=[]
data_y=[]for i in range(0,len(y)-length):x1=x[i:i+length]y1=y[i+length]data_x.append(x1)data_y.append(y1)x_train,x_test,y_train,y_test=train_test_split(data_x,data_y,test_size=0.3,shuffle=False)x_train=Variable(torch.Tensor(x_train).float())
y_train=Variable(torch.Tensor(y_train).float())
x_test=torch.Tensor(x_test).float()
y_test=torch.Tensor(y_test).float()class LSTM(torch.nn.Module):def __init__(self, input_size):super(LSTM, self).__init__()self.h_size=input_sizeself.lstm=torch.nn.LSTM(input_size=self.h_size,hidden_size=self.h_size,num_layers=2,batch_first=True)self.fc=torch.nn.Linear(self.h_size,1)def forward(self,x):_,(h,_)=self.lstm(x)# 取lstm最后一层的最后一个时间步的输出，并展开成二维h=h[-1].view(-1,self.h_size)out=self.fc(h)return outmodel=LSTM(num_class)loss=torch.nn.MSELoss() # 股票预测是回归算法
optim=torch.optim.Adam(model.parameters(),lr=0.01)for ei in range(200):optim.zero_grad()h=model(x_train)cost=loss(h,y_train)cost.backward()optim.step()if ei % 50 == 0:print(ei,'损失值:',cost.data.numpy())#真实值与预测值画图
h_test=model(x_test)
plt.plot(h_test.data.numpy(),c='r')
plt.plot(y_test,c='b')
plt.show()