lstm_model.py

import torch
import torch.nn as nn
import numpy as np
class StockLSTM(nn.Module):
    def __init__(self,input_size,hidden_size,num_layers):
        super(StockLSTM,self).__init__()
        self.lstm=nn.LSTM(input_size,hidden_size,num_layers,batch_first=True)
        self.dropout=nn.Dropout(0.2)
        self.fc=nn.Linear(hidden_size,1)
        self.sigmoid=nn.Sigmoid()
    def forward(self,x):
        out, _=self.lstm(x)
        out=self.dropout(out)
        out=out[:,-1,:]
        out=self.fc(out)
        out=self.sigmoid(out)
        return out
def prepare_sequences(df,predcitors):
    x=[]
    y=[]
    for i in range(60,len(df)):

        x.append(df[predcitors].iloc[i-60:i].values)
        y.append(df["Target"].iloc[i])
    return x,y
def train_lstm(df,predictors,hidden_size=64,num_layers=2,epoch=50,lr=0.01):
    x,y=prepare_sequences(df,predictors)
    x_tensor=torch.FloatTensor(np.array(x))
    y_tensor=torch.FloatTensor(np.array(y)).unsqueeze(1)
    model=StockLSTM(input_size=len(predictors),hidden_size=hidden_size,num_layers=num_layers)
    optimizer=torch.optim.Adam(model.parameters(),lr=lr)
    for i in range(epoch):
        optimizer.zero_grad()
        output=model(x_tensor)
        criterion=nn.BCELoss()
        loss=criterion(output,y_tensor)
        loss.backward()
        optimizer.step()
    return model
def predict_lstm(model,df,predictors,sequence_length=60):
    input=df[predictors].iloc[-sequence_length:].values
    input=torch.FloatTensor(input).unsqueeze(0)
    with torch.no_grad():
        prob=(model(input).item())
    return prob