model_train.py

import pandas as pd
from sqlalchemy import create_engine
import torch
from sklearn.preprocessing import MinMaxScaler
import numpy as np
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader


device = "cuda" if torch.cuda.is_available() else "cpu"
print("Using device: ", device )

# Database connection info
DATABASE_URL = "postgresql://postgres:123456789@localhost/trader_master"

engine = create_engine(DATABASE_URL)

df = pd.read_sql("SELECT * FROM candle_tsm_minute_5;", engine)

df['time'] = pd.to_datetime(df['ts'], unit='s')

df = df.sort_values('time')

# data = df['c'].values.reshape(-1, 1) # 2D input

data = df[["o", "h", "l", "c", "v"]].values # shape(N, 5)

scaler = MinMaxScaler()
data_scaled = scaler.fit_transform(data)

# # INPUT_LEN = 60 * 24 * 15 # 1440 * 15 = 21600
INPUT_LEN = 60 * 5
PRED_INDEX = list(range(1, 13))

def create_sequence(data, input_len, pred_index):
    X, y = [], []
    max_h = max(pred_index)

    for i in range(len(data) - input_len - max_h):
        X.append(data[i : i + input_len])

        y_sub = []
        for h in pred_index:
            close_future = data[i + input_len + h -1][3]
            y_sub.append(close_future)

        y.append(y_sub)

    return np.array(X), np.array(y)

X, y = create_sequence(data_scaled, INPUT_LEN, PRED_INDEX)
print("X:", X.shape, "Y:", y.shape)

# Train/test split
split = int(len(X) * 0.95)
X_train, X_test = X[:split], X[split:]
y_train, y_test = y[:split], y[split:]

X_train = torch.tensor(X_train, dtype=torch.float32).to(device)
y_train = torch.tensor(y_train, dtype=torch.float32).to(device)
X_test = torch.tensor(X_test, dtype=torch.float32).to(device)
y_test = torch.tensor(y_test, dtype=torch.float32).to(device)

batch_size = 512

train_dataset = TensorDataset(X_train, y_train)
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)

class LSTMModel(nn.Module):
    def __init__(self):
        super(LSTMModel, self).__init__()

        self.lstm1 = nn.LSTM(
            input_size = 5,
            hidden_size = 64,
            num_layers = 1,
            batch_first = True
        )

        self.dropout1 = nn.Dropout(0.2)

        self.lstm2 = nn.LSTM(
            input_size = 64,
            hidden_size = 64,
            num_layers = 1,
            batch_first = True
        )

        self.dropout2 = nn.Dropout(0.2)

        self.fc1 = nn.Linear(64, 32)
        self.fc2 = nn.Linear(32, 12)

    def forward(self, x):
        x, _ = self.lstm1(x)
        x = self.dropout1(x)

        x, _ = self.lstm2(x)
        x = self.dropout2(x)

        x = x[:, -1, :]
        x = self.fc1(x)
        x = self.fc2(x)
        return x
    
model = LSTMModel().to(device)
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# Training
EPOCHS = 20

for epoch in range(EPOCHS):
    model.train()

    for batch_x, batch_y in train_loader:
        optimizer.zero_grad()

        y_pred = model(batch_x)
        loss = criterion(y_pred, batch_y)
        loss.backward()
        optimizer.step()

    print(f"Epoch {epoch}/{EPOCHS} Loss: {loss.item():.6f}")

# Predict
model.eval()

MODEL_PATH = "./model_train/model/tsm_model.pt"
torch.save(model.state_dict(), MODEL_PATH)

last_seq = data_scaled[-INPUT_LEN:]
last_seq = torch.tensor(last_seq, dtype=torch.float32).unsqueeze(0).to(device)

future_scaled = model(last_seq).detach().cpu().numpy()[0]

dummy = np.zeros((12, 5))
dummy[:, 3] = future_scaled

future = scaler.inverse_transform(dummy)[:, 3]

print(f"Predicted future close (5 -> 60 minutes): {future}")