train_deep.py

"""
10-hour deep learning training with PyTorch.

Trains a neural network with many epochs and early stopping.
"""

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from algo_trader.research.dataset import build_ml_table, load_features_labels
from algo_trader.common.logger import setup_logger, logger
from algo_trader.common.config import settings
from datetime import datetime
import numpy as np
import pandas as pd
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import json

setup_logger(level="INFO", json=False)

print("=" * 60)
print("🧠 10-HOUR DEEP LEARNING TRAINING")
print("=" * 60)
print(f"Start time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
print("=" * 60)
print()


# Define neural network
class DeepTrader(nn.Module):
    """Deep neural network for crypto trading."""
    
    def __init__(self, input_size: int):
        super().__init__()
        self.network = nn.Sequential(
            nn.Linear(input_size, 256),
            nn.BatchNorm1d(256),
            nn.ReLU(),
            nn.Dropout(0.3),
            
            nn.Linear(256, 128),
            nn.BatchNorm1d(128),
            nn.ReLU(),
            nn.Dropout(0.3),
            
            nn.Linear(128, 64),
            nn.BatchNorm1d(64),
            nn.ReLU(),
            nn.Dropout(0.2),
            
            nn.Linear(64, 32),
            nn.BatchNorm1d(32),
            nn.ReLU(),
            nn.Dropout(0.2),
            
            nn.Linear(32, 1),
            nn.Sigmoid()
        )
    
    def forward(self, x):
        return self.network(x)


# Load data
print("Loading data...")
features, labels = load_features_labels(timeframe="1h")
X, y, meta = build_ml_table(features, labels, target="direction_4")

print(f"Dataset: {len(X):,} samples, {len(X.columns)} features")
print()

# Scale features
scaler = StandardScaler()
X_scaled = scaler.fit_transform(X.values)

# Train/val split
X_train, X_val, y_train, y_val = train_test_split(
    X_scaled, y.values, test_size=0.2, random_state=42, stratify=y
)

print(f"Training: {len(X_train):,} samples")
print(f"Validation: {len(X_val):,} samples")
print()

# Convert to PyTorch tensors
X_train_t = torch.FloatTensor(X_train)
y_train_t = torch.FloatTensor(y_train).reshape(-1, 1)
X_val_t = torch.FloatTensor(X_val)
y_val_t = torch.FloatTensor(y_val).reshape(-1, 1)

# Create data loaders
train_dataset = TensorDataset(X_train_t, y_train_t)
val_dataset = TensorDataset(X_val_t, y_val_t)

use_cuda = torch.cuda.is_available()
train_loader = DataLoader(train_dataset, batch_size=64, shuffle=True, pin_memory=use_cuda)
val_loader = DataLoader(val_dataset, batch_size=256, pin_memory=use_cuda)

# Initialize model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print(f"Using device: {device}")
print()

model = DeepTrader(input_size=X_train.shape[1]).to(device)
criterion = nn.BCELoss()
optimizer = optim.Adam(model.parameters(), lr=0.001, weight_decay=1e-5)
scheduler = optim.lr_scheduler.ReduceLROnPlateau(
    optimizer, mode='max', factor=0.5, patience=20
)

# Training configuration
max_epochs = 10000  # Will run for hours with early stopping
patience = 100
best_val_auc = 0
patience_counter = 0

print("🚀 Starting training...")
print(f"Max epochs: {max_epochs:,}")
print(f"Early stopping patience: {patience}")
print()

start_time = datetime.now()
history = []

for epoch in range(max_epochs):
    # Training phase
    model.train()
    train_loss = 0
    for batch_X, batch_y in train_loader:
        batch_X, batch_y = batch_X.to(device), batch_y.to(device)
        
        optimizer.zero_grad()
        outputs = model(batch_X)
        loss = criterion(outputs, batch_y)
        loss.backward()
        optimizer.step()
        
        train_loss += loss.item()
    
    train_loss /= len(train_loader)
    
    # Validation phase
    model.eval()
    val_loss = 0
    val_preds = []
    val_targets = []
    
    with torch.no_grad():
        for batch_X, batch_y in val_loader:
            batch_X, batch_y = batch_X.to(device), batch_y.to(device)
            outputs = model(batch_X)
            loss = criterion(outputs, batch_y)
            val_loss += loss.item()
            
            val_preds.extend(outputs.cpu().numpy())
            val_targets.extend(batch_y.cpu().numpy())
    
    val_loss /= len(val_loader)
    
    # Calculate AUC
    from sklearn.metrics import roc_auc_score
    val_auc = roc_auc_score(val_targets, val_preds)
    
    # Update learning rate
    scheduler.step(val_auc)
    
    # Log progress
    if (epoch + 1) % 10 == 0:
        elapsed = (datetime.now() - start_time).total_seconds() / 3600
        print(f"Epoch {epoch+1:4d}/{max_epochs} | "
              f"Train Loss: {train_loss:.4f} | "
              f"Val Loss: {val_loss:.4f} | "
              f"Val AUC: {val_auc:.4f} | "
              f"Time: {elapsed:.2f}h")
    
    # Save history
    history.append({
        'epoch': epoch + 1,
        'train_loss': train_loss,
        'val_loss': val_loss,
        'val_auc': val_auc,
        'lr': optimizer.param_groups[0]['lr'],
    })
    
    # Early stopping
    if val_auc > best_val_auc:
        best_val_auc = val_auc
        patience_counter = 0
        
        # Save best model
        torch.save({
            'epoch': epoch,
            'model_state_dict': model.state_dict(),
            'optimizer_state_dict': optimizer.state_dict(),
            'best_val_auc': best_val_auc,
            'scaler': scaler,
        }, settings.artifacts_dir / 'best_deep_model.pt')
    else:
        patience_counter += 1
        if patience_counter >= patience:
            print(f"\n⚠️ Early stopping triggered at epoch {epoch+1}")
            break

end_time = datetime.now()
duration = (end_time - start_time).total_seconds() / 3600

print()
print("=" * 60)
print("✅ TRAINING COMPLETED!")
print("=" * 60)
print(f"Duration: {duration:.2f} hours")
print(f"Total epochs: {len(history)}")
print(f"Best validation AUC: {best_val_auc:.4f}")
print()

# Save history
history_df = pd.DataFrame(history)
history_path = settings.artifacts_dir / 'deep_training_history.csv'
history_df.to_csv(history_path, index=False)
print(f"Training history saved to: {history_path}")

# Save config
config = {
    'model_type': 'DeepTrader',
    'architecture': '256-128-64-32-1',
    'input_features': len(X.columns),
    'total_epochs': len(history),
    'best_val_auc': float(best_val_auc),
    'duration_hours': duration,
    'feature_names': list(X.columns),
}

config_path = settings.artifacts_dir / 'deep_model_config.json'
with open(config_path, 'w') as f:
    json.dump(config, f, indent=2)

print(f"Model config saved to: {config_path}")
print(f"Best model saved to: {settings.artifacts_dir / 'best_deep_model.pt'}")
print("=" * 60)