322 lines
11 KiB
Rust
322 lines
11 KiB
Rust
//! Barlow Twins Demo
|
|
//!
|
|
//! Demonstrates Barlow Twins self-supervised learning with various configurations
|
|
//! and backbone architectures.
|
|
|
|
use rtx_transformers::prelude::*;
|
|
use rtx_transformers::ssl::*;
|
|
|
|
fn main() -> Result<()> {
|
|
println!("🔥 Barlow Twins Self-Supervised Learning Demo");
|
|
println!("===============================================");
|
|
|
|
let device = Device::cuda(0).unwrap_or(Device::default());
|
|
|
|
// Demo 1: Basic Barlow Twins training
|
|
demo_basic_training(&device)?;
|
|
|
|
// Demo 2: Custom projector configuration
|
|
demo_custom_projector(&device)?;
|
|
|
|
// Demo 3: Different backbone architectures
|
|
demo_different_backbones(&device)?;
|
|
|
|
// Demo 4: Integration with SSL framework
|
|
demo_ssl_integration(&device)?;
|
|
|
|
// Demo 5: Loss analysis and visualization
|
|
demo_loss_analysis(&device)?;
|
|
|
|
println!("✅ All Barlow Twins demos completed successfully!");
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Demo 1: Basic Barlow Twins training
|
|
fn demo_basic_training(device: &Device) -> Result<()> {
|
|
println!("\n📚 Demo 1: Basic Barlow Twins Training");
|
|
println!("--------------------------------------");
|
|
|
|
// Create Barlow Twins config with default hyperparameters
|
|
let config = BarlowTwinsConfig::default().with_backbone_dim(2048);
|
|
|
|
println!("Config: {:?}", config);
|
|
|
|
// Create backbone and trainer
|
|
let backbone = Box::new(VisionBackbone::resnet50(device));
|
|
let mut trainer = BarlowTwinsTrainer::new(backbone, config, device)?;
|
|
|
|
// Create fake batch of images (batch_size=32, channels=3, height=224, width=224)
|
|
let batch_size = 32;
|
|
let images = Tensor::randn(vec![batch_size, 3, 224, 224], DType::F32, device)?;
|
|
|
|
println!("Training on batch shape: {:?}", images.shape());
|
|
|
|
// Training step
|
|
trainer.train();
|
|
let result = trainer.train_step(&images, Some(42))?;
|
|
|
|
let loss_value = result.loss.to_vec::<f32>()?[0];
|
|
println!("Loss: {:.6}", loss_value);
|
|
println!("Invariance loss: {:.6}", result.invariance_loss);
|
|
println!("Redundancy loss: {:.6}", result.redundancy_loss);
|
|
|
|
// Feature extraction
|
|
trainer.eval();
|
|
let features = trainer.extract_features(&images)?;
|
|
println!("Extracted features shape: {:?}", features.shape());
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Demo 2: Custom projector configuration
|
|
fn demo_custom_projector(device: &Device) -> Result<()> {
|
|
println!("\n🔧 Demo 2: Custom Projector Configuration");
|
|
println!("------------------------------------------");
|
|
|
|
// Custom projector with different architecture
|
|
let custom_config = BarlowTwinsConfig::new(2048, vec![4096, 2048, 1024])
|
|
.with_lambda_coeff(0.01) // Higher regularization
|
|
.with_batch_norm(false) // Without batch norm
|
|
.with_scale_loss(0.1); // Different scaling
|
|
|
|
println!("Custom config: {:?}", custom_config);
|
|
|
|
let backbone = Box::new(VisionBackbone::resnet50(device));
|
|
let mut trainer = BarlowTwinsTrainer::new(backbone, custom_config, device)?;
|
|
|
|
let images = Tensor::randn(vec![16, 3, 224, 224], DType::F32, device)?;
|
|
|
|
trainer.train();
|
|
let result = trainer.train_step(&images, Some(123))?;
|
|
|
|
let loss_value = result.loss.to_vec::<f32>()?[0];
|
|
println!("Custom config loss: {:.6}", loss_value);
|
|
|
|
// Analyze cross-correlation matrix
|
|
let cross_corr = result.cross_correlation;
|
|
let diagonal = extract_diagonal(&cross_corr)?;
|
|
let diagonal_data = diagonal.to_vec::<f32>()?;
|
|
let diagonal_mean = diagonal_data.iter().sum::<f32>() / diagonal_data.len() as f32;
|
|
|
|
println!("Cross-correlation matrix shape: {:?}", cross_corr.shape());
|
|
println!("Diagonal mean: {:.4} (should approach 1.0)", diagonal_mean);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Demo 3: Different backbone architectures
|
|
fn demo_different_backbones(device: &Device) -> Result<()> {
|
|
println!("\n🏗️ Demo 3: Different Backbone Architectures");
|
|
println!("---------------------------------------------");
|
|
|
|
let batch_size = 16;
|
|
let images = Tensor::randn(vec![batch_size, 3, 224, 224], DType::F32, device)?;
|
|
|
|
// ResNet-50 style (2048-dim features)
|
|
println!("Testing ResNet-50 style backbone:");
|
|
let resnet_config = BarlowTwinsConfig::default().with_backbone_dim(2048);
|
|
let resnet_backbone = Box::new(VisionBackbone::resnet50(device));
|
|
let mut resnet_trainer = BarlowTwinsTrainer::new(resnet_backbone, resnet_config, device)?;
|
|
|
|
resnet_trainer.train();
|
|
let resnet_result = resnet_trainer.train_step(&images, Some(42))?;
|
|
let resnet_loss = resnet_result.loss.to_vec::<f32>()?[0];
|
|
println!(" ResNet loss: {:.6}", resnet_loss);
|
|
|
|
// ViT-Base style (768-dim features)
|
|
println!("Testing ViT-Base style backbone:");
|
|
let vit_config = BarlowTwinsConfig::new(768, vec![3072, 3072, 3072]).with_lambda_coeff(0.005);
|
|
let vit_backbone = Box::new(VisionBackbone::vit_base(device));
|
|
let mut vit_trainer = BarlowTwinsTrainer::new(vit_backbone, vit_config, device)?;
|
|
|
|
vit_trainer.train();
|
|
let vit_result = vit_trainer.train_step(&images, Some(42))?;
|
|
let vit_loss = vit_result.loss.to_vec::<f32>()?[0];
|
|
println!(" ViT loss: {:.6}", vit_loss);
|
|
|
|
// Feature extraction comparison
|
|
resnet_trainer.eval();
|
|
vit_trainer.eval();
|
|
|
|
let resnet_features = resnet_trainer.extract_features(&images)?;
|
|
let vit_features = vit_trainer.extract_features(&images)?;
|
|
|
|
println!(" ResNet features: {:?}", resnet_features.shape());
|
|
println!(" ViT features: {:?}", vit_features.shape());
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Demo 4: Integration with SSL framework
|
|
fn demo_ssl_integration(device: &Device) -> Result<()> {
|
|
println!("\n🔄 Demo 4: SSL Framework Integration");
|
|
println!("------------------------------------");
|
|
|
|
// Create Barlow Twins SSL config
|
|
let barlow_config = BarlowTwinsConfig::new(2048, vec![8192, 8192, 8192])
|
|
.with_lambda_coeff(0.005)
|
|
.with_batch_norm(true);
|
|
|
|
let ssl_config = SSLTrainingConfig {
|
|
method: SSLMethod::BarlowTwins(barlow_config),
|
|
learning_rate: 1e-3,
|
|
batch_size: 32,
|
|
epochs: 10,
|
|
warmup_epochs: 2,
|
|
weight_decay: 1e-4,
|
|
eval_freq: 5,
|
|
augmentation: AugmentationConfig {
|
|
flip_prob: 0.5,
|
|
color_jitter: 0.4,
|
|
crop_scale: (0.08, 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!("SSL config created for Barlow Twins");
|
|
|
|
// Create SSL trainer
|
|
let backbone = VisionBackbone::resnet50(device);
|
|
let mut trainer = SSLTrainer::new(
|
|
backbone, None, // Barlow Twins doesn't need separate target backbone
|
|
ssl_config, device,
|
|
)?;
|
|
|
|
let images = Tensor::randn(vec![32, 3, 224, 224], DType::F32, device)?;
|
|
|
|
// Training loop simulation
|
|
for epoch in 0..5 {
|
|
trainer.set_epoch(epoch);
|
|
trainer.train();
|
|
|
|
let metrics = trainer.train_step(&images, Some(epoch as u64))?;
|
|
|
|
println!("Epoch {}: Loss = {:.6}", epoch, metrics.train_loss);
|
|
|
|
if let MethodMetrics::BarlowTwins {
|
|
invariance_loss,
|
|
redundancy_loss,
|
|
diagonal_mean,
|
|
off_diagonal_rms,
|
|
} = metrics.method_metrics
|
|
{
|
|
println!(
|
|
" Invariance: {:.4}, Redundancy: {:.4}",
|
|
invariance_loss, redundancy_loss
|
|
);
|
|
println!(
|
|
" Diagonal mean: {:.4}, Off-diag RMS: {:.4}",
|
|
diagonal_mean, off_diagonal_rms
|
|
);
|
|
}
|
|
}
|
|
|
|
// Evaluation
|
|
trainer.eval();
|
|
let features = trainer.extract_features(&images)?;
|
|
println!("Final features shape: {:?}", features.shape());
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Demo 5: Loss analysis and visualization
|
|
fn demo_loss_analysis(device: &Device) -> Result<()> {
|
|
println!("\n📊 Demo 5: Loss Analysis");
|
|
println!("------------------------");
|
|
|
|
let batch_size = 64;
|
|
let feature_dim = 1024;
|
|
let lambda_coeff = 0.005;
|
|
let scale_loss = 1.0 / 32.0;
|
|
|
|
// Test with random embeddings
|
|
println!("Analyzing loss with random embeddings:");
|
|
let y1 = Tensor::randn(vec![batch_size, feature_dim], DType::F32, device)?;
|
|
let y2 = Tensor::randn(vec![batch_size, feature_dim], DType::F32, device)?;
|
|
|
|
let loss_result = compute_barlow_twins_loss(&y1, &y2, lambda_coeff, scale_loss)?;
|
|
let loss_value = loss_result.loss.to_vec::<f32>()?[0];
|
|
|
|
println!(" Total loss: {:.6}", loss_value);
|
|
println!(" Invariance loss: {:.6}", loss_result.invariance_loss);
|
|
println!(" Redundancy loss: {:.6}", loss_result.redundancy_loss);
|
|
|
|
// Test with identical embeddings (should give low loss)
|
|
println!("\nAnalyzing loss with identical embeddings:");
|
|
let y_identical = normalize_features(&y1)?;
|
|
let identical_loss =
|
|
compute_barlow_twins_loss(&y_identical, &y_identical, lambda_coeff, scale_loss)?;
|
|
let identical_loss_value = identical_loss.loss.to_vec::<f32>()?[0];
|
|
|
|
println!(" Total loss: {:.6}", identical_loss_value);
|
|
println!(" Invariance loss: {:.6}", identical_loss.invariance_loss);
|
|
println!(" Redundancy loss: {:.6}", identical_loss.redundancy_loss);
|
|
|
|
// Test symmetry
|
|
println!("\nTesting loss symmetry:");
|
|
let loss_12 = compute_barlow_twins_loss(&y1, &y2, lambda_coeff, scale_loss)?;
|
|
let loss_21 = compute_barlow_twins_loss(&y2, &y1, lambda_coeff, scale_loss)?;
|
|
|
|
let loss_12_val = loss_12.loss.to_vec::<f32>()?[0];
|
|
let loss_21_val = loss_21.loss.to_vec::<f32>()?[0];
|
|
let symmetry_diff = (loss_12_val - loss_21_val).abs();
|
|
|
|
println!(" Loss(y1, y2): {:.6}", loss_12_val);
|
|
println!(" Loss(y2, y1): {:.6}", loss_21_val);
|
|
println!(" Symmetry difference: {:.8} (should be ~0)", symmetry_diff);
|
|
|
|
// Cross-correlation matrix analysis
|
|
println!("\nCross-correlation matrix analysis:");
|
|
let cross_corr = compute_cross_correlation_matrix(&y1, &y2)?;
|
|
let diagonal = extract_diagonal(&cross_corr)?;
|
|
let diagonal_data = diagonal.to_vec::<f32>()?;
|
|
|
|
let diagonal_mean = diagonal_data.iter().sum::<f32>() / diagonal_data.len() as f32;
|
|
let diagonal_std = {
|
|
let variance = diagonal_data
|
|
.iter()
|
|
.map(|&x| (x - diagonal_mean).powi(2))
|
|
.sum::<f32>()
|
|
/ diagonal_data.len() as f32;
|
|
variance.sqrt()
|
|
};
|
|
|
|
println!(" Matrix shape: {:?}", cross_corr.shape());
|
|
println!(
|
|
" Diagonal mean: {:.4} ± {:.4}",
|
|
diagonal_mean, diagonal_std
|
|
);
|
|
|
|
Ok(())
|
|
}
|
|
|
|
/// Simple test backbone for demonstrations
|
|
#[derive(Debug, Clone)]
|
|
struct TestBackbone {
|
|
output_dim: usize,
|
|
}
|
|
|
|
impl TestBackbone {
|
|
#[allow(dead_code)]
|
|
fn new(output_dim: usize) -> Self {
|
|
Self { output_dim }
|
|
}
|
|
}
|
|
|
|
impl Backbone for TestBackbone {
|
|
fn forward(&self, input: &Tensor) -> Result<Tensor> {
|
|
let batch_size = input.shape()[0];
|
|
let device = input.device();
|
|
|
|
// Simple projection to output_dim
|
|
Tensor::randn(vec![batch_size, self.output_dim], DType::F32, device)
|
|
}
|
|
|
|
fn output_dim(&self) -> usize {
|
|
self.output_dim
|
|
}
|
|
}
|