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

Add comprehensive project documentation and fix data pipeline

Browse Source 
- Add detailed README with architecture diagram and usage instructions
- Add API, configuration, and development documentation
- Fix price data column handling for yfinance auto_adjust=True
- Fix model feature dimension indexing and temporal attention batching
- Add missing imports and position tracking in paper broker
- Add python-dotenv support for environment variables
- Update .gitignore with Python artifacts and environment files
This commit is contained in:
Florian Egger
2026-05-26 14:10:48 +02:00
parent 4bf7394a0a
commit 0cf37e786a
17 changed files with 1536 additions and 277 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+541
View File
@@ -0,0 +1,541 @@
# StockGNN R9700 — AMD-Optimized Graph Neural Network Trading System
A real-time intraday stock trading system built around a **Graph Neural Network (GNN)** with **Temporal Attention**, optimized for the **AMD Radeon R9700 AI Pro** GPU (32GB). The system ingests alternative data (news sentiment, social media), price bars, and corporate actions to generate trade signals across a universe of stocks.
---
## 📋 Table of Contents
- [Overview](#overview)
- [Architecture](#architecture)
- [Key Features](#key-features)
- [Project Structure](#project-structure)
- [Quick Start](#quick-start)
- [Configuration](#configuration)
- [Training](#training)
- [Live Trading](#live-trading)
- [Web Dashboard](#web-dashboard)
- [API Reference](#api-reference)
- [AMD Optimizations](#amd-optimizations)
- [Memory Management](#memory-management)
- [Data Pipeline](#data-pipeline)
- [Trading System](#trading-system)
- [Benchmarking](#benchmarking)
- [Development](#development)
- [Troubleshooting](#troubleshooting)
- [License](#license)
---
## Overview
StockGNN uses a **Graph Attention Network (GAT)** with:
- **Temporal Attention** over multi-source feature sequences (price, news, social media)
- **Corporate Action Awareness** for splits, dividends, and delistings
- **Intraday GNN** variant with stateful LSTM for real-time inference
- **Online Learning** to adapt to regime changes during market hours
- **AMD-specific ROCm optimizations** (bf16 mixed precision, gradient checkpointing, pinned memory)
---
## Architecture
```
┌─────────────────────────────────────────────────────────────────┐
│ Data Sources │
│ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────────┐ │
│ │ Polygon │ │ yfinance │ │ NewsAPI │ │ Twitter/ │ │
│ │ WebSocket│ │ Historical│ │ Sentiment│ │ Reddit │ │
│ └────┬─────┘ └────┬─────┘ └────┬─────┘ └──────┬───────┘ │
└───────┼────────────┼─────────────┼───────────────┼─────────────┘
│ │ │ │
▼ ▼ ▼ ▼
┌─────────────────────────────────────────────────────────────────┐
│ Data Pipeline │
│ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────────┐ │
│ │ Price │ │ Corporate│ │ Sector │ │ Alternative │ │
│ │ Data │ │ Actions │ │ Data │ │ Data │ │
│ └────┬─────┘ └────┬─────┘ └────┬─────┘ └──────┬───────┘ │
│ │ │ │ │ │
│ ▼ ▼ ▼ ▼ │
│ ┌───────────────────────────────────────────────────────┐ │
│ │ SQLite Database (WAL mode) │ │
│ │ price_data | corporate_actions | features | trades │ │
│ └────────────────────────┬──────────────────────────────┘ │
└─────────────────────────┼──────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ GNN Models │
│ ┌────────────────────┐ ┌────────────────────┐ │
│ │ CorporateAction- │ │ IntradayGNN │ │
│ │ AwareGNN │ │ (stateful LSTM) │ │
│ │ │ │ │ │
│ │ TemporalAttention │ │ TemporalAttention │ │
│ │ GAT Conv Layers │ │ GAT Conv Layers │ │
│ │ LSTM │ │ LSTM │ │
│ │ Alternative Data │ │ Real-time Features │ │
│ │ Fusion │ │ │ │
│ └────────────────────┘ └────────────────────┘ │
└─────────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ Trading Execution │
│ ┌──────────────┐ ┌──────────────┐ ┌──────────────┐ │
│ │ Paper Broker │ │ IB Broker │ │ Risk Manager │ │
│ │ (Simulation) │ │ (Interactive │ │ (Position │ │
│ │ │ │ Brokers) │ │ Sizing) │ │
│ └──────────────┘ └──────────────┘ └──────────────┘ │
└─────────────────────────────────────────────────────────────────┘
│
▼
┌─────────────────────────────────────────────────────────────────┐
│ Web Dashboard │
│ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ │
│ │ Metrics │ │ Charts │ │ Trading │ │ Logs │ │
│ │ (Live) │ │ (Memory, │ │ Controls │ │ (System) │ │
│ │ │ │ Portfolio)│ │ │ │ │ │
│ └──────────┘ └──────────┘ └──────────┘ └──────────┘ │
└─────────────────────────────────────────────────────────────────┘
```
---
## Key Features
| Feature | Description |
|---------|-------------|
| **AMD Optimizations** | ROCm bf16 mixed precision, gradient checkpointing, Flash Attention, memory-aware batching |
| **GNN Architecture** | Graph Attention Networks with temporal sequence processing and corporate action awareness |
| **Live Data** | Polygon.io WebSocket for real-time trades/quotes with automatic reconnection |
| **Alternative Data** | News sentiment (FinBERT) and social media (Twitter/Reddit) integrated into node features |
| **Online Learning** | Periodic model updates during market hours to adapt to regime shifts |
| **Risk Management** | Max drawdown, daily loss limits, volatility-target position sizing |
| **Web Dashboard** | Single-page application with real-time WebSocket metrics and Chart.js visualizations |
| **Backtesting** | Full backtesting framework with benchmark comparison and performance metrics |
---
## Project Structure
```
stock_gnn_r9700/
├── config.py # Central configuration
├── requirements.txt # Python dependencies
├── main.py # Main training & backtesting script
├── live_trading.py # Live intraday trading system
├── benchmark.py # Performance benchmarking
├── README.md # This file
│
├── data/
│ ├── raw/ # Raw downloaded data
│ ├── processed/ # Processed datasets
│ └── external/ # External datasets (delisted stocks)
│
├── models/ # Saved model weights
│
├── src/
│ ├── amd/
│ │ └── optimizations.py # AMDOptimizer, AMDSparseAttention, AMDGATConv
│ ├── data/
│ │ ├── pipeline.py # StockDataPipeline (SQLite, yfinance, features)
│ │ ├── live_data.py # LiveDataService (Polygon WebSocket)
│ │ ├── news_processor.py # News data processor
│ │ ├── social_processor.py # Social media processor
│ │ └── sentiment.py # Sentiment analyzer
│ ├── models/
│ │ ├── gnn_model.py # CorporateActionAwareGNN + TemporalAttention
│ │ ├── intraday_gnn.py # IntradayGNN (stateful LSTM variant)
│ │ └── trainer.py # GNNTrainer with AMP & gradient clipping
│ ├── trading/
│ │ ├── broker.py # Abstract broker interface
│ │ ├── paper_broker.py # Paper trading simulation
│ │ ├── ib_broker.py # Interactive Brokers implementation
│ │ └── real_time_trader.py # Real-time trading execution
│ ├── evaluation/
│ │ ├── backtester.py # GNNBacktester
│ │ ├── intraday_backtester.py
│ │ └── metrics.py # Sharpe, Sortino, Calmar, IR, etc.
│ ├── utils/
│ │ ├── memory_manager.py # MemoryManager (AMD GPU memory tracking)
│ │ ├── helpers.py # generate_intraday_timestamps
│ │ ├── visualization.py # Plotting utilities
│ │ └── text_processing.py # Text cleaning
│ └── web/
│ ├── app.py # FastAPI entrypoint
│ ├── services/
│ │ └── state.py # Shared application state
│ ├── api/
│ │ ├── dashboard.py # Dashboard API endpoints
│ │ ├── trading_endpoints.py
│ │ ├── models_endpoints.py
│ │ └── data_endpoints.py
│ ├── templates/
│ │ └── index.html # Main dashboard template
│ └── static/
│ ├── css/style.css # Dark-themed responsive UI
│ └── js/app.js # Chart.js, WebSocket, API client
│
└── notebooks/ # Jupyter notebooks for exploration
```
---
## Quick Start
### Prerequisites
- Python 3.10+
- AMD GPU with ROCm 5.6+ (or NVIDIA GPU with CUDA)
- 32GB+ GPU memory recommended (configurable)
### Installation
```bash
# Clone repository
git clone <repository-url>
cd stock_gnn_r9700
# Create virtual environment
python -m venv venv
source venv/bin/activate # Windows: venv\Scripts\activate
# Install dependencies
pip install -r requirements.txt
# For AMD ROCm specifically:
# pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/rocm5.6
# pip install torch-geometric torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-2.1.0+rocm5.6.html
```
### Run Training
```bash
python main.py
```
This will:
1. Initialize the data pipeline and download historical data
2. Create training/validation datasets
3. Train the GNN model with AMD optimizations
4. Run backtesting on validation data
5. Generate performance plots and metrics
### Run Live Trading (Paper)
```bash
python live_trading.py
```
### Run Web Dashboard
```bash
python -m uvicorn src.web.app:app --host 0.0.0.0 --port 8000
```
Then open `http://localhost:8000` in your browser.
---
## Configuration
All settings are centralized in `config.py`. Key parameters:
```python
# GPU Settings
DEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'
AMD_GPU = True
GPU_MEMORY_LIMIT = 0.9 # Use 90% of 32GB = ~28.8GB
MIXED_PRECISION = True
PRECISION = 'bf16' # 'fp16' or 'bf16'
ROCM_OPT_LEVEL = 'O2'
# Model Architecture
HIDDEN_CHANNELS = 128
NUM_HEADS = 16
DROPOUT = 0.3
LEARNING_RATE = 0.0005
BATCH_SIZE = 128
SEQUENCE_LENGTH = 60
# Intraday Trading
TRADING_FREQUENCY = '5min' # '1min', '5min', '15min', '30min', '1h'
INITIAL_CAPITAL = 100000
MAX_POSITION_SIZE = 0.03 # 3% of portfolio per position
# Live Data
DATA_PROVIDER = 'polygon' # 'polygon', 'alphavantage', 'ib'
POLYGON_API_KEY = 'your_key'
```
---
## Training
The training loop (`src/models/trainer.py`) supports:
- **Mixed Precision Training** (bf16/fp16) via `torch.cuda.amp`
- **Gradient Checkpointing** for memory-efficient large models
- **AdamW Optimizer** with weight decay
- **ReduceLROnPlateau** scheduler
- **Automatic batch skipping** when GPU memory is constrained
- **Per-epoch benchmarking** and checkpoint saving
Training is triggered via the Web Dashboard or `main.py`.
---
## Live Trading
The live trading system (`live_trading.py`) runs an async event loop with:
1. **LiveDataService** — connects to Polygon.io WebSocket for real-time market data
2. **RealTimeTrader** — generates signals on configurable intervals (1min–1h)
3. **PaperTradingBroker** — simulates execution (or swap for InteractiveBrokersBroker)
4. **Online Learning** — periodic model updates every hour
5. **Memory Monitor** — automatic GPU cache clearing when usage exceeds 85%
---
## Web Dashboard
The dashboard provides real-time monitoring and control:
| Feature | Description |
|---------|-------------|
| **Live Metrics** | GPU memory %, portfolio value, cash, model status via WebSocket |
| **Charts** | Memory usage and portfolio value Chart.js graphs (60-point history) |
| **Trading Controls** | Start/stop trading, view positions/orders, close positions |
| **Model Controls** | Train, save, load, benchmark models |
| **Data Controls** | View tickers, trigger data updates |
| **System Logs** | Real-time log streaming with clear button |
---
## API Reference
### Dashboard
| Method | Endpoint | Description |
|--------|----------|-------------|
| GET | `/api/dashboard/metrics` | Current metrics (memory, account, model) |
| GET | `/api/dashboard/logs?limit=100` | Recent log entries |
| POST | `/api/dashboard/logs/clear` | Clear stored logs |
### Trading
| Method | Endpoint | Description |
|--------|----------|-------------|
| GET | `/api/trading/status` | Trading status, cash, positions |
| POST | `/api/trading/start` | Start live trading |
| POST | `/api/trading/stop` | Stop live trading |
| GET | `/api/trading/orders` | List all orders |
| POST | `/api/trading/order` | Submit manual order |
| POST | `/api/trading/cancel/{id}` | Cancel order |
| POST | `/api/trading/positions/close/{ticker}` | Close position |
| POST | `/api/trading/positions/close-all` | Close all positions |
### Models
| Method | Endpoint | Description |
|--------|----------|-------------|
| GET | `/api/models/status` | Model configuration status |
| POST | `/api/models/train` | Start training run |
| POST | `/api/models/train/stop` | Stop training |
| POST | `/api/models/save` | Save model weights |
| POST | `/api/models/load` | Load model weights |
| POST | `/api/models/benchmark` | Run performance benchmark |
### Data
| Method | Endpoint | Description |
|--------|----------|-------------|
| GET | `/api/data/tickers` | List tracked tickers |
| GET | `/api/data/pipeline/status` | Pipeline status |
| POST | `/api/data/update` | Trigger data update |
| GET | `/api/data/features/{ticker}` | Latest features for ticker |
| GET | `/api/data/price/{ticker}` | Latest price for ticker |
| GET | `/api/data/prices/{ticker}?limit=30` | Historical prices |
### WebSocket
| Event | Direction | Description |
|-------|-----------|-------------|
| `metrics` | Server → Client | Real-time metrics broadcast (every 2s) |
| `ping` | Client → Server | Keep-alive / latency check |
---
## AMD Optimizations
`src/amd/optimizations.py` provides:
- **AMDOptimizer**: Configures ROCm, enables Flash Attention, sets memory limits
- **AMDSparseAttention**: Custom attention implementation optimized for AMD GPUs
- **AMDGATConv**: GAT layer using `torch_geometric.nn.MessagePassing` with `softmax`
Key optimizations applied:
| Technique | Benefit |
|-----------|---------|
| bf16 Mixed Precision | 2x memory reduction, faster compute on MI200/RX7000 |
| Gradient Checkpointing | Train larger models with same GPU memory |
| Flash Attention | O(N) memory vs O(N²) for long sequences |
| SiLU Activation | Better AMD GPU utilization vs ReLU |
| Pinned Memory | Faster CPU→GPU data transfers |
---
## Memory Management
`src/utils/memory_manager.py` provides:
- **GPU Memory Limiting** — caps usage to 90% of total memory
- **Automatic Cache Clearing** — triggers at configurable thresholds (default 85%)
- **Per-Operation Memory Checks** — skips batches if insufficient memory
- **Background Monitoring** — logs usage every 60 seconds
- **Model Memory Estimation** — estimates activation memory before training
---
## Data Pipeline
`src/data/pipeline.py` handles:
1. **SQLite Database** with WAL mode for concurrent access
2. **Price Data** from yfinance (daily OHLCV, auto-adjusted)
3. **Corporate Actions** (splits, dividends, delistings)
4. **Sector/Industry** classifications
5. **Alternative Data** (news sentiment, social media volume)
6. **Graph Construction** — edges based on sector correlation
The pipeline supports incremental updates: only fetches new data since last run.
---
## Trading System
### Signal Generation
The trading system generates signals when:
- Market is open (9:30 AM – 4:00 PM ET)
- Interval matches `TRADING_FREQUENCY` (e.g., every 5 minutes)
- Model prediction > +0.2% → BUY
- Model prediction < -0.2% and position exists → SELL
### Risk Management
| Rule | Default |
|------|---------|
| Max Position Size | 3% of portfolio |
| Max Daily Loss | 1% |
| Max Drawdown | 5% |
| Max Daily Positions | 50 |
| Max Hold Time | 4 hours |
| Min Hold Time | 10 minutes |
### Execution
- **VWAP** execution algorithm (configurable)
- **Paper Broker** for simulation
- **Interactive Brokers** adapter for live trading
---
## Benchmarking
`benchmark.py` measures:
| Metric | Description |
|--------|-------------|
| Inference Time | Average forward pass latency |
| Training Time | Average full training step latency |
| Throughput | Samples/second for inference and training |
| Memory Usage | Peak GPU memory during benchmark |
Run with:
```bash
python benchmark.py
```
Results are logged to `benchmark_r9700.log`.
---
## Development
### Adding a New Data Source
1. Create a processor in `src/data/` (e.g., `src/data/earnings_processor.py`)
2. Implement `fetch_earnings(tickers, start, end)` and `get_earnings_features(ticker, date)`
3. Register in `StockDataPipeline`
4. Update `config.py` with API keys and feature definitions
### Adding a New Model Variant
1. Inherit from `nn.Module` in `src/models/`
2. Use `torch.utils.checkpoint.checkpoint` for memory efficiency
3. Wrap with `AMDOptimizer.optimize_model()` before training
4. Register in `GNNTrainer`
### Adding a New Trading Strategy
1. Subclass `Broker` for execution
2. Implement `_generate_trading_signals()` in `RealTimeTrader`
3. Configure thresholds in `config.py`
---
## Troubleshooting
### GPU Out of Memory
- Reduce `BATCH_SIZE` in `config.py`
- Reduce `SEQUENCE_LENGTH` or `HIDDEN_CHANNELS`
- Enable gradient checkpointing (`MIXED_PRECISION = True`)
- Lower `GPU_MEMORY_LIMIT` to trigger earlier cache clearing
### ROCm Installation Issues
```bash
# Verify ROCm is installed
rocminfo
# Check PyTorch sees the GPU
python -c "import torch; print(torch.cuda.is_available())"
```
### Data Pipeline Empty
- Ensure `yfinance` can connect to Yahoo Finance
- Check `data/processed/` directory permissions
- Verify ticker symbols are valid
### WebSocket Disconnects
- Check `WEBSOCKET_MAX_RETRIES` and `WEBSOCKET_RECONNECT_DELAY` in config
- Verify Polygon.io API key is valid
- Check firewall rules for WebSocket connections
### Model Not Loading
- Verify `models/stock_gnn_r9700.pt` exists
- Check `config.py` `MODEL_DIR` path
- Ensure `DEVICE` matches training device (CPU vs CUDA)
---
## License
MIT License — see [LICENSE](LICENSE) for details.
---
## Acknowledgments
- PyTorch Geometric for GNN implementations
- yfinance for historical market data
- Polygon.io for real-time market data
- Chart.js for dashboard visualizations