rustytorch/crates/training/rtx-transformers/examples/moco_v3_example.rs

//! MoCo v3 (Momentum Contrast v3) Example
//!
//! Demonstrates how to use MoCo v3 for self-supervised visual representation learning.
//! This example shows the complete training pipeline including:
//! - Configuration setup
//! - Model creation
//! - Training loop
//! - Feature extraction for downstream tasks

use rtx_transformers::prelude::*;

fn main() -> Result<()> {
    // Initialize the framework
    rtx_transformers::init()?;

    println!("🔥 RTX MoCo v3 Self-Supervised Learning Example");
    println!("================================================");

    let device = Device::cuda(0).unwrap_or(Device::default());

    // Configure MoCo v3 with research-recommended hyperparameters
    let moco_config = MoCoV3Config::new(768, 256)
        .with_tau(0.999) // Momentum coefficient for EMA
        .with_temperature(0.07) // Temperature for InfoNCE loss
        .with_queue_size(65536); // Large queue for negative samples

    println!("✅ MoCo v3 Configuration:");
    println!("   - Feature dim: {}", moco_config.feature_dim);
    println!("   - Predictor dim: {}", moco_config.predictor_dim);
    println!("   - Momentum τ: {}", moco_config.tau);
    println!("   - Temperature: {}", moco_config.temperature);
    println!("   - Queue size: {}", moco_config.queue_size);

    // Create SSL training configuration
    let ssl_config = SSLTrainingConfig {
        method: SSLMethod::MoCoV3(moco_config),
        learning_rate: 1e-3,
        batch_size: 256, // Larger batch size recommended for MoCo v3
        epochs: 200,     // Longer training than supervised learning
        warmup_epochs: 10,
        weight_decay: 1e-4,
        eval_freq: 10,
        augmentation: AugmentationConfig {
            flip_prob: 0.5,
            color_jitter: 0.4,
            crop_scale: (0.2, 1.0),
            blur_prob: 0.1,
            normalize_mean: vec![0.485, 0.456, 0.406],
            normalize_std: vec![0.229, 0.224, 0.225],
        },
    };

    println!("\n✅ Training Configuration:");
    println!("   - Batch size: {}", ssl_config.batch_size);
    println!("   - Learning rate: {}", ssl_config.learning_rate);
    println!("   - Epochs: {}", ssl_config.epochs);

    // Create vision encoders (ResNet-50 style)
    let query_encoder = VisionBackbone::resnet50(&device);
    let key_encoder = VisionBackbone::resnet50(&device);

    println!("\n✅ Created Query and Key Encoders (ResNet-50 style)");

    // Create SSL trainer
    let mut trainer = SSLTrainer::new(query_encoder, Some(key_encoder), ssl_config, &device)?;

    println!("✅ Created MoCo v3 SSL Trainer");

    // Simulate training data (in practice, load from ImageNet)
    println!("\n🔥 Starting Training Loop...");
    println!("Note: This is a demonstration with synthetic data");

    for epoch in 0..5 {
        // Reduced epochs for demo
        trainer.set_epoch(epoch);
        trainer.train();

        // Simulate batch of images (3 channels, 224x224)
        let images = Tensor::randn(&[32, 3, 224, 224], DType::F32, &device)?;

        // Perform one training step
        let metrics = trainer.train_step(&images, Some(epoch as u64))?;

        if let MethodMetrics::MoCoV3 {
            similarity,
            queue_usage,
        } = metrics.method_metrics
        {
            println!(
                "Epoch {:2}: Loss = {:.6}, Pos Sim = {:.4}, Queue = {:.1}%",
                epoch,
                metrics.train_loss,
                similarity,
                queue_usage * 100.0
            );
        }
    }

    println!("\n✅ Training Complete!");

    // Feature extraction for downstream tasks
    println!("\n🔥 Feature Extraction for Downstream Tasks");
    trainer.eval();

    // Extract features from test images
    let test_images = Tensor::randn(&[10, 3, 224, 224], DType::F32, &device)?;
    let features = trainer.extract_features(&test_images)?;

    println!("✅ Extracted features: shape = {:?}", features.shape());
    println!("   These features can be used for:");
    println!("   - Linear classification");
    println!("   - k-NN evaluation");
    println!("   - Transfer learning");
    println!("   - Similarity search");

    // Demonstrate SSL evaluation
    println!("\n🔥 SSL Evaluation Demo");
    let evaluator = SSLEvaluator::new(&device);

    // Create dummy labeled data for evaluation
    let train_images = Tensor::randn(&[100, 3, 224, 224], DType::F32, &device)?;
    let train_labels = Tensor::randint(0, 10, &[100], DType::F32, &device)?;
    let val_images = Tensor::randn(&[20, 3, 224, 224], DType::F32, &device)?;
    let val_labels = Tensor::randint(0, 10, &[20], DType::F32, &device)?;

    // Linear probing evaluation
    let accuracy = evaluator.linear_probe(
        &mut trainer,
        &train_images,
        &train_labels,
        &val_images,
        &val_labels,
    )?;

    println!("✅ Linear Probing Accuracy: {:.1}%", accuracy * 100.0);

    // k-NN evaluation
    let knn_accuracy = evaluator.knn_evaluation(
        &mut trainer,
        &train_images,
        &train_labels,
        &val_images,
        &val_labels,
        5, // k=5
    )?;

    println!("✅ k-NN (k=5) Accuracy: {:.1}%", knn_accuracy * 100.0);

    println!("\n🎉 MoCo v3 Example Complete!");
    println!("Key advantages of MoCo v3:");
    println!("   ✓ No batch shuffling needed (simpler than v1/v2)");
    println!("   ✓ Predictor head improves representation quality");
    println!("   ✓ Works well with larger batch sizes");
    println!("   ✓ Strong performance on Vision Transformers");

    Ok(())
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_moco_v3_example() {
        // Test that our example compiles and runs without errors
        let result = main();
        assert!(result.is_ok(), "MoCo v3 example should run successfully");
    }
}