train.py

"""
Main Training Script for Anomaly Detection Model.

Trains the autoencoder with checkpoint support and all MLOps features.
"""

import argparse
import os
import sys
from pathlib import Path

import torch
import yaml
from loguru import logger

# Add src to path
sys.path.insert(0, str(Path(__file__).parent))

from src.data.data_loader import create_dataloaders
from src.data.data_generator import BearingFaultSimulator
from src.models.autoencoder import VibrationAutoencoder
from src.training.trainer import Trainer, TrainingConfig
from src.utils.config_loader import load_config
from src.utils.logger import setup_logger


def parse_args():
    """Parse command line arguments."""
    parser = argparse.ArgumentParser(
        description="Train anomaly detection model"
    )
    
    parser.add_argument(
        "--config",
        type=str,
        default="config/config.yaml",
        help="Path to configuration file",
    )
    
    parser.add_argument(
        "--resume",
        type=str,
        default=None,
        help="Path to checkpoint to resume from",
    )
    
    parser.add_argument(
        "--qat",
        action="store_true",
        help="Enable quantization-aware training",
    )
    
    parser.add_argument(
        "--pretrained",
        type=str,
        default=None,
        help="Path to pretrained model for fine-tuning",
    )
    
    parser.add_argument(
        "--force-regenerate",
        action="store_true",
        help="Force regenerate cached data",
    )
    
    parser.add_argument(
        "--device",
        type=str,
        default=None,
        help="Device to use (cuda/cpu)",
    )
    
    parser.add_argument(
        "--seed",
        type=int,
        default=None,
        help="Random seed",
    )
    
    return parser.parse_args()


def set_seed(seed: int):
    """Set random seed for reproducibility."""
    import random
    import numpy as np
    
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    
    if torch.cuda.is_available():
        torch.backends.cudnn.deterministic = True
        torch.backends.cudnn.benchmark = False
    
    logger.info(f"Random seed set to {seed}")


def generate_synthetic_data(config: dict):
    """Generate synthetic training data if configured."""
    data_config = config["data"]
    
    if not data_config.get("generate_synthetic", False):
        return
    
    logger.info("Generating synthetic data...")
    
    # Create generator
    simulator = BearingFaultSimulator(
        sampling_rate=data_config["sampling_rate"],
    )
    
    # Generate normal samples
    normal_samples = []
    for i in range(data_config["num_normal_samples"]):
        signal = simulator.generate_normal_signal(
            duration=data_config["signal_duration"],
        )
        normal_samples.append(signal)
        
        if (i + 1) % 1000 == 0:
            logger.info(f"Generated {i + 1}/{data_config['num_normal_samples']} normal samples")
    
    # Generate anomaly samples
    anomaly_samples = []
    anomaly_types = data_config["anomaly_types"]
    samples_per_type = data_config["num_anomaly_samples"] // len(anomaly_types)
    
    for anomaly_type in anomaly_types:
        for i in range(samples_per_type):
            signal = simulator.generate_fault_signal(
                duration=data_config["signal_duration"],
                fault_type=anomaly_type,
            )
            anomaly_samples.append(signal)
    
    logger.info(
        f"Generated {len(normal_samples)} normal and "
        f"{len(anomaly_samples)} anomaly samples"
    )
    
    # Save to disk
    synthetic_dir = Path(data_config["synthetic_dir"])
    synthetic_dir.mkdir(parents=True, exist_ok=True)
    
    np.save(synthetic_dir / "normal_samples.npy", np.array(normal_samples))
    np.save(synthetic_dir / "anomaly_samples.npy", np.array(anomaly_samples))
    
    logger.info(f"Synthetic data saved to {synthetic_dir}")


def main():
    """Main training function."""
    args = parse_args()
    
    # Load configuration
    logger.info(f"Loading configuration from {args.config}")
    config = load_config(args.config)
    
    # Override config with CLI arguments
    if args.device:
        config["hardware"]["device"] = args.device
    if args.seed:
        config["project"]["seed"] = args.seed
    
    # Setup logger
    setup_logger(
        log_dir=config["logging"]["log_dir"],
        level=config["logging"]["level"],
    )
    
    # Set random seed
    set_seed(config["project"]["seed"])
    
    # Device setup
    device = config["hardware"]["device"]
    if device == "cuda" and not torch.cuda.is_available():
        logger.warning("CUDA not available, falling back to CPU")
        device = "cpu"
    
    logger.info(f"Using device: {device}")
    
    # Generate synthetic data if needed
    if config["data"]["generate_synthetic"]:
        synthetic_path = Path(config["data"]["synthetic_dir"]) / "normal_samples.npy"
        if not synthetic_path.exists() or args.force_regenerate:
            generate_synthetic_data(config)
    
    # Create dataloaders with caching
    logger.info("Creating dataloaders...")
    train_loader, val_loader, test_loader, preprocessor = create_dataloaders(
        config=config,
        force_regenerate=args.force_regenerate,
    )
    
    logger.info(
        f"Data loaded - Train: {len(train_loader.dataset)}, "
        f"Val: {len(val_loader.dataset)}, Test: {len(test_loader.dataset)}"
    )
    
    # Create model
    logger.info("Creating model...")
    feature_dim = preprocessor.get_feature_dimension()
    
    model = VibrationAutoencoder(
        input_dim=feature_dim,
        encoder_hidden_dims=config["model"]["encoder"]["hidden_dims"],
        latent_dim=config["model"]["encoder"]["latent_dim"],
        decoder_hidden_dims=config["model"]["decoder"]["hidden_dims"],
        activation=config["model"]["encoder"]["activation"],
        dropout=config["model"]["encoder"]["dropout"],
        batch_norm=config["model"]["encoder"]["batch_norm"],
    )
    
    logger.info(f"Model created - Parameters: {sum(p.numel() for p in model.parameters()):,}")
    
    # Load pretrained weights if specified
    if args.pretrained:
        logger.info(f"Loading pretrained weights from {args.pretrained}")
        checkpoint = torch.load(args.pretrained, map_location=device)
        model.load_state_dict(checkpoint["model_state_dict"])
    
    # Create training config
    training_config = TrainingConfig.from_config(config)
    
    # Enable QAT if requested
    if args.qat:
        logger.info("Quantization-Aware Training enabled")
        training_config.use_qat = True
    
    # Create trainer
    trainer = Trainer(
        model=model,
        config=training_config,
        device=device,
    )
    
    # Resume from checkpoint if specified
    start_epoch = 0
    if args.resume:
        logger.info(f"Resuming from checkpoint: {args.resume}")
        start_epoch = trainer.load_checkpoint(args.resume)
        logger.info(f"Resuming from epoch {start_epoch}")
    
    # Train model
    logger.info("Starting training...")
    trainer.fit(
        train_loader=train_loader,
        val_loader=val_loader,
        start_epoch=start_epoch,
    )
    
    # Evaluate on test set
    logger.info("Evaluating on test set...")
    test_loss = trainer.validate(test_loader)
    logger.info(f"Test Loss: {test_loss:.6f}")
    
    # Save final model
    final_checkpoint_path = os.path.join(
        config["checkpoint"]["dirpath"],
        "final_model.ckpt",
    )
    trainer.save_checkpoint(
        filepath=final_checkpoint_path,
        epoch=training_config.epochs,
    )
    
    logger.info(f"Training completed! Final model saved to {final_checkpoint_path}")


if __name__ == "__main__":
    main()