360 lines
14 KiB
Python
360 lines
14 KiB
Python
"""
|
||
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)
|