From 2673e27a40afe34b4e455282bf0423cc8d4b85da Mon Sep 17 00:00:00 2001 From: Florian Egger Date: Tue, 26 May 2026 15:17:00 +0200 Subject: [PATCH] Fix head divisibility validation, benchmark data format, and target computation bugs - 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) --- src/amd/optimizations.py | 8 ++++- src/data/pipeline.py | 64 +++++++++++++++++---------------- src/models/gnn_model.py | 8 ++++- src/models/intraday_gnn.py | 11 ++++-- src/models/trainer.py | 10 ++++-- src/web/api/dashboard.py | 3 -- src/web/api/models_endpoints.py | 8 +++-- 7 files changed, 68 insertions(+), 44 deletions(-) diff --git a/src/amd/optimizations.py b/src/amd/optimizations.py index 28026dd..203b703 100644 --- a/src/amd/optimizations.py +++ b/src/amd/optimizations.py @@ -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: diff --git a/src/data/pipeline.py b/src/data/pipeline.py index 4b8ac31..d303fe3 100644 --- a/src/data/pipeline.py +++ b/src/data/pipeline.py @@ -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) diff --git a/src/models/gnn_model.py b/src/models/gnn_model.py index 9fa532c..fe177da 100644 --- a/src/models/gnn_model.py +++ b/src/models/gnn_model.py @@ -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 diff --git a/src/models/intraday_gnn.py b/src/models/intraday_gnn.py index 904482c..b3b4acc 100644 --- a/src/models/intraday_gnn.py +++ b/src/models/intraday_gnn.py @@ -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) diff --git a/src/models/trainer.py b/src/models/trainer.py index f35ff4e..bf5bb11 100644 --- a/src/models/trainer.py +++ b/src/models/trainer.py @@ -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(), ): diff --git a/src/web/api/dashboard.py b/src/web/api/dashboard.py index 793762e..9e0b350 100644 --- a/src/web/api/dashboard.py +++ b/src/web/api/dashboard.py @@ -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 diff --git a/src/web/api/models_endpoints.py b/src/web/api/models_endpoints.py index 7922538..d66f73b 100644 --- a/src/web/api/models_endpoints.py +++ b/src/web/api/models_endpoints.py @@ -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()