Files
ahad-quant/transformer_gru_model.py
2026-06-25 14:00:20 +03:00

360 lines
14 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters
This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.
"""
AHAD QUANT Forex V5 — TransformerGRU Hybrid Model
Combines a Bidirectional GRU (short-range pattern capture) with a
Transformer Encoder (long-range self-attention) for binary classification.
Architecture
───────────────────────────────────────────────────────────────
Input : (B, T, 62)
Layer 1: Input projection — Linear(62, D_MODEL) + GELU
→ (B, T, D_MODEL)
Layer 2: Positional Encoding (sinusoidal, fixed)
→ (B, T, D_MODEL)
Layer 3: Bidirectional GRU — hidden=GRU_HIDDEN, 2 layers
— captures local sequential dependencies in both directions
→ (B, T, D_MODEL) [GRU_HIDDEN*2 projected back to D_MODEL]
Layer 4: Transformer Encoder — N_ATT_LAYERS × (MHA + FFN + LayerNorm)
— long-range temporal attention on top of GRU hidden states
→ (B, T, D_MODEL)
Layer 5: Aggregation — mean pooling + last token concatenated
→ (B, D_MODEL*2)
Layer 6: Classification head
Linear → GELU → Dropout → Linear(1) — logit
Output : (B,) logit — apply sigmoid for probability
Training util functions
───────────────────────
train_tgru(X_seq_tr, y_tr, X_seq_va, y_va, ...) → TransformerGRU
predict_tgru_proba(model, X_seq_tensor) → np.ndarray (B,)
"""
import math
import time
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
# ─── Hyper-parameters (can be overridden via kwargs in train_tgru) ───────────
D_MODEL = 128 # main embedding dimension
GRU_HIDDEN = 64 # GRU hidden per direction (×2 bidir = D_MODEL)
N_ATT_LAYERS = 3 # Transformer encoder depth
N_HEADS = 4 # attention heads (D_MODEL must be divisible)
FFN_DIM = 256 # feedforward expansion in Transformer
DROPOUT = 0.10
BATCH_SIZE = 512
LR = 5e-4
WEIGHT_DECAY = 1e-4
MAX_EPOCHS = 30
PATIENCE = 6
GRAD_CLIP = 1.0
# ─── Positional Encoding (sinusoidal) ────────────────────────────────────────
class SinusoidalPositionalEncoding(nn.Module):
"""
Fixed sinusoidal positional encoding added to embeddings.
PE(pos, 2i) = sin(pos / 10000^(2i/d_model))
PE(pos, 2i+1) = cos(pos / 10000^(2i/d_model))
"""
def __init__(self, d_model: int, max_len: int = 512, dropout: float = 0.1):
super().__init__()
self.drop = nn.Dropout(dropout)
pe = torch.zeros(max_len, d_model)
pos = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
div = torch.exp(
torch.arange(0, d_model, 2, dtype=torch.float) * (-math.log(10000.0) / d_model)
)
pe[:, 0::2] = torch.sin(pos * div)
pe[:, 1::2] = torch.cos(pos * div)
self.register_buffer("pe", pe.unsqueeze(0)) # (1, max_len, d_model)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""x : (B, T, d_model)"""
return self.drop(x + self.pe[:, : x.size(1)])
# ─── Transformer Encoder Block ───────────────────────────────────────────────
class TransformerEncoderBlock(nn.Module):
"""
Standard Pre-LN Transformer encoder block:
x → LayerNorm → MHA → residual
→ LayerNorm → FFN → residual
Pre-LayerNorm (before attention) provides more stable training.
"""
def __init__(self, d_model: int, n_heads: int, ffn_dim: int,
dropout: float = DROPOUT):
super().__init__()
self.norm1 = nn.LayerNorm(d_model)
self.attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)
self.norm2 = nn.LayerNorm(d_model)
self.ffn = nn.Sequential(
nn.Linear(d_model, ffn_dim),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(ffn_dim, d_model),
nn.Dropout(dropout),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
# Self-attention sub-layer
h, _ = self.attn(self.norm1(x), self.norm1(x), self.norm1(x))
x = x + h
# Feed-forward sub-layer
x = x + self.ffn(self.norm2(x))
return x
# ─── Main model ──────────────────────────────────────────────────────────────
class TransformerGRU(nn.Module):
"""
TransformerGRU for binary classification on (B, T, F) time-series.
Output: (B,) raw logit — use with BCEWithLogitsLoss.
"""
def __init__(
self,
num_features : int = 62,
seq_len : int = 168,
d_model : int = D_MODEL,
gru_hidden : int = GRU_HIDDEN,
n_att_layers : int = N_ATT_LAYERS,
n_heads : int = N_HEADS,
ffn_dim : int = FFN_DIM,
dropout : float = DROPOUT,
):
super().__init__()
assert d_model == gru_hidden * 2, (
f"D_MODEL ({d_model}) must equal GRU_HIDDEN*2 ({gru_hidden*2}) "
"so bidirectional GRU output aligns with d_model"
)
self.d_model = d_model
self.seq_len = seq_len
# ── 1. Input projection ──────────────────────────────────────────────
self.input_proj = nn.Sequential(
nn.Linear(num_features, d_model),
nn.GELU(),
nn.Dropout(dropout),
)
# ── 2. Positional encoding ───────────────────────────────────────────
self.pos_enc = SinusoidalPositionalEncoding(d_model, max_len=seq_len + 2, dropout=dropout)
# ── 3. Bidirectional GRU ─────────────────────────────────────────────
# Input: (B, T, d_model) → Output: (B, T, gru_hidden*2 = d_model)
self.gru = nn.GRU(
input_size = d_model,
hidden_size = gru_hidden,
num_layers = 2,
batch_first = True,
bidirectional = True,
dropout = dropout,
)
self.gru_norm = nn.LayerNorm(d_model)
# Residual projection (input_proj output → d_model already, residual fits)
# ── 4. Transformer encoder ───────────────────────────────────────────
self.transformer = nn.ModuleList([
TransformerEncoderBlock(d_model, n_heads, ffn_dim, dropout)
for _ in range(n_att_layers)
])
self.final_norm = nn.LayerNorm(d_model)
# ── 5. Classification head ───────────────────────────────────────────
# Aggregate: mean pool + last token → concat → d_model*2
self.head = nn.Sequential(
nn.Linear(d_model * 2, d_model),
nn.GELU(),
nn.Dropout(dropout),
nn.Linear(d_model, 1),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""
x : (B, T, num_features)
→ (B,) logit
"""
# 1. Input projection + positional encoding
h = self.input_proj(x) # (B, T, d_model)
h = self.pos_enc(h)
# 2. Bidirectional GRU — add residual
gru_out, _ = self.gru(h) # (B, T, d_model) [bidir → gru_h*2]
h = self.gru_norm(h + gru_out) # add+norm residual
# 3. Transformer encoder
for block in self.transformer:
h = block(h) # (B, T, d_model)
h = self.final_norm(h)
# 4. Aggregation: mean pool + last token
mean_pool = h.mean(dim=1) # (B, d_model)
last_tok = h[:, -1, :] # (B, d_model)
pooled = torch.cat([mean_pool, last_tok], dim=-1) # (B, d_model*2)
# 5. Head
return self.head(pooled).squeeze(-1) # (B,)
# ─── Training utilities ──────────────────────────────────────────────────────
def _make_loader(X: np.ndarray, y: np.ndarray,
batch_size: int, shuffle: bool) -> DataLoader:
X_t = torch.FloatTensor(X)
y_t = torch.FloatTensor(y)
return DataLoader(TensorDataset(X_t, y_t),
batch_size=batch_size, shuffle=shuffle,
num_workers=0, pin_memory=torch.cuda.is_available())
def train_tgru(
X_seq_tr : np.ndarray,
y_tr : np.ndarray,
X_seq_va : np.ndarray,
y_va : np.ndarray,
*,
# Architecture overrides
d_model : int = D_MODEL,
gru_hidden : int = GRU_HIDDEN,
n_att_layers : int = N_ATT_LAYERS,
n_heads : int = N_HEADS,
ffn_dim : int = FFN_DIM,
dropout : float = DROPOUT,
# Training overrides
batch_size : int = BATCH_SIZE,
lr : float = LR,
weight_decay : float = WEIGHT_DECAY,
max_epochs : int = MAX_EPOCHS,
patience : int = PATIENCE,
grad_clip : float = GRAD_CLIP,
) -> "TransformerGRU":
"""
Train a TransformerGRU on pre-normalised sequence data.
Parameters
----------
X_seq_tr / X_seq_va : (N, T, 62) float32 — already normalised
y_tr / y_va : (N,) float32 binary labels
Returns
-------
Best model (lowest val loss) as a CPU-resident TransformerGRU.
"""
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
print(f" [TGRU] device={device} | "
f"train={len(y_tr):,} val={len(y_va):,} samples")
_, T, F = X_seq_tr.shape
model = TransformerGRU(
num_features = F,
seq_len = T,
d_model = d_model,
gru_hidden = gru_hidden,
n_att_layers = n_att_layers,
n_heads = n_heads,
ffn_dim = ffn_dim,
dropout = dropout,
).to(device)
n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
print(f" [TGRU] parameters: {n_params:,}")
optimiser = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimiser, T_max=max_epochs)
criterion = nn.BCEWithLogitsLoss()
scaler_amp = torch.cuda.amp.GradScaler(enabled=torch.cuda.is_available())
train_loader = _make_loader(X_seq_tr, y_tr, batch_size, shuffle=True)
val_loader = _make_loader(X_seq_va, y_va, batch_size * 2, shuffle=False)
best_val_loss = float("inf")
best_state = None
no_improve = 0
t0 = time.time()
for epoch in range(1, max_epochs + 1):
# ── Train ────────────────────────────────────────────────────────────
model.train()
train_loss = 0.0
for X_b, y_b in train_loader:
X_b, y_b = X_b.to(device), y_b.to(device)
optimiser.zero_grad()
with torch.cuda.amp.autocast(enabled=torch.cuda.is_available()):
logits = model(X_b)
loss = criterion(logits, y_b)
scaler_amp.scale(loss).backward()
scaler_amp.unscale_(optimiser)
nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
scaler_amp.step(optimiser)
scaler_amp.update()
train_loss += loss.item() * len(y_b)
train_loss /= len(y_tr)
# ── Validate ─────────────────────────────────────────────────────────
model.eval()
val_loss = 0.0
correct = 0
with torch.no_grad():
for X_b, y_b in val_loader:
X_b, y_b = X_b.to(device), y_b.to(device)
logits = model(X_b)
val_loss += criterion(logits, y_b).item() * len(y_b)
correct += ((logits > 0).float() == y_b).sum().item()
val_loss /= len(y_va)
val_acc = correct / len(y_va)
scheduler.step()
elapsed = time.time() - t0
print(f" [TGRU] epoch {epoch:3d}/{max_epochs} "
f"train={train_loss:.4f} val={val_loss:.4f} "
f"val_acc={val_acc:.4f} ({elapsed:.0f}s)")
# ── Early stopping ───────────────────────────────────────────────────
if val_loss < best_val_loss - 1e-5:
best_val_loss = val_loss
best_state = {k: v.cpu().clone() for k, v in model.state_dict().items()}
no_improve = 0
else:
no_improve += 1
if no_improve >= patience:
print(f" [TGRU] early stopping at epoch {epoch}")
break
model.load_state_dict(best_state)
model.cpu()
print(f" [TGRU] training complete — best val_loss={best_val_loss:.4f}")
return model
@torch.no_grad()
def predict_tgru_proba(
model : TransformerGRU,
X_seq_t : torch.Tensor,
batch_size: int = 1024,
device : str = "cpu",
) -> np.ndarray:
"""
Run inference on a (N, T, F) float tensor.
Returns probability array of shape (N,).
"""
model.eval()
model.to(device)
probs = []
for start in range(0, len(X_seq_t), batch_size):
chunk = X_seq_t[start : start + batch_size].to(device)
logits = model(chunk)
probs.append(torch.sigmoid(logits).cpu().numpy())
model.cpu()
return np.concatenate(probs).astype(np.float32)