fegger/trading_gnn
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix head divisibility validation, benchmark data format, and target computation bugs

Browse Source 
- Add embed_dim/feature_dim divisibility checks in attention modules
- Auto-select largest valid head count in IntradayGNN when num_features isn't divisible by 8
- Fix benchmark to pass PyG Data objects instead of tuples to model
- Fix target computation to use precomputed next trading date and filter valid tickers
- Add missing device argument to autocast calls
- Remove stale global app_state reference in dashboard API
- Clean up various formatting (line wrapping, whitespace)
This commit is contained in:
Florian Egger
2026-05-26 15:17:00 +02:00
parent 46657c7ffe
commit 2673e27a40
7 changed files with 68 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/amd/optimizations.py
+7 -1
View File
@@ -165,6 +165,10 @@ class AMDSparseAttention(nn.Module):
def __init__(self, embed_dim, num_heads, dropout=0.1):
super().__init__()
if embed_dim % num_heads != 0:
raise ValueError(
f"embed_dim ({embed_dim}) must be divisible by num_heads ({num_heads})"
)
self.embed_dim = embed_dim
self.num_heads = num_heads
self.head_dim = embed_dim // num_heads
@@ -291,7 +295,9 @@ class AMDGATConv(MessagePassing):
alpha = (alpha_src, alpha_dst)
# Propagate — pass edge_attr so correlation weights reach message()
out = self.propagate(edge_index, x=(x_src, x_dst), alpha=alpha, edge_attr=edge_attr, size=size)
out = self.propagate(
edge_index, x=(x_src, x_dst), alpha=alpha, edge_attr=edge_attr, size=size
)
# Concatenate or average heads
if self.concat:
src/data/pipeline.py
+33 -31
View File
@@ -309,7 +309,9 @@ class StockDataPipeline:
"high": row["High"],
"low": row["Low"],
"close": row["Close"],
"adj_close": row["Close"], # auto_adjust=True, 'Close' is already adjusted
"adj_close": row[
"Close"
], # auto_adjust=True, 'Close' is already adjusted
"volume": row["Volume"],
}
)
@@ -581,9 +583,7 @@ class StockDataPipeline:
if len(seq_arr) > config.SEQUENCE_LENGTH:
seq_arr = seq_arr[-config.SEQUENCE_LENGTH :]
elif len(seq_arr) < config.SEQUENCE_LENGTH:
pad = np.zeros(
(config.SEQUENCE_LENGTH - len(seq_arr), 5), dtype=np.float32
)
pad = np.zeros((config.SEQUENCE_LENGTH - len(seq_arr), 5), dtype=np.float32)
seq_arr = np.vstack([pad, seq_arr])
# News features (current date, broadcast across all timesteps)
@@ -710,9 +710,7 @@ class StockDataPipeline:
continue
# Pre-fetch PIT data once per ticker (avoids O(n²) repeated calls)
pit_cache = {
t: self.get_point_in_time_data(t, date) for t in valid_tickers
}
pit_cache = {t: self.get_point_in_time_data(t, date) for t in valid_tickers}
# Build 3-D node feature matrix
node_sequences: List[np.ndarray] = []
@@ -834,16 +832,33 @@ class StockDataPipeline:
# Compute once per day — stock universe and PIT data don't change intraday
pd_date = pd.Timestamp(date)
current_tickers = [
t for t in tickers
t
for t in tickers
if t in self.price_data
and not self.price_data[t].empty
and self.price_data[t].index[0] <= pd_date <= self.price_data[t].index[-1]
and self.price_data[t].index[0]
<= pd_date
<= self.price_data[t].index[-1]
]
if not current_tickers:
continue
date_dt = datetime.strptime(date, "%Y-%m-%d")
pit_cache = {t: self.get_point_in_time_data(t, date_dt) for t in current_tickers}
pit_cache = {
t: self.get_point_in_time_data(t, date_dt) for t in current_tickers
}
# Determine next trading date once per day
next_date = self._next_trading_date(pd_date)
if next_date is None:
continue
# Pre-filter tickers that have target data on the next trading date
target_valid_tickers = [
t for t in current_tickers if next_date in self.price_data[t].index
]
if not target_valid_tickers:
continue
# Get all timestamps for this trading day
timestamps = generate_intraday_timestamps(date)
@@ -856,7 +871,7 @@ class StockDataPipeline:
sequence_features = []
valid_tickers = []
for ticker in current_tickers:
for ticker in target_valid_tickers:
# Check memory before processing ticker
if not self.memory_manager.ensure_memory(50 * 1024**2): # 50MB
logger.warning(f"Skipping {ticker} due to memory constraints")
@@ -881,8 +896,7 @@ class StockDataPipeline:
if isinstance(daily_data, pd.Series):
# If only one day of data, create a sequence with the same values
features = [
daily_data["Close"] / daily_data["Open"]
- 1, # Return
daily_data["Close"] / daily_data["Open"] - 1, # Return
0.2, # Volatility (placeholder)
0.0, # Momentum (placeholder)
np.log(daily_data["Volume"] + 1), # Log volume
@@ -980,26 +994,14 @@ class StockDataPipeline:
)
# Create target (returns over prediction horizon)
y = []
y_vals = []
for ticker in valid_tickers:
# For intraday, we would predict the next few minutes
# For this example, we'll predict the next day's return
next_date = datetime.strptime(date, "%Y-%m-%d") + timedelta(days=1)
if (
ticker in self.price_data
and next_date.strftime("%Y-%m-%d")
in self.price_data[ticker].index
):
current_price = self.price_data[ticker].loc[date]["Close"]
future_price = self.price_data[ticker].loc[
next_date.strftime("%Y-%m-%d")
]["Close"]
ret = future_price / current_price - 1
y.append(ret)
else:
y.append(0) # Default value
current_price = self.price_data[ticker].loc[date]["Close"]
future_price = self.price_data[ticker].loc[next_date]["Close"]
ret = future_price / current_price - 1
y_vals.append(ret)
y = torch.tensor(y, dtype=torch.float32).unsqueeze(1)
y = torch.tensor(y_vals, dtype=torch.float32).unsqueeze(1)
# Create Data object
data = Data(x=x, edge_index=edge_index, edge_attr=edge_weight, y=y)
src/models/gnn_model.py
+7 -1
View File
@@ -16,6 +16,10 @@ class TemporalAttention(nn.Module):
def __init__(self, feature_dim: int, num_heads: int = 8, dropout: float = 0.1):
super().__init__()
if feature_dim % num_heads != 0:
raise ValueError(
f"feature_dim ({feature_dim}) must be divisible by num_heads ({num_heads})"
)
self.feature_dim = feature_dim
self.num_heads = num_heads
self.head_dim = feature_dim // num_heads
@@ -308,7 +312,9 @@ class CorporateActionAwareGNN(nn.Module):
price_attended = self.temporal_attention(self.price_processor(price_features))
news_attended = self.temporal_attention(self.news_processor(news_features))
social_attended = self.temporal_attention(self.social_processor(social_features))
social_attended = self.temporal_attention(
self.social_processor(social_features)
)
alternative_features = torch.cat(
[price_attended, news_attended, social_attended], dim=1
src/models/intraday_gnn.py
+8 -3
View File
@@ -24,7 +24,12 @@ class IntradayGNN(nn.Module):
self.num_features = num_features
# Temporal attention for sequence processing
self.temporal_attention = TemporalAttention(num_features, num_heads=8)
# num_features must be divisible by num_heads; 14 is not divisible by 8.
# Pick the largest divisor of num_features up to 8.
num_heads = max(
h for h in range(1, min(9, num_features + 1)) if num_features % h == 0
)
self.temporal_attention = TemporalAttention(num_features, num_heads=num_heads)
# Feature processing modules with AMD optimizations
self.feature_processor = nn.Sequential(
@@ -161,8 +166,8 @@ class IntradayGNN(nn.Module):
lstm_out = lstm_out.squeeze(1) # (N, HIDDEN_CHANNELS)
# Per-feature gating: sigmoid weights scale each channel independently
attention_weights = self.attention(lstm_out) # (N, HIDDEN_CHANNELS)
attended = lstm_out * attention_weights # (N, HIDDEN_CHANNELS)
attention_weights = self.attention(lstm_out) # (N, HIDDEN_CHANNELS)
attended = lstm_out * attention_weights # (N, HIDDEN_CHANNELS)
return self.linear(attended)
src/models/trainer.py
+8 -2
View File
@@ -143,7 +143,9 @@ class GNNTrainer:
self.memory_manager.empty_cache()
continue
epoch_train_loss = epoch_train_loss / num_train_batches if num_train_batches else 0.0
epoch_train_loss = (
epoch_train_loss / num_train_batches if num_train_batches else 0.0
)
self.train_losses.append(epoch_train_loss)
# Validation
@@ -185,7 +187,9 @@ class GNNTrainer:
self.memory_manager.empty_cache()
continue
epoch_val_loss = epoch_val_loss / num_val_batches if num_val_batches else 0.0
epoch_val_loss = (
epoch_val_loss / num_val_batches if num_val_batches else 0.0
)
self.val_losses.append(epoch_val_loss)
# Update learning rate scheduler
@@ -234,6 +238,7 @@ class GNNTrainer:
data = data.to(self.device, non_blocking=config.PIN_MEMORY)
with autocast(
config.DEVICE,
enabled=config.MIXED_PRECISION,
dtype=self.amd_optimizer.get_precision_dtype(),
):
@@ -343,6 +348,7 @@ class GNNTrainer:
data = data.to(self.device, non_blocking=config.PIN_MEMORY)
with autocast(
config.DEVICE,
enabled=config.MIXED_PRECISION,
dtype=self.amd_optimizer.get_precision_dtype(),
):
src/web/api/dashboard.py
-3
View File
@@ -13,9 +13,6 @@ from src.web.services.state import AppState
logger = logging.getLogger(__name__)
router = APIRouter()
# Reference to global app state (injected via module import in app.py)
app_state: AppState = None # type: ignore
def _get_state() -> AppState:
from src.web.app import app_state as _state
src/web/api/models_endpoints.py
+5 -3
View File
@@ -8,6 +8,7 @@ from datetime import datetime
from typing import Dict
from fastapi import APIRouter
from torch_geometric.data import Data
from config import config
@@ -151,16 +152,17 @@ async def run_benchmark() -> Dict:
edge_attr = torch.randn(num_edges, 1).to(config.DEVICE)
# Warm-up
dummy_data = Data(x=x, edge_index=edge_index, edge_attr=edge_attr)
for _ in range(10):
with torch.no_grad():
_ = model((x, edge_index, edge_attr))
_ = model(dummy_data)
# Benchmark inference
start = time.time()
num_runs = 100
for _ in range(num_runs):
with torch.no_grad():
_ = model((x, edge_index, edge_attr))
_ = model(dummy_data)
inference_time = (time.time() - start) / num_runs
# Benchmark training
@@ -172,7 +174,7 @@ async def run_benchmark() -> Dict:
start = time.time()
for _ in range(num_runs):
optimizer.zero_grad()
out = model((x, edge_index, edge_attr))
out = model(dummy_data)
loss = criterion(out, y)
loss.backward()
optimizer.step()