misc

NN/models/enhanced_cnn.py

@@ -0,0 +1,413 @@
+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)