221 lines
7.6 KiB
Python
221 lines
7.6 KiB
Python
import math
|
|
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from torch_geometric.nn import GATConv
|
|
|
|
from config import config
|
|
from src.amd.optimizations import AMDGATConv, AMDSparseAttention
|
|
from src.models.gnn_model import TemporalAttention
|
|
|
|
|
|
class IntradayGNN(nn.Module):
|
|
"""
|
|
Intraday GNN model optimized for AMD Radeon R9700 AI Pro
|
|
"""
|
|
|
|
def __init__(
|
|
self, num_features: int, sequence_length: int = config.SEQUENCE_LENGTH
|
|
):
|
|
super(IntradayGNN, self).__init__()
|
|
|
|
self.sequence_length = sequence_length
|
|
self.num_features = num_features
|
|
|
|
# Temporal attention for sequence processing
|
|
self.temporal_attention = TemporalAttention(num_features, num_heads=8)
|
|
|
|
# Feature processing modules with AMD optimizations
|
|
self.feature_processor = nn.Sequential(
|
|
nn.Linear(num_features, config.HIDDEN_CHANNELS),
|
|
nn.SiLU(), # Swish activation often works better on AMD GPUs
|
|
nn.Linear(config.HIDDEN_CHANNELS, config.HIDDEN_CHANNELS),
|
|
nn.LayerNorm(config.HIDDEN_CHANNELS),
|
|
)
|
|
|
|
# Graph attention layers with AMD optimizations
|
|
try:
|
|
self.conv1 = AMDGATConv(
|
|
config.HIDDEN_CHANNELS,
|
|
config.HIDDEN_CHANNELS,
|
|
heads=config.NUM_HEADS,
|
|
concat=True,
|
|
dropout=config.DROPOUT,
|
|
add_self_loops=True,
|
|
)
|
|
except Exception:
|
|
self.conv1 = GATConv(
|
|
config.HIDDEN_CHANNELS,
|
|
config.HIDDEN_CHANNELS,
|
|
heads=config.NUM_HEADS,
|
|
concat=True,
|
|
dropout=config.DROPOUT,
|
|
add_self_loops=True,
|
|
)
|
|
|
|
try:
|
|
self.conv2 = AMDGATConv(
|
|
config.HIDDEN_CHANNELS * config.NUM_HEADS,
|
|
config.HIDDEN_CHANNELS,
|
|
heads=1,
|
|
concat=False,
|
|
dropout=config.DROPOUT,
|
|
add_self_loops=True,
|
|
)
|
|
except Exception:
|
|
self.conv2 = GATConv(
|
|
config.HIDDEN_CHANNELS * config.NUM_HEADS,
|
|
config.HIDDEN_CHANNELS,
|
|
heads=1,
|
|
concat=False,
|
|
dropout=config.DROPOUT,
|
|
add_self_loops=True,
|
|
)
|
|
|
|
# LSTM for temporal dependencies with AMD optimizations
|
|
self.lstm = nn.LSTM(
|
|
input_size=config.HIDDEN_CHANNELS,
|
|
hidden_size=config.HIDDEN_CHANNELS,
|
|
num_layers=2,
|
|
batch_first=True,
|
|
dropout=config.DROPOUT if config.NUM_HEADS > 1 else 0,
|
|
bidirectional=False,
|
|
)
|
|
|
|
# Initialize LSTM weights
|
|
self._init_lstm_weights()
|
|
|
|
# Attention mechanism for final prediction
|
|
self.attention = nn.Sequential(
|
|
nn.Linear(config.HIDDEN_CHANNELS, config.HIDDEN_CHANNELS),
|
|
nn.SiLU(),
|
|
nn.Linear(config.HIDDEN_CHANNELS, 1),
|
|
nn.Softmax(dim=1),
|
|
)
|
|
|
|
# Final prediction layer
|
|
self.linear = nn.Linear(config.HIDDEN_CHANNELS, 1)
|
|
|
|
# State for online learning
|
|
self.hidden_state = None
|
|
|
|
# Initialize weights
|
|
self._init_weights()
|
|
|
|
def _init_weights(self):
|
|
"""Initialize weights with AMD-friendly initialization"""
|
|
# Initialize linear layer weights
|
|
nn.init.xavier_uniform_(self.linear.weight)
|
|
nn.init.zeros_(self.linear.bias)
|
|
|
|
# Initialize attention layers
|
|
for layer in self.attention:
|
|
if isinstance(layer, nn.Linear):
|
|
nn.init.xavier_uniform_(layer.weight)
|
|
nn.init.zeros_(layer.bias)
|
|
|
|
def _init_lstm_weights(self):
|
|
"""Initialize LSTM weights with orthogonal initialization for better convergence"""
|
|
for name, param in self.lstm.named_parameters():
|
|
if "weight_ih" in name:
|
|
nn.init.orthogonal_(param)
|
|
elif "weight_hh" in name:
|
|
nn.init.orthogonal_(param)
|
|
elif "bias" in name:
|
|
nn.init.zeros_(param)
|
|
# Set forget gate bias to 1 for better gradient flow
|
|
n = param.size(0)
|
|
param.data[n // 4 : n // 2].fill_(1)
|
|
|
|
def forward(self, data):
|
|
x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr
|
|
|
|
# x shape: (num_stocks, sequence_length, num_features)
|
|
batch_size, seq_len, num_features = x.size()
|
|
|
|
# Apply temporal attention to each stock's sequence
|
|
temporal_features = []
|
|
for i in range(batch_size):
|
|
stock_sequence = x[i].unsqueeze(0) # (1, sequence_length, num_features)
|
|
temporal_feature = self.temporal_attention(stock_sequence)
|
|
temporal_features.append(temporal_feature)
|
|
|
|
# Stack temporal features
|
|
temporal_features = torch.cat(
|
|
temporal_features, dim=0
|
|
) # (num_stocks, feature_dim)
|
|
|
|
# Process features
|
|
processed_features = self.feature_processor(temporal_features)
|
|
|
|
# Process through GNN with gradient checkpointing for memory efficiency
|
|
x = self._gnn_forward(processed_features, edge_index, edge_attr)
|
|
|
|
# Process through LSTM for temporal dependencies
|
|
# Reshape for LSTM: (num_stocks, 1, hidden_size)
|
|
lstm_input = x.unsqueeze(1)
|
|
|
|
# If we have a hidden state from previous prediction, use it
|
|
if self.hidden_state is not None and config.STATEFUL_PREDICTION:
|
|
lstm_out, self.hidden_state = self.lstm(lstm_input, self.hidden_state)
|
|
else:
|
|
lstm_out, self.hidden_state = self.lstm(lstm_input)
|
|
|
|
# Remove sequence dimension
|
|
lstm_out = lstm_out.squeeze(1)
|
|
|
|
# Apply attention to LSTM outputs
|
|
attention_weights = self.attention(lstm_out)
|
|
attended = (lstm_out * attention_weights).sum(dim=1, keepdim=True)
|
|
|
|
# Final prediction
|
|
return self.linear(attended)
|
|
|
|
def _gnn_forward(self, x, edge_index, edge_attr):
|
|
"""Forward pass through GNN with gradient checkpointing for memory efficiency"""
|
|
# Gradient checkpointing for memory efficiency on AMD GPUs
|
|
x = torch.utils.checkpoint.checkpoint(
|
|
self.conv1, x, edge_index, edge_attr, preserve_rng_state=False
|
|
)
|
|
x = F.silu(x) # Swish activation often works better than ReLU on AMD GPUs
|
|
x = F.dropout(x, p=config.DROPOUT, training=self.training)
|
|
|
|
x = torch.utils.checkpoint.checkpoint(
|
|
self.conv2, x, edge_index, edge_attr, preserve_rng_state=False
|
|
)
|
|
|
|
return x
|
|
|
|
def reset_state(self):
|
|
"""Reset the hidden state of the LSTM"""
|
|
self.hidden_state = None
|
|
|
|
def get_attention_weights(self, data):
|
|
"""Get attention weights for interpretability"""
|
|
x, edge_index, edge_attr = data.x, data.edge_index, data.edge_attr
|
|
|
|
# Process through the model up to attention
|
|
batch_size, seq_len, num_features = x.size()
|
|
temporal_features = []
|
|
|
|
for i in range(batch_size):
|
|
stock_sequence = x[i].unsqueeze(0)
|
|
temporal_feature = self.temporal_attention(stock_sequence)
|
|
temporal_features.append(temporal_feature)
|
|
|
|
temporal_features = torch.cat(temporal_features, dim=0)
|
|
processed_features = self.feature_processor(temporal_features)
|
|
|
|
x = self.conv1(processed_features, edge_index, edge_attr)
|
|
x = F.silu(x)
|
|
x = self.conv2(x, edge_index, edge_attr)
|
|
|
|
lstm_input = x.unsqueeze(1)
|
|
lstm_out, _ = self.lstm(lstm_input)
|
|
lstm_out = lstm_out.squeeze(1)
|
|
|
|
attention_weights = self.attention(lstm_out)
|
|
|
|
return attention_weights
|