1.查看数据
2. 结果
3.代码
#!/usr/bin/env python
# coding: utf-8
# In[1]:
get_ipython().run_line_magic('config', 'Completer.use_jedi = False')
# In[2]:
import torch
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import torch.nn as nn
from torch.utils.data import Dataset,DataLoader
from sklearn.preprocessing import MinMaxScaler
# In[3]:
filePath = './000001SH_index.csv'
seq_len = 8
batch_size = 64
input_size = 1
hidden_size = 8
num_layers = 1
# In[4]:
sc = MinMaxScaler(feature_range=(-1,1))
# In[5]:
raw_data = pd.read_csv(filePath)
data = raw_data.loc[:,['close']]
data = pd.DataFrame(sc.fit_transform(data))
X = []
Y = []
for i in range(data.shape[0]-seq_len):
X.append(np.array(data.iloc[i:i+seq_len,:].values,dtype=np.float32))
Y.append(np.array(data.iloc[i+seq_len,0],dtype=np.float32))
x_train,x_test = X[:int(len(Y)*0.9)],X[int(len(Y)*0.9):]
y_train,y_test = Y[:int(len(Y)*0.9)],Y[int(len(Y)*0.9):]
print(np.array(x_test).shape)
print(np.array(x_train).shape)
# In[6]:
raw_data.close[:int(len(Y)*0.9)].plot(figsize=(16,4))
raw_data.close[int(len(Y)*0.9):].plot(figsize=(16,4))
plt.show()
# In[7]:
class MyDataSet(Dataset):
def __init__(self,X,Y):
super().__init__()
self.xx,self.yy = X,Y
def __getitem__(self,index):
return self.xx[index],self.yy[index]
def __len__(self):
return len(self.yy)
trainDataSet = MyDataSet(x_train,y_train)
testDataSet = MyDataSet(x_test,y_test)
trainLoader = DataLoader(dataset=trainDataSet, batch_size=batch_size,shuffle=True)
testLoader = DataLoader(dataset=testDataSet, batch_size=batch_size,shuffle=True)
# In[8]:
class LSTM(nn.Module):
def __init__(self,input_size,hidden_size,num_layers):
super().__init__()
self.input_size = input_size
self.hidden_size = hidden_size
self.num_layers = num_layers
self.lstm = nn.LSTM(input_size = self.input_size,hidden_size = self.hidden_size,num_layers=self.num_layers,batch_first=True)
self.liner = nn.Linear(in_features=self.hidden_size,out_features=1)
def forward(self, x):
_, (hidden, cell) = self.lstm(x)
out = self.liner(hidden)
return out.reshape(-1,1)
model = LSTM(input_size,hidden_size,num_layers)
# In[9]:
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
loss_function = nn.MSELoss()
# In[10]:
epochs = 200
"""
data_x batch_size, seq_size, input_size=1
data_y batch_size, 1
"""
for epoch in range(epochs):
for i,(data_x,data_y) in enumerate(trainLoader):
pred = model(data_x)
pred = pred.reshape(-1)
loss = loss_function(pred, data_y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
if epoch % 10 ==0:
print("epoch:{},loss========{}".format(epoch+10,loss.item()))
# In[11]:
y_pred = []
y_label = []
with torch.no_grad():
pred = model(torch.from_numpy(np.array(X)))
y_pred.extend(sc.inverse_transform(pred))
y_label = np.array(Y).reshape(-1,1)
y_label = sc.inverse_transform(y_label)
plt.figure(figsize=(16,4))
plt.plot(y_label)
plt.plot(y_pred)
plt.show()
plt.figure(figsize=(16,4))
plt.plot(y_label[0:100],label='raw data')
plt.plot(y_pred[0:100], label='pred data')
plt.legend()
plt.show()