413 lines
16 KiB
Python
413 lines
16 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.optim as optim
|
|
import numpy as np
|
|
import os
|
|
import logging
|
|
import torch.nn.functional as F
|
|
from typing import List, Tuple, Dict, Any, Optional, Union
|
|
|
|
# Configure logger
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class ResidualBlock(nn.Module):
|
|
"""
|
|
Residual block with pre-activation (BatchNorm -> ReLU -> Conv)
|
|
"""
|
|
def __init__(self, in_channels, out_channels, stride=1):
|
|
super(ResidualBlock, self).__init__()
|
|
self.bn1 = nn.BatchNorm1d(in_channels)
|
|
self.conv1 = nn.Conv1d(in_channels, out_channels, kernel_size=3, stride=stride, padding=1, bias=False)
|
|
self.bn2 = nn.BatchNorm1d(out_channels)
|
|
self.conv2 = nn.Conv1d(out_channels, out_channels, kernel_size=3, stride=1, padding=1, bias=False)
|
|
|
|
# Shortcut connection to match dimensions
|
|
self.shortcut = nn.Sequential()
|
|
if stride != 1 or in_channels != out_channels:
|
|
self.shortcut = nn.Sequential(
|
|
nn.Conv1d(in_channels, out_channels, kernel_size=1, stride=stride, bias=False)
|
|
)
|
|
|
|
def forward(self, x):
|
|
out = F.relu(self.bn1(x))
|
|
shortcut = self.shortcut(out)
|
|
out = self.conv1(out)
|
|
out = self.conv2(F.relu(self.bn2(out)))
|
|
out += shortcut
|
|
return out
|
|
|
|
class SelfAttention(nn.Module):
|
|
"""
|
|
Self-attention mechanism for sequential data
|
|
"""
|
|
def __init__(self, dim):
|
|
super(SelfAttention, self).__init__()
|
|
self.query = nn.Linear(dim, dim)
|
|
self.key = nn.Linear(dim, dim)
|
|
self.value = nn.Linear(dim, dim)
|
|
self.scale = torch.sqrt(torch.tensor(dim, dtype=torch.float32))
|
|
|
|
def forward(self, x):
|
|
# x shape: [batch_size, seq_len, dim]
|
|
batch_size, seq_len, dim = x.size()
|
|
|
|
q = self.query(x) # [batch_size, seq_len, dim]
|
|
k = self.key(x) # [batch_size, seq_len, dim]
|
|
v = self.value(x) # [batch_size, seq_len, dim]
|
|
|
|
# Calculate attention scores
|
|
scores = torch.matmul(q, k.transpose(-2, -1)) / self.scale # [batch_size, seq_len, seq_len]
|
|
|
|
# Apply softmax to get attention weights
|
|
attention = F.softmax(scores, dim=-1) # [batch_size, seq_len, seq_len]
|
|
|
|
# Apply attention to values
|
|
out = torch.matmul(attention, v) # [batch_size, seq_len, dim]
|
|
|
|
return out, attention
|
|
|
|
class EnhancedCNN(nn.Module):
|
|
"""
|
|
Enhanced CNN model with residual connections and attention mechanisms
|
|
for improved trading decision making
|
|
"""
|
|
def __init__(self, input_shape, n_actions, confidence_threshold=0.5):
|
|
super(EnhancedCNN, self).__init__()
|
|
|
|
# Store dimensions
|
|
self.input_shape = input_shape
|
|
self.n_actions = n_actions
|
|
self.confidence_threshold = confidence_threshold
|
|
|
|
# Calculate input dimensions
|
|
if isinstance(input_shape, (list, tuple)):
|
|
if len(input_shape) == 3: # [channels, height, width]
|
|
self.channels, self.height, self.width = input_shape
|
|
self.feature_dim = self.height * self.width
|
|
elif len(input_shape) == 2: # [timeframes, features]
|
|
self.channels = input_shape[0]
|
|
self.features = input_shape[1]
|
|
self.feature_dim = self.features * self.channels
|
|
elif len(input_shape) == 1: # [features]
|
|
self.channels = 1
|
|
self.features = input_shape[0]
|
|
self.feature_dim = self.features
|
|
else:
|
|
raise ValueError(f"Unsupported input shape: {input_shape}")
|
|
else: # single integer
|
|
self.channels = 1
|
|
self.features = input_shape
|
|
self.feature_dim = input_shape
|
|
|
|
# Build network
|
|
self._build_network()
|
|
|
|
# Initialize device
|
|
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
|
self.to(self.device)
|
|
|
|
logger.info(f"EnhancedCNN initialized with input shape: {input_shape}, actions: {n_actions}")
|
|
|
|
def _build_network(self):
|
|
"""Build the enhanced neural network with current feature dimensions"""
|
|
|
|
# 1D CNN for sequential data
|
|
if self.channels > 1:
|
|
# Reshape expected: [batch, timeframes, features]
|
|
self.conv_layers = nn.Sequential(
|
|
nn.Conv1d(self.channels, 64, kernel_size=3, padding=1),
|
|
nn.BatchNorm1d(64),
|
|
nn.ReLU(),
|
|
nn.Dropout(0.2),
|
|
|
|
ResidualBlock(64, 128),
|
|
nn.MaxPool1d(kernel_size=2, stride=2),
|
|
nn.Dropout(0.3),
|
|
|
|
ResidualBlock(128, 256),
|
|
nn.MaxPool1d(kernel_size=2, stride=2),
|
|
nn.Dropout(0.4),
|
|
|
|
ResidualBlock(256, 512),
|
|
nn.AdaptiveAvgPool1d(1) # Global average pooling
|
|
)
|
|
# Feature dimension after conv layers
|
|
self.conv_features = 512
|
|
else:
|
|
# For 1D vectors, skip the convolutional part
|
|
self.conv_layers = None
|
|
self.conv_features = 0
|
|
|
|
# Fully connected layers for all cases
|
|
# We'll use deeper layers with skip connections
|
|
if self.conv_layers is None:
|
|
# For 1D inputs without conv preprocessing
|
|
self.fc1 = nn.Linear(self.feature_dim, 512)
|
|
self.features_dim = 512
|
|
else:
|
|
# For data processed by conv layers
|
|
self.fc1 = nn.Linear(self.conv_features, 512)
|
|
self.features_dim = 512
|
|
|
|
# Common feature extraction layers
|
|
self.fc_layers = nn.Sequential(
|
|
self.fc1,
|
|
nn.ReLU(),
|
|
nn.Dropout(0.4),
|
|
nn.Linear(512, 512),
|
|
nn.ReLU(),
|
|
nn.Dropout(0.4),
|
|
nn.Linear(512, 256),
|
|
nn.ReLU()
|
|
)
|
|
|
|
# Dueling architecture
|
|
self.advantage_stream = nn.Sequential(
|
|
nn.Linear(256, 128),
|
|
nn.ReLU(),
|
|
nn.Linear(128, self.n_actions)
|
|
)
|
|
|
|
self.value_stream = nn.Sequential(
|
|
nn.Linear(256, 128),
|
|
nn.ReLU(),
|
|
nn.Linear(128, 1)
|
|
)
|
|
|
|
# Extrema detection head with increased capacity
|
|
self.extrema_head = nn.Sequential(
|
|
nn.Linear(256, 128),
|
|
nn.ReLU(),
|
|
nn.Dropout(0.3),
|
|
nn.Linear(128, 3) # 0=bottom, 1=top, 2=neither
|
|
)
|
|
|
|
# Price prediction heads with increased capacity
|
|
self.price_pred_immediate = nn.Sequential(
|
|
nn.Linear(256, 64),
|
|
nn.ReLU(),
|
|
nn.Linear(64, 3) # Up, Down, Sideways
|
|
)
|
|
|
|
self.price_pred_midterm = nn.Sequential(
|
|
nn.Linear(256, 64),
|
|
nn.ReLU(),
|
|
nn.Linear(64, 3) # Up, Down, Sideways
|
|
)
|
|
|
|
self.price_pred_longterm = nn.Sequential(
|
|
nn.Linear(256, 64),
|
|
nn.ReLU(),
|
|
nn.Linear(64, 3) # Up, Down, Sideways
|
|
)
|
|
|
|
# Value prediction with increased capacity
|
|
self.price_pred_value = nn.Sequential(
|
|
nn.Linear(256, 128),
|
|
nn.ReLU(),
|
|
nn.Dropout(0.3),
|
|
nn.Linear(128, 4) # % change for different timeframes
|
|
)
|
|
|
|
# Additional attention layer for feature refinement
|
|
self.attention = SelfAttention(256)
|
|
|
|
def _check_rebuild_network(self, features):
|
|
"""Check if network needs to be rebuilt for different feature dimensions"""
|
|
if features != self.feature_dim:
|
|
logger.info(f"Rebuilding network for new feature dimension: {features} (was {self.feature_dim})")
|
|
self.feature_dim = features
|
|
self._build_network()
|
|
# Move to device after rebuilding
|
|
self.to(self.device)
|
|
return True
|
|
return False
|
|
|
|
def forward(self, x):
|
|
"""Forward pass through the network"""
|
|
batch_size = x.size(0)
|
|
|
|
# Process different input shapes
|
|
if len(x.shape) > 2:
|
|
# Handle 3D input [batch, timeframes, features]
|
|
if self.conv_layers is not None:
|
|
# Reshape for 1D convolution:
|
|
# [batch, timeframes, features] -> [batch, timeframes, features*1]
|
|
if len(x.shape) == 3:
|
|
x = x.permute(0, 1, 2) # Ensure shape is [batch, timeframes, features]
|
|
x_reshaped = x.permute(0, 1, 2) # [batch, timeframes, features]
|
|
|
|
# Check if the feature dimension has changed and rebuild if necessary
|
|
if x_reshaped.size(1) * x_reshaped.size(2) != self.feature_dim:
|
|
total_features = x_reshaped.size(1) * x_reshaped.size(2)
|
|
self._check_rebuild_network(total_features)
|
|
|
|
# Apply convolutions
|
|
x_conv = self.conv_layers(x_reshaped)
|
|
# Flatten: [batch, channels, 1] -> [batch, channels]
|
|
x_flat = x_conv.view(batch_size, -1)
|
|
else:
|
|
# If no conv layers, just flatten
|
|
x_flat = x.view(batch_size, -1)
|
|
else:
|
|
# For 2D input [batch, features]
|
|
x_flat = x
|
|
|
|
# Check if dimensions have changed
|
|
if x_flat.size(1) != self.feature_dim:
|
|
self._check_rebuild_network(x_flat.size(1))
|
|
|
|
# Apply FC layers
|
|
features = self.fc_layers(x_flat)
|
|
|
|
# Add attention for feature refinement
|
|
features_3d = features.unsqueeze(1) # [batch, 1, features]
|
|
features_attended, _ = self.attention(features_3d)
|
|
features_refined = features_attended.squeeze(1) # [batch, features]
|
|
|
|
# Calculate advantage and value
|
|
advantage = self.advantage_stream(features_refined)
|
|
value = self.value_stream(features_refined)
|
|
|
|
# Combine for Q-values (Dueling architecture)
|
|
q_values = value + advantage - advantage.mean(dim=1, keepdim=True)
|
|
|
|
# Get extrema predictions
|
|
extrema_pred = self.extrema_head(features_refined)
|
|
|
|
# Price movement predictions
|
|
price_immediate = self.price_pred_immediate(features_refined)
|
|
price_midterm = self.price_pred_midterm(features_refined)
|
|
price_longterm = self.price_pred_longterm(features_refined)
|
|
price_values = self.price_pred_value(features_refined)
|
|
|
|
# Package price predictions
|
|
price_predictions = {
|
|
'immediate': price_immediate,
|
|
'midterm': price_midterm,
|
|
'longterm': price_longterm,
|
|
'values': price_values
|
|
}
|
|
|
|
return q_values, extrema_pred, price_predictions, features_refined
|
|
|
|
def act(self, state, explore=True):
|
|
"""
|
|
Choose action based on state with confidence thresholding
|
|
"""
|
|
state_tensor = torch.FloatTensor(state).unsqueeze(0).to(self.device)
|
|
|
|
with torch.no_grad():
|
|
q_values, _, _, _ = self(state_tensor)
|
|
|
|
# Apply softmax to get action probabilities
|
|
action_probs = F.softmax(q_values, dim=1)
|
|
|
|
# Get action with highest probability
|
|
action = action_probs.argmax(dim=1).item()
|
|
action_confidence = action_probs[0, action].item()
|
|
|
|
# Check if confidence exceeds threshold
|
|
if action_confidence < self.confidence_threshold:
|
|
# Force HOLD action (typically action 2)
|
|
action = 2 # Assume 2 is HOLD
|
|
logger.info(f"Action {action} confidence {action_confidence:.4f} below threshold {self.confidence_threshold}, forcing HOLD")
|
|
|
|
return action, action_confidence
|
|
|
|
def save(self, path):
|
|
"""Save model weights and architecture"""
|
|
os.makedirs(os.path.dirname(path), exist_ok=True)
|
|
torch.save({
|
|
'state_dict': self.state_dict(),
|
|
'input_shape': self.input_shape,
|
|
'n_actions': self.n_actions,
|
|
'feature_dim': self.feature_dim,
|
|
'confidence_threshold': self.confidence_threshold
|
|
}, f"{path}.pt")
|
|
logger.info(f"Enhanced CNN model saved to {path}.pt")
|
|
|
|
def load(self, path):
|
|
"""Load model weights and architecture"""
|
|
try:
|
|
checkpoint = torch.load(f"{path}.pt", map_location=self.device)
|
|
self.input_shape = checkpoint['input_shape']
|
|
self.n_actions = checkpoint['n_actions']
|
|
self.feature_dim = checkpoint['feature_dim']
|
|
if 'confidence_threshold' in checkpoint:
|
|
self.confidence_threshold = checkpoint['confidence_threshold']
|
|
self._build_network()
|
|
self.load_state_dict(checkpoint['state_dict'])
|
|
self.to(self.device)
|
|
logger.info(f"Enhanced CNN model loaded from {path}.pt")
|
|
return True
|
|
except Exception as e:
|
|
logger.error(f"Error loading model: {str(e)}")
|
|
return False
|
|
|
|
# Additional utility for example sifting
|
|
class ExampleSiftingDataset:
|
|
"""
|
|
Dataset that selectively keeps high-quality examples for training
|
|
to improve model performance
|
|
"""
|
|
def __init__(self, max_examples=50000):
|
|
self.examples = []
|
|
self.labels = []
|
|
self.rewards = []
|
|
self.max_examples = max_examples
|
|
self.min_reward_threshold = -0.05 # Minimum reward to keep an example
|
|
|
|
def add_example(self, state, action, reward, next_state, done):
|
|
"""Add a new training example with reward-based filtering"""
|
|
# Only keep examples with rewards above the threshold
|
|
if reward > self.min_reward_threshold:
|
|
self.examples.append((state, action, reward, next_state, done))
|
|
self.rewards.append(reward)
|
|
|
|
# Sort by reward and keep only the top examples
|
|
if len(self.examples) > self.max_examples:
|
|
# Sort by reward (highest first)
|
|
sorted_indices = np.argsort(self.rewards)[::-1]
|
|
# Keep top examples
|
|
self.examples = [self.examples[i] for i in sorted_indices[:self.max_examples]]
|
|
self.rewards = [self.rewards[i] for i in sorted_indices[:self.max_examples]]
|
|
|
|
# Update the minimum reward threshold to be the minimum in our kept examples
|
|
self.min_reward_threshold = min(self.rewards)
|
|
|
|
def get_batch(self, batch_size):
|
|
"""Get a batch of examples, prioritizing better examples"""
|
|
if not self.examples:
|
|
return None
|
|
|
|
# Calculate selection probabilities based on rewards
|
|
rewards = np.array(self.rewards)
|
|
# Shift rewards to be positive for probability calculation
|
|
min_reward = min(rewards)
|
|
shifted_rewards = rewards - min_reward + 0.1 # Add small constant
|
|
probs = shifted_rewards / shifted_rewards.sum()
|
|
|
|
# Sample batch indices with reward-based probabilities
|
|
indices = np.random.choice(
|
|
len(self.examples),
|
|
size=min(batch_size, len(self.examples)),
|
|
p=probs,
|
|
replace=False
|
|
)
|
|
|
|
# Create batch
|
|
batch = [self.examples[i] for i in indices]
|
|
states, actions, rewards, next_states, dones = zip(*batch)
|
|
|
|
return {
|
|
'states': np.array(states),
|
|
'actions': np.array(actions),
|
|
'rewards': np.array(rewards),
|
|
'next_states': np.array(next_states),
|
|
'dones': np.array(dones)
|
|
}
|
|
|
|
def __len__(self):
|
|
return len(self.examples) |