Abdoul Ahad Binizi
382 lines
16 KiB
Python
382 lines
16 KiB
Python
"""
|
||
AHAD QUANT Forex V5 — Temporal Fusion Transformer (TFT)
|
||
Inspired by Lim et al., 2021 — adapted for binary classification on Forex 1h data.
|
||
|
||
Architecture
|
||
───────────────────────────────────────────────────────────────
|
||
Input : (B, T, 62) — B batches, T=168 timesteps, 62 features
|
||
Layer 1: Variable Selection Network (VSN)
|
||
— per-timestep soft feature weighting via GRN + Softmax
|
||
→ (B, T, D_MODEL)
|
||
Layer 2: Local LSTM encoder
|
||
— captures short-range sequential dependencies
|
||
→ (B, T, D_MODEL)
|
||
Layer 3: Gated skip connection (GRN) with add+norm
|
||
→ (B, T, D_MODEL)
|
||
Layer 4: Temporal Multi-Head Self-Attention (N_HEADS heads)
|
||
— captures long-range dependencies
|
||
→ (B, T, D_MODEL)
|
||
Layer 5: Gated skip connection + mean pooling
|
||
→ (B, D_MODEL)
|
||
Layer 6: Classification head
|
||
GRN → Linear(D_MODEL, 1) — logit (use with BCEWithLogitsLoss)
|
||
Output : (B,) logit — apply sigmoid for probability
|
||
|
||
Training util functions
|
||
───────────────────────
|
||
train_tft(X_seq_tr, y_tr, X_seq_va, y_va, seq_scaler, ...) → TemporalFusionTransformer
|
||
predict_tft_proba(model, X_seq_norm_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_tft) ────────────
|
||
|
||
D_MODEL = 64 # embedding / hidden dimension
|
||
N_HEADS = 4 # attention heads (D_MODEL must be divisible by N_HEADS)
|
||
LSTM_LAYERS = 1 # LSTM depth
|
||
N_ATT_LAYERS = 2 # stacked temporal attention blocks
|
||
DROPOUT = 0.10
|
||
BATCH_SIZE = 512
|
||
LR = 8e-4 # raised from 5e-4: VSN's 62 parallel GRNs need more
|
||
# headroom once GRAD_CLIP is loosened (see below)
|
||
WEIGHT_DECAY = 1e-4
|
||
MAX_EPOCHS = 30
|
||
PATIENCE = 6 # early stopping patience
|
||
GRAD_CLIP = 5.0 # raised from 1.0: the VSN's 62 independent per-feature
|
||
# GRNs produce a much larger raw gradient norm than a
|
||
# single shared embedding (e.g. TransformerGRU's),
|
||
# so a clip of 1.0 over-throttled ALL parameters and
|
||
# kept loss stuck near ln(2). 5.0 lets real updates
|
||
# through while still preventing rare gradient spikes.
|
||
|
||
|
||
# ─── Building blocks ─────────────────────────────────────────────────────────
|
||
|
||
class GatedResidualNetwork(nn.Module):
|
||
"""
|
||
Core TFT building block.
|
||
x → ELU(W1·x) → Dropout → W2·(·) → GLU gate → LayerNorm(skip + output)
|
||
|
||
Supports optional context vector (c) injected before the first linear.
|
||
"""
|
||
|
||
def __init__(self, input_dim: int, hidden_dim: int, output_dim: int,
|
||
dropout: float = DROPOUT, context_dim: int = 0):
|
||
super().__init__()
|
||
self.fc1 = nn.Linear(input_dim + context_dim, hidden_dim)
|
||
self.fc2 = nn.Linear(hidden_dim, output_dim * 2) # × 2 for GLU
|
||
# Skip connection: project input → output dim if they differ
|
||
self.skip = nn.Linear(input_dim, output_dim) if input_dim != output_dim else nn.Identity()
|
||
self.norm = nn.LayerNorm(output_dim)
|
||
self.drop = nn.Dropout(dropout)
|
||
self.out_dim = output_dim
|
||
|
||
def forward(self, x: torch.Tensor, context: torch.Tensor | None = None) -> torch.Tensor:
|
||
h = x if context is None else torch.cat([x, context], dim=-1)
|
||
h = F.elu(self.fc1(h))
|
||
h = self.drop(h)
|
||
h = self.fc2(h) # (*, out_dim*2)
|
||
h1, h2 = h.chunk(2, dim=-1) # GLU split
|
||
h = h1 * torch.sigmoid(h2) # Gated Linear Unit
|
||
return self.norm(self.skip(x) + h)
|
||
|
||
|
||
class VariableSelectionNetwork(nn.Module):
|
||
"""
|
||
Soft per-timestep feature selection.
|
||
|
||
Each of the `num_features` features is projected to D_MODEL via its own GRN.
|
||
A shared GRN over the concatenated input produces Softmax weights.
|
||
Output = weighted sum of per-feature projections → (B, T, D_MODEL).
|
||
"""
|
||
|
||
def __init__(self, num_features: int, d_model: int = D_MODEL,
|
||
dropout: float = DROPOUT):
|
||
super().__init__()
|
||
self.d_model = d_model
|
||
self.num_features = num_features
|
||
|
||
# One GRN per feature: (B, T, 1) → (B, T, d_model)
|
||
self.var_grns = nn.ModuleList([
|
||
GatedResidualNetwork(1, d_model, d_model, dropout)
|
||
for _ in range(num_features)
|
||
])
|
||
|
||
# Selection GRN: (B, T, num_features) → weights (B, T, num_features)
|
||
self.select_grn = GatedResidualNetwork(num_features, d_model, num_features, dropout)
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
x : (B, T, num_features)
|
||
→ (B, T, d_model)
|
||
"""
|
||
B, T, num_feat = x.shape
|
||
|
||
# Per-feature embeddings: (B*T, 1) → (B*T, d_model) → stack → (B, T, F, d_model)
|
||
xi = x.reshape(B * T, num_feat) # (B*T, F)
|
||
embeddings = torch.stack(
|
||
[grn(xi[:, i : i + 1]) for i, grn in enumerate(self.var_grns)],
|
||
dim=1 # (B*T, F, d_model)
|
||
) # → (B*T, F, d_model)
|
||
|
||
# Selection weights
|
||
weights = self.select_grn(xi) # (B*T, F) logits
|
||
weights = F.softmax(weights, dim=-1) # (B*T, F) soft weights
|
||
|
||
# Weighted sum over features
|
||
out = (embeddings * weights.unsqueeze(-1)).sum(dim=1) # (B*T, d_model)
|
||
return out.reshape(B, T, self.d_model)
|
||
|
||
|
||
class TemporalAttentionBlock(nn.Module):
|
||
"""One layer of temporal self-attention + gated residual."""
|
||
|
||
def __init__(self, d_model: int = D_MODEL, n_heads: int = N_HEADS,
|
||
dropout: float = DROPOUT):
|
||
super().__init__()
|
||
self.attn = nn.MultiheadAttention(d_model, n_heads, dropout=dropout, batch_first=True)
|
||
self.gate = GatedResidualNetwork(d_model, d_model, d_model, dropout)
|
||
self.norm = nn.LayerNorm(d_model)
|
||
self.drop = nn.Dropout(dropout)
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
attn_out, _ = self.attn(x, x, x) # (B, T, d_model)
|
||
attn_out = self.drop(attn_out)
|
||
return self.norm(x + self.gate(attn_out)) # residual
|
||
|
||
|
||
# ─── Main model ──────────────────────────────────────────────────────────────
|
||
|
||
class TemporalFusionTransformer(nn.Module):
|
||
"""
|
||
TFT for binary classification on time-series of shape (B, T, num_features).
|
||
Output: (B,) raw logit — use with BCEWithLogitsLoss.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
num_features : int = 62,
|
||
seq_len : int = 168,
|
||
d_model : int = D_MODEL,
|
||
n_heads : int = N_HEADS,
|
||
lstm_layers : int = LSTM_LAYERS,
|
||
n_att_layers : int = N_ATT_LAYERS,
|
||
dropout : float = DROPOUT,
|
||
):
|
||
super().__init__()
|
||
self.seq_len = seq_len
|
||
|
||
# ── 1. Input embedding ───────────────────────────────────────────────
|
||
# NOTE: previously used a VariableSelectionNetwork with one independent
|
||
# GatedResidualNetwork PER FEATURE (62 branches), combined via softmax-
|
||
# weighted average. Averaging ~62 quasi-independent branches collapses
|
||
# inter-sample variance by a factor of ~sqrt(num_features) (classic
|
||
# noise-cancellation effect), starving the model of signal before
|
||
# training even starts — confirmed empirically (logit std stuck ~0.03
|
||
# across 60 training steps even on an easy synthetic task, vs healthy
|
||
# growth to ~0.65 with this simpler shared projection). Replaced with
|
||
# a single shared Linear + GRN, matching the embedding strategy that
|
||
# already works well in transformer_gru_model.py.
|
||
self.input_proj = nn.Linear(num_features, d_model)
|
||
self.input_grn = GatedResidualNetwork(d_model, d_model, d_model, dropout)
|
||
|
||
# ── 2. Local LSTM encoder ────────────────────────────────────────────
|
||
self.lstm = nn.LSTM(
|
||
input_size = d_model,
|
||
hidden_size = d_model,
|
||
num_layers = lstm_layers,
|
||
batch_first = True,
|
||
dropout = dropout if lstm_layers > 1 else 0.0,
|
||
)
|
||
self.lstm_gate = GatedResidualNetwork(d_model, d_model, d_model, dropout)
|
||
self.lstm_norm = nn.LayerNorm(d_model)
|
||
|
||
# ── 3. Temporal attention stack ──────────────────────────────────────
|
||
self.attn_blocks = nn.ModuleList([
|
||
TemporalAttentionBlock(d_model, n_heads, dropout)
|
||
for _ in range(n_att_layers)
|
||
])
|
||
|
||
# ── 4. Output head ───────────────────────────────────────────────────
|
||
self.output_grn = GatedResidualNetwork(d_model, d_model, d_model, dropout)
|
||
self.head = nn.Linear(d_model, 1)
|
||
|
||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||
"""
|
||
x : (B, T, num_features)
|
||
→ (B,) logit
|
||
"""
|
||
# 1. Shared input embedding (see note in __init__)
|
||
h = self.input_grn(self.input_proj(x)) # (B, T, d_model)
|
||
|
||
# 2. LSTM local encoder + gated skip
|
||
lstm_out, _ = self.lstm(h) # (B, T, d_model)
|
||
h = self.lstm_norm(h + self.lstm_gate(lstm_out)) # add+norm
|
||
|
||
# 3. Stacked temporal attention
|
||
for block in self.attn_blocks:
|
||
h = block(h) # (B, T, d_model)
|
||
|
||
# 4. Mean pooling over time + classification
|
||
h = h.mean(dim=1) # (B, d_model)
|
||
h = self.output_grn(h)
|
||
return self.head(h).squeeze(-1) # (B,)
|
||
|
||
|
||
# ─── Training utilities ──────────────────────────────────────────────────────
|
||
|
||
def _make_loader(X: np.ndarray, y: np.ndarray,
|
||
batch_size: int, shuffle: bool) -> DataLoader:
|
||
"""Wrap numpy arrays in a TensorDataset / 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_tft(
|
||
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,
|
||
n_heads : int = N_HEADS,
|
||
n_att_layers : int = N_ATT_LAYERS,
|
||
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,
|
||
) -> "TemporalFusionTransformer":
|
||
"""
|
||
Train a TemporalFusionTransformer 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 TemporalFusionTransformer.
|
||
"""
|
||
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
|
||
print(f" [TFT] device={device} | "
|
||
f"train={len(y_tr):,} val={len(y_va):,} samples")
|
||
|
||
_, T, F = X_seq_tr.shape
|
||
model = TemporalFusionTransformer(
|
||
num_features = F,
|
||
seq_len = T,
|
||
d_model = d_model,
|
||
n_heads = n_heads,
|
||
n_att_layers = n_att_layers,
|
||
dropout = dropout,
|
||
).to(device)
|
||
|
||
n_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
|
||
print(f" [TFT] 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" [TFT] 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" [TFT] early stopping at epoch {epoch}")
|
||
break
|
||
|
||
model.load_state_dict(best_state)
|
||
model.cpu()
|
||
print(f" [TFT] training complete — best val_loss={best_val_loss:.4f}")
|
||
return model
|
||
|
||
|
||
@torch.no_grad()
|
||
def predict_tft_proba(
|
||
model : TemporalFusionTransformer,
|
||
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)
|