324 lines
12 KiB
Rust
324 lines
12 KiB
Rust
//! Comprehensive Meta-Learning Example
|
||
//!
|
||
//! This example demonstrates how to use the RTX Transformers meta-learning
|
||
//! implementation for few-shot learning tasks using both MAML and
|
||
//! Prototypical Networks.
|
||
|
||
use rtx_transformers::meta::*;
|
||
use rtx_tensor::{Device, Tensor};
|
||
use std::collections::HashMap;
|
||
|
||
fn main() -> Result<(), Box<dyn std::error::Error>> {
|
||
println!("🚀 RTX Transformers Meta-Learning Comprehensive Example");
|
||
println!("========================================================\n");
|
||
|
||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||
|
||
// Step 1: Create a more realistic few-shot dataset
|
||
println!("📊 Step 1: Creating Few-Shot Learning Dataset");
|
||
println!("-----------------------------------------------");
|
||
|
||
let num_classes = 10;
|
||
let samples_per_class = 20;
|
||
let feature_dim = 128;
|
||
|
||
let dataset = FewShotDataset::synthetic(num_classes, samples_per_class, feature_dim, &device)?;
|
||
println!("✓ Created dataset:");
|
||
println!(" - {} classes", dataset.num_classes);
|
||
println!(" - {} samples per class", samples_per_class);
|
||
println!(" - {} feature dimensions", dataset.feature_dim);
|
||
println!();
|
||
|
||
// Step 2: Configure different few-shot scenarios
|
||
println!("🎯 Step 2: Few-Shot Learning Scenarios");
|
||
println!("--------------------------------------");
|
||
|
||
let scenarios = vec![
|
||
("5-way 1-shot", 5, 1, 5),
|
||
("5-way 5-shot", 5, 5, 5),
|
||
("10-way 1-shot", 10, 1, 3),
|
||
("3-way 2-shot", 3, 2, 7),
|
||
];
|
||
|
||
for (name, n_way, k_shot, query_per_class) in &scenarios {
|
||
println!("📋 Scenario: {}", name);
|
||
let episode = dataset.sample_episode(*n_way, *k_shot, *query_per_class, &device)?;
|
||
|
||
println!(" Support set: {} samples ({} × {} shots)",
|
||
episode.support_set.len(), n_way, k_shot);
|
||
println!(" Query set: {} samples ({} × {} queries)",
|
||
episode.query_set.len(), n_way, query_per_class);
|
||
|
||
// Verify episode structure
|
||
let (support_x, support_y) = episode.support_batch(&device)?;
|
||
let (query_x, query_y) = episode.query_batch(&device)?;
|
||
|
||
println!(" Support batch shape: {:?}", support_x.shape());
|
||
println!(" Query batch shape: {:?}", query_x.shape());
|
||
println!();
|
||
}
|
||
|
||
// Step 3: MAML Demonstration
|
||
println!("🧠 Step 3: MAML (Model-Agnostic Meta-Learning)");
|
||
println!("===============================================");
|
||
|
||
let maml_config = MAMLConfig {
|
||
inner_lr: 0.01,
|
||
outer_lr: 0.001,
|
||
inner_steps: 3,
|
||
first_order: false,
|
||
meta_batch_size: 8,
|
||
};
|
||
|
||
println!("📋 MAML Configuration:");
|
||
println!(" Inner learning rate: {}", maml_config.inner_lr);
|
||
println!(" Outer learning rate: {}", maml_config.outer_lr);
|
||
println!(" Inner steps: {}", maml_config.inner_steps);
|
||
println!(" First-order approximation: {}", maml_config.first_order);
|
||
println!();
|
||
|
||
let mut maml = MAML::new(feature_dim, 64, 5, maml_config, &device)?;
|
||
|
||
// Generate training episodes for MAML
|
||
let mut sampler = EpisodeSampler::new(dataset.clone(), EpisodeSamplerConfig {
|
||
n_way: 5,
|
||
k_shot: 1,
|
||
query_per_class: 3,
|
||
num_episodes: 16,
|
||
seed: Some(42),
|
||
});
|
||
|
||
println!("🏋️ Training MAML...");
|
||
let training_episodes = sampler.sample_episodes(8, &device)?;
|
||
let maml_stats = maml.meta_update(training_episodes, &device)?;
|
||
|
||
println!("✓ MAML Training Results:");
|
||
println!(" Meta-updates: {}", maml_stats.meta_updates);
|
||
println!(" Tasks processed: {}", maml_stats.tasks_processed);
|
||
println!(" Average inner loss: {:.4}", maml_stats.avg_inner_loss);
|
||
println!(" Average outer loss: {:.4}", maml_stats.avg_outer_loss);
|
||
println!(" Average query accuracy: {:.4}", maml_stats.avg_query_accuracy);
|
||
println!();
|
||
|
||
// Test fast adaptation
|
||
println!("⚡ Testing MAML Fast Adaptation...");
|
||
let test_episode = sampler.sample_episode(&device)?;
|
||
let adapted_network = maml.fast_adaptation(&test_episode, &device)?;
|
||
|
||
let (query_x, query_y) = test_episode.query_batch(&device)?;
|
||
let predictions = adapted_network.forward(&query_x)?;
|
||
let accuracy = maml.compute_accuracy(&predictions, &query_y)?;
|
||
|
||
println!("✓ Fast Adaptation Results:");
|
||
println!(" Adapted in {} inner steps", maml.config.inner_steps);
|
||
println!(" Query accuracy after adaptation: {:.4}", accuracy);
|
||
println!();
|
||
|
||
// Step 4: Prototypical Networks Demonstration
|
||
println!("🎯 Step 4: Prototypical Networks");
|
||
println!("================================");
|
||
|
||
// Test different distance metrics
|
||
let distance_metrics = vec![
|
||
("Euclidean", DistanceMetric::Euclidean),
|
||
("Cosine", DistanceMetric::Cosine),
|
||
("Manhattan", DistanceMetric::Manhattan),
|
||
];
|
||
|
||
let mut proto_results = HashMap::new();
|
||
|
||
for (metric_name, metric) in &distance_metrics {
|
||
println!("📏 Testing {} Distance Metric", metric_name);
|
||
|
||
let proto_config = PrototypicalConfig {
|
||
learning_rate: 0.001,
|
||
feature_dim: 32,
|
||
distance_metric: metric.clone(),
|
||
temperature: 1.0,
|
||
batch_size: 4,
|
||
};
|
||
|
||
let proto_net = PrototypicalNetworks::new(feature_dim, proto_config, &device)?;
|
||
|
||
// Test on multiple episodes
|
||
let mut accuracies = Vec::new();
|
||
let mut losses = Vec::new();
|
||
|
||
for _ in 0..5 {
|
||
let episode = dataset.sample_episode(3, 2, 4, &device)?;
|
||
let (logits, _, accuracy) = proto_net.classify(&episode, &device)?;
|
||
let (_, query_y) = episode.query_batch(&device)?;
|
||
let loss = proto_net.compute_loss(&logits, &query_y)?;
|
||
|
||
accuracies.push(accuracy);
|
||
losses.push(loss.to_vec()?[0]);
|
||
}
|
||
|
||
let mean_accuracy = accuracies.iter().sum::<f32>() / accuracies.len() as f32;
|
||
let mean_loss = losses.iter().sum::<f32>() / losses.len() as f32;
|
||
|
||
proto_results.insert(metric_name, (mean_accuracy, mean_loss));
|
||
|
||
println!(" ✓ Mean accuracy: {:.4}", mean_accuracy);
|
||
println!(" ✓ Mean loss: {:.4}", mean_loss);
|
||
println!();
|
||
}
|
||
|
||
// Step 5: Comparison and Analysis
|
||
println!("📊 Step 5: Method Comparison & Analysis");
|
||
println!("=======================================");
|
||
|
||
// Compare distance metrics
|
||
println!("🔍 Distance Metric Comparison:");
|
||
let mut best_metric = "";
|
||
let mut best_accuracy = 0.0;
|
||
|
||
for (metric, (accuracy, loss)) in &proto_results {
|
||
println!(" {} - Accuracy: {:.4}, Loss: {:.4}", metric, accuracy, loss);
|
||
if *accuracy > best_accuracy {
|
||
best_accuracy = *accuracy;
|
||
best_metric = metric;
|
||
}
|
||
}
|
||
println!(" 🏆 Best performing metric: {} ({:.4} accuracy)", best_metric, best_accuracy);
|
||
println!();
|
||
|
||
// Step 6: Complete Pipeline Demonstration
|
||
println!("🔄 Step 6: Complete Meta-Learning Pipeline");
|
||
println!("==========================================");
|
||
|
||
let pipeline_config = MetaLearningConfig {
|
||
input_dim: feature_dim,
|
||
hidden_dim: 64,
|
||
output_dim: num_classes,
|
||
algorithm: MetaLearningAlgorithm::Hybrid,
|
||
sampler_config: EpisodeSamplerConfig {
|
||
n_way: 5,
|
||
k_shot: 2,
|
||
query_per_class: 3,
|
||
num_episodes: 20,
|
||
seed: Some(12345),
|
||
},
|
||
maml_config: MAMLConfig {
|
||
inner_lr: 0.02,
|
||
outer_lr: 0.001,
|
||
inner_steps: 2,
|
||
first_order: false,
|
||
meta_batch_size: 4,
|
||
},
|
||
proto_config: PrototypicalConfig {
|
||
learning_rate: 0.001,
|
||
feature_dim: 32,
|
||
distance_metric: DistanceMetric::Euclidean,
|
||
temperature: 0.5,
|
||
batch_size: 4,
|
||
},
|
||
};
|
||
|
||
println!("🚀 Running Hybrid Pipeline (MAML + Prototypical)...");
|
||
let mut pipeline = MetaLearningPipeline::new(dataset, pipeline_config, &device)?;
|
||
let training_results = pipeline.train(10, &device)?;
|
||
|
||
println!("✓ Pipeline Training Complete:");
|
||
if let Some(maml_stats) = &training_results.maml_stats {
|
||
println!(" MAML - Meta-updates: {}, Accuracy: {:.4}",
|
||
maml_stats.meta_updates, maml_stats.avg_query_accuracy);
|
||
}
|
||
if let Some(proto_stats) = &training_results.proto_stats {
|
||
println!(" Prototypical - Episodes: {}, Accuracy: {:.4}",
|
||
proto_stats.episodes_processed, proto_stats.avg_accuracy);
|
||
}
|
||
println!();
|
||
|
||
// Generate test episodes for evaluation
|
||
println!("🧪 Evaluating on Test Episodes...");
|
||
let test_episodes = pipeline.sampler.sample_episodes(5, &device)?;
|
||
let evaluation_results = pipeline.evaluate(test_episodes, &device)?;
|
||
|
||
if let Some(metrics) = &evaluation_results.evaluation_metrics {
|
||
println!("✓ Evaluation Results:");
|
||
println!(" Test episodes: {}", metrics.num_episodes);
|
||
println!(" Mean accuracy: {:.4}", metrics.mean_accuracy);
|
||
println!(" Mean loss: {:.4}", metrics.mean_loss);
|
||
println!(" 95% Confidence interval: ({:.4}, {:.4})",
|
||
metrics.confidence_interval.0, metrics.confidence_interval.1);
|
||
}
|
||
|
||
if let Some(comparison) = &evaluation_results.comparison {
|
||
println!(" Method comparison: {} is better by {:.4} accuracy",
|
||
comparison.better_approach, comparison.accuracy_difference.abs());
|
||
}
|
||
println!();
|
||
|
||
// Step 7: Advanced Features Demo
|
||
println!("🎛️ Step 7: Advanced Features");
|
||
println!("============================");
|
||
|
||
// Task distribution sampling
|
||
println!("🎲 Task Distribution Sampling:");
|
||
let base_config = EpisodeSamplerConfig {
|
||
n_way: 3,
|
||
k_shot: 1,
|
||
query_per_class: 2,
|
||
num_episodes: 1,
|
||
seed: Some(999),
|
||
};
|
||
|
||
let mut task_dist = TaskDistribution::multi_task(base_config);
|
||
|
||
for i in 0..3 {
|
||
let config = task_dist.sample_task_config();
|
||
println!(" Task {}: {}-way {}-shot", i + 1, config.n_way, config.k_shot);
|
||
}
|
||
println!();
|
||
|
||
// Episode consistency check
|
||
println!("🔍 Episode Consistency Check:");
|
||
let mut sampler = EpisodeSampler::new(
|
||
FewShotDataset::synthetic(5, 10, 16, &device)?,
|
||
EpisodeSamplerConfig {
|
||
n_way: 3,
|
||
k_shot: 1,
|
||
query_per_class: 2,
|
||
seed: Some(777),
|
||
..Default::default()
|
||
}
|
||
);
|
||
|
||
let episode1 = sampler.sample_episode(&device)?;
|
||
sampler.reset();
|
||
let episode2 = sampler.sample_episode(&device)?;
|
||
|
||
let consistent = episode1.support_labels == episode2.support_labels;
|
||
println!(" Reproducible sampling: {}", if consistent { "✓ PASS" } else { "✗ FAIL" });
|
||
println!();
|
||
|
||
// Final Summary
|
||
println!("🎉 Meta-Learning Implementation Summary");
|
||
println!("======================================");
|
||
println!("✅ MAML (Model-Agnostic Meta-Learning):");
|
||
println!(" - Inner/outer loop optimization implemented");
|
||
println!(" - Fast adaptation for new tasks");
|
||
println!(" - First-order approximation (FOMAML) support");
|
||
println!();
|
||
println!("✅ Prototypical Networks:");
|
||
println!(" - Multiple distance metrics (Euclidean, Cosine, Manhattan)");
|
||
println!(" - Prototype computation from support sets");
|
||
println!(" - Temperature scaling for probability calibration");
|
||
println!();
|
||
println!("✅ Episode Sampling & Evaluation:");
|
||
println!(" - N-way K-shot episode generation");
|
||
println!(" - Reproducible sampling with seeding");
|
||
println!(" - Comprehensive evaluation metrics");
|
||
println!(" - Task distribution sampling");
|
||
println!();
|
||
println!("✅ Integration Features:");
|
||
println!(" - Transformer model integration ready");
|
||
println!(" - Complete meta-learning pipeline");
|
||
println!(" - Hybrid algorithm support");
|
||
println!(" - Production-ready implementation");
|
||
println!();
|
||
println!("🚀 RTX Transformers Meta-Learning: IMPLEMENTATION COMPLETE!");
|
||
|
||
Ok(())
|
||
} |