main.py

"""
main.py
Main entry point for training DREAMS
"""

import torch
import argparse
import json
import os
from torch.utils.data import DataLoader, random_split

# Import modules
from model import MultiModalTransformerModel
from train import Trainer, MultiModalDataset


def load_config(config_path: str) -> dict:
    """Load configuration from JSON file"""
    with open(config_path, 'r') as f:
        config = json.load(f)
    return config


def setup_data_loaders(config: dict):
    """Setup train and validation data loaders"""
    # Create dataset
    dataset = MultiModalDataset(
        data_path=config.get('data_path'),
        num_samples=config.get('num_samples', 10000)
    )
    
    # Split into train and validation
    val_split = config.get('val_split', 0.1)
    val_size = int(len(dataset) * val_split)
    train_size = len(dataset) - val_size
    
    train_dataset, val_dataset = random_split(
        dataset, [train_size, val_size],
        generator=torch.Generator().manual_seed(config.get('seed', 42))
    )
    
    # Create data loaders
    train_loader = DataLoader(
        train_dataset,
        batch_size=config['batch_size'],
        shuffle=True,
        num_workers=config.get('num_workers', 4),
        pin_memory=True
    )
    
    val_loader = DataLoader(
        val_dataset,
        batch_size=config['batch_size'],
        shuffle=False,
        num_workers=config.get('num_workers', 4),
        pin_memory=True
    )
    
    return train_loader, val_loader


def main(args):
    """Main training function"""
    
    # Load configuration
    if args.config:
        config = load_config(args.config)
    else:
        # Default configuration
        config = {
            # Data dimensions
            'input_dim_o': 128,
            'input_dim_p': 64,
            'input_dim_t': 32,
            'input_dim_x': 100,
            
            # Model architecture
            'hidden_dim': 256,
            'z_dim': 128,
            'memory_dim': 128,
            'd_model': 512,
            'nhead': 8,
            'num_layers': 6,
            'dim_feedforward': 2048,
            'memory_size': 1000,
            'num_platforms': 10,
            'num_flows': 4,
            
            # Output dimensions
            'output_dim_y': 1,
            'output_dim_n': 5,
            
            # Training parameters
            'batch_size': 32,
            'num_epochs': 100,
            'learning_rate': 1e-4,
            'lr_encoder': 5e-5,
            'lr_backbone': 5e-5,
            'weight_decay': 1e-5,
            'max_grad_norm': 1.0,
            
            # Loss weights
            'lambda_dis': 0.1,
            'lambda_rep': 0.05,
            'beta_kl': 1.0,
            
            # Optimizer and scheduler
            'optimizer': 'adamw',
            'scheduler': 'cosine',
            'min_lr': 1e-6,
            
            # Replay buffer
            'buffer_type': 'standard',
            'buffer_capacity': 10000,
            
            # DREAMS augmentation
            'augmentation_config': {
                'noise_scale': 0.01,
                'mixup_alpha': 0.2,
                'cutmix_prob': 0.5,
                'temporal_shift_range': 2,
                'feature_dropout': 0.1
            },
            
            # Other settings
            'val_split': 0.1,
            'num_workers': 4,
            'log_interval': 10,
            'save_interval': 10,
            'seed': 42,
            'log_dir': 'runs/experiment',
            
            # Model-specific settings
            'use_learned_weight': True,
            'use_film_embedding': True,
            'dropout': 0.1,
            'lstm_layers': 2,
            'memory_heads': 8,
            'head_hidden_dims': [256, 128]
        }
    
    # Override config with command line arguments
    if args.batch_size:
        config['batch_size'] = args.batch_size
    if args.epochs:
        config['num_epochs'] = args.epochs
    if args.lr:
        config['learning_rate'] = args.lr
    if args.device:
        device = args.device
    else:
        device = 'cuda' if torch.cuda.is_available() else 'cpu'
    
    # Set random seed
    torch.manual_seed(config['seed'])
    if torch.cuda.is_available():
        torch.cuda.manual_seed(config['seed'])
    
    # Create output directory
    os.makedirs(args.output_dir, exist_ok=True)
    
    # Save configuration
    config_save_path = os.path.join(args.output_dir, 'config.json')
    with open(config_save_path, 'w') as f:
        json.dump(config, f, indent=2)
    print(f"Configuration saved to {config_save_path}")
    
    # Setup data loaders
    print("Setting up data loaders...")
    train_loader, val_loader = setup_data_loaders(config)
    print(f"Train samples: {len(train_loader.dataset)}")
    print(f"Validation samples: {len(val_loader.dataset)}")
    
    # Initialize model
    print("Initializing model...")
    model = MultiModalTransformerModel(config)
    print(f"Model initialized with {sum(p.numel() for p in model.parameters()):,} parameters")
    
    # Initialize trainer
    trainer = Trainer(model, config, device=device)
    
    # Load checkpoint if specified
    if args.resume:
        print(f"Loading checkpoint from {args.resume}")
        start_epoch, _ = trainer.load_checkpoint(args.resume)
        print(f"Resumed from epoch {start_epoch}")
    
    # Train model
    print("Starting training...")
    trained_model = trainer.train(train_loader, val_loader)
    
    # Save final model
    final_model_path = os.path.join(args.output_dir, 'final_model.pth')
    torch.save(trained_model.state_dict(), final_model_path)
    print(f"Final model saved to {final_model_path}")
    
    print("Training completed successfully!")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Train Multi-Modal Transformer Model')
    
    # Configuration
    parser.add_argument('--config', type=str, default=None,
                       help='Path to configuration JSON file')
    
    # Training parameters
    parser.add_argument('--batch-size', type=int, default=None,
                       help='Batch size for training')
    parser.add_argument('--epochs', type=int, default=None,
                       help='Number of training epochs')
    parser.add_argument('--lr', type=float, default=None,
                       help='Learning rate')
    
    # Device
    parser.add_argument('--device', type=str, default=None,
                       choices=['cuda', 'cpu'],
                       help='Device to use for training')
    
    # Checkpointing
    parser.add_argument('--resume', type=str, default=None,
                       help='Path to checkpoint to resume training')
    parser.add_argument('--output-dir', type=str, default='output',
                       help='Output directory for saving models and logs')
    
    args = parser.parse_args()
    
    main(args)