pytorch股票预测

#导包
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=5

data_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_size
        self.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 out

model=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()