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)
|