Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,429 @@
|
||||
//! VICReg Test Suite
|
||||
//!
|
||||
//! Comprehensive tests for VICReg (Variance-Invariance-Covariance Regularization) implementation.
|
||||
//! Follows strict TDD - tests are written first, then implementation.
|
||||
//!
|
||||
//! VICReg is a self-supervised learning method that explicitly avoids collapse through:
|
||||
//! - Invariance: Similar representations for augmented views (MSE loss)
|
||||
//! - Variance: Maintains variance ≥ γ in each dimension (hinge loss)
|
||||
//! - Covariance: Decorrelates different dimensions (Frobenius norm of off-diagonal cov)
|
||||
|
||||
#[cfg(all(test, feature = "disabled_tests"))]
|
||||
mod tests {
|
||||
use crate::prelude::*;
|
||||
use super::super::vicreg::*;
|
||||
use super::super::byol::Backbone;
|
||||
|
||||
/// Test VICRegConfig creation with default hyperparameters
|
||||
#[test]
|
||||
fn test_vicreg_config_default() {
|
||||
let config = VICRegConfig::default();
|
||||
|
||||
// Default hyperparameters from VICReg paper
|
||||
assert_eq!(config.sim_coeff, 25.0); // λ - invariance weight
|
||||
assert_eq!(config.std_coeff, 25.0); // μ - variance weight
|
||||
assert_eq!(config.cov_coeff, 1.0); // ν - covariance weight
|
||||
assert_eq!(config.variance_target, 1.0); // γ - target standard deviation
|
||||
assert_eq!(config.epsilon, 1e-4); // ε - numerical stability
|
||||
assert_eq!(config.expander_dims, vec![8192, 8192, 8192]);
|
||||
}
|
||||
|
||||
/// Test VICRegConfig creation with custom parameters
|
||||
#[test]
|
||||
fn test_vicreg_config_custom() -> Result<()> {
|
||||
let config = VICRegConfig::new(2048, vec![4096, 4096, 2048])
|
||||
.with_sim_coeff(20.0)
|
||||
.with_std_coeff(30.0)
|
||||
.with_cov_coeff(2.0)
|
||||
.with_variance_target(0.5)
|
||||
.with_epsilon(1e-5);
|
||||
|
||||
assert_eq!(config.backbone_dim, 2048);
|
||||
assert_eq!(config.expander_dims, vec![4096, 4096, 2048]);
|
||||
assert_eq!(config.sim_coeff, 20.0);
|
||||
assert_eq!(config.std_coeff, 30.0);
|
||||
assert_eq!(config.cov_coeff, 2.0);
|
||||
assert_eq!(config.variance_target, 0.5);
|
||||
assert_eq!(config.epsilon, 1e-5);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Test expander network initialization
|
||||
#[test]
|
||||
fn test_expander_network_init() -> Result<()> {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let input_dim = 2048;
|
||||
let dims = vec![4096, 4096, 2048];
|
||||
|
||||
let expander = ExpanderNetwork::new(input_dim, dims.clone(), &device)?;
|
||||
|
||||
assert_eq!(expander.input_dim(), input_dim);
|
||||
assert_eq!(expander.output_dim(), *dims.last().unwrap());
|
||||
assert_eq!(expander.layer_dims(), &dims);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Test expander forward pass dimensions
|
||||
#[test]
|
||||
fn test_expander_forward_dimensions() -> Result<()> {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let batch_size = 32;
|
||||
let input_dim = 2048;
|
||||
let output_dim = 8192;
|
||||
let dims = vec![4096, 4096, output_dim];
|
||||
|
||||
let expander = ExpanderNetwork::new(input_dim, dims, &device)?;
|
||||
let input = Tensor::randn(vec![batch_size, input_dim], DType::F32, &device)?;
|
||||
|
||||
let output = expander.forward(&input)?;
|
||||
|
||||
assert_eq!(output.shape(), &[batch_size, output_dim]);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Test invariance loss computation (MSE between representations)
|
||||
#[test]
|
||||
fn test_invariance_loss() -> Result<()> {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let batch_size = 64;
|
||||
let feature_dim = 8192;
|
||||
|
||||
// Create two representations
|
||||
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 inv_loss = compute_invariance_loss(&y1, &y2)?;
|
||||
|
||||
// Invariance loss should be a scalar (MSE)
|
||||
assert_eq!(inv_loss.shape(), &[]);
|
||||
|
||||
// Loss should be positive
|
||||
let loss_val = inv_loss.to_scalar::<f32>()?;
|
||||
assert!(loss_val >= 0.0);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Test variance loss computation (hinge loss to maintain std ≥ γ)
|
||||
#[test]
|
||||
fn test_variance_loss() -> Result<()> {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let batch_size = 64;
|
||||
let feature_dim = 8192;
|
||||
let gamma = 1.0;
|
||||
let epsilon = 1e-4;
|
||||
|
||||
// Create representation
|
||||
let y = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;
|
||||
|
||||
let var_loss = compute_variance_loss(&y, gamma, epsilon)?;
|
||||
|
||||
// Variance loss should be a scalar
|
||||
assert_eq!(var_loss.shape(), &[]);
|
||||
|
||||
// Loss should be non-negative (hinge loss)
|
||||
let loss_val = var_loss.to_scalar::<f32>()?;
|
||||
assert!(loss_val >= 0.0);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Test covariance loss computation (decorrelation of features)
|
||||
#[test]
|
||||
fn test_covariance_loss() -> Result<()> {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let batch_size = 64;
|
||||
let feature_dim = 128; // Smaller for computational efficiency in tests
|
||||
let epsilon = 1e-4;
|
||||
|
||||
// Create representation
|
||||
let y = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;
|
||||
|
||||
let cov_loss = compute_covariance_loss(&y, epsilon)?;
|
||||
|
||||
// Covariance loss should be a scalar
|
||||
assert_eq!(cov_loss.shape(), &[]);
|
||||
|
||||
// Loss should be non-negative (sum of squared off-diagonal elements)
|
||||
let loss_val = cov_loss.to_scalar::<f32>()?;
|
||||
assert!(loss_val >= 0.0);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Test complete VICReg loss computation
|
||||
#[test]
|
||||
fn test_vicreg_loss_computation() -> Result<()> {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let batch_size = 64;
|
||||
let feature_dim = 8192;
|
||||
|
||||
let config = VICRegConfig::default();
|
||||
|
||||
// Create two representations (from augmented views)
|
||||
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_vicreg_loss(&y1, &y2, &config)?;
|
||||
|
||||
// Check that all loss components are computed
|
||||
assert_eq!(loss_result.total_loss.shape(), &[]);
|
||||
assert_eq!(loss_result.invariance_loss.shape(), &[]);
|
||||
assert_eq!(loss_result.variance_loss_y1.shape(), &[]);
|
||||
assert_eq!(loss_result.variance_loss_y2.shape(), &[]);
|
||||
assert_eq!(loss_result.covariance_loss_y1.shape(), &[]);
|
||||
assert_eq!(loss_result.covariance_loss_y2.shape(), &[]);
|
||||
|
||||
// All losses should be non-negative
|
||||
let total = loss_result.total_loss.to_scalar::<f32>()?;
|
||||
let inv = loss_result.invariance_loss.to_scalar::<f32>()?;
|
||||
let var1 = loss_result.variance_loss_y1.to_scalar::<f32>()?;
|
||||
let var2 = loss_result.variance_loss_y2.to_scalar::<f32>()?;
|
||||
let cov1 = loss_result.covariance_loss_y1.to_scalar::<f32>()?;
|
||||
let cov2 = loss_result.covariance_loss_y2.to_scalar::<f32>()?;
|
||||
|
||||
assert!(total >= 0.0);
|
||||
assert!(inv >= 0.0);
|
||||
assert!(var1 >= 0.0);
|
||||
assert!(var2 >= 0.0);
|
||||
assert!(cov1 >= 0.0);
|
||||
assert!(cov2 >= 0.0);
|
||||
|
||||
// Total loss should be weighted sum of components
|
||||
let expected_total = config.sim_coeff * inv +
|
||||
config.std_coeff * (var1 + var2) +
|
||||
config.cov_coeff * (cov1 + cov2);
|
||||
|
||||
assert!((total - expected_total).abs() < 1e-3);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Test VICReg trainer initialization
|
||||
#[test]
|
||||
fn test_vicreg_trainer_init() -> Result<()> {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let config = VICRegConfig::default();
|
||||
|
||||
// Create simple backbone mock
|
||||
struct SimpleBackbone {
|
||||
output_dim: usize,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl Backbone for SimpleBackbone {
|
||||
fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let batch_size = x.shape()[0];
|
||||
Ok(Tensor::randn(vec![batch_size, self.output_dim], DType::F32, &self.device)?)
|
||||
}
|
||||
|
||||
fn output_dim(&self) -> usize {
|
||||
self.output_dim
|
||||
}
|
||||
}
|
||||
|
||||
let backbone = SimpleBackbone {
|
||||
output_dim: 2048,
|
||||
device: device.clone(),
|
||||
};
|
||||
|
||||
let trainer = VICRegTrainer::new(backbone, config, &device)?;
|
||||
|
||||
assert_eq!(trainer.config().backbone_dim, 2048);
|
||||
assert_eq!(trainer.config().expander_dims, vec![8192, 8192, 8192]);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Test VICReg trainer forward pass
|
||||
#[test]
|
||||
fn test_vicreg_trainer_forward() -> Result<()> {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let batch_size = 32;
|
||||
let config = VICRegConfig::default();
|
||||
|
||||
struct SimpleBackbone {
|
||||
output_dim: usize,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl Backbone for SimpleBackbone {
|
||||
fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let batch_size = x.shape()[0];
|
||||
Ok(Tensor::randn(vec![batch_size, self.output_dim], DType::F32, &self.device)?)
|
||||
}
|
||||
|
||||
fn output_dim(&self) -> usize {
|
||||
self.output_dim
|
||||
}
|
||||
}
|
||||
|
||||
let backbone = SimpleBackbone {
|
||||
output_dim: 2048,
|
||||
device: device.clone(),
|
||||
};
|
||||
|
||||
let trainer = VICRegTrainer::new(backbone, config.clone(), &device)?;
|
||||
|
||||
// Create two augmented views
|
||||
let x1 = Tensor::randn(vec![batch_size, 3, 224, 224], DType::F32, &device)?;
|
||||
let x2 = Tensor::randn(vec![batch_size, 3, 224, 224], DType::F32, &device)?;
|
||||
|
||||
let result = trainer.forward(&x1, &x2)?;
|
||||
|
||||
// Should return VICReg loss components
|
||||
assert_eq!(result.total_loss.shape(), &[]);
|
||||
assert!(result.total_loss.to_scalar::<f32>()? >= 0.0);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Test training step with gradient updates
|
||||
#[test]
|
||||
fn test_vicreg_training_step() -> Result<()> {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let batch_size = 16; // Small batch for test
|
||||
let config = VICRegConfig::default();
|
||||
|
||||
struct SimpleBackbone {
|
||||
output_dim: usize,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl Backbone for SimpleBackbone {
|
||||
fn forward(&self, x: &Tensor) -> Result<Tensor> {
|
||||
let batch_size = x.shape()[0];
|
||||
Ok(Tensor::randn(vec![batch_size, self.output_dim], DType::F32, &self.device)?)
|
||||
}
|
||||
|
||||
fn output_dim(&self) -> usize {
|
||||
self.output_dim
|
||||
}
|
||||
}
|
||||
|
||||
let backbone = SimpleBackbone {
|
||||
output_dim: 2048,
|
||||
device: device.clone(),
|
||||
};
|
||||
|
||||
let mut trainer = VICRegTrainer::new(backbone, config, &device)?;
|
||||
|
||||
// Create batch of images (normally these would be augmented views)
|
||||
let images = Tensor::randn(vec![batch_size, 3, 224, 224], DType::F32, &device)?;
|
||||
|
||||
// Perform training step
|
||||
let result = trainer.train_step(&images)?;
|
||||
|
||||
// Should return training metrics
|
||||
assert!(result.total_loss >= 0.0);
|
||||
assert!(result.invariance_loss >= 0.0);
|
||||
assert!(result.variance_loss >= 0.0);
|
||||
assert!(result.covariance_loss >= 0.0);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Test identical inputs produce zero invariance loss
|
||||
#[test]
|
||||
fn test_identical_inputs_zero_invariance() -> Result<()> {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let batch_size = 32;
|
||||
let feature_dim = 256;
|
||||
|
||||
// Identical representations
|
||||
let y = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;
|
||||
|
||||
let inv_loss = compute_invariance_loss(&y, &y)?;
|
||||
let loss_val = inv_loss.to_scalar::<f32>()?;
|
||||
|
||||
// Invariance loss should be very close to zero for identical inputs
|
||||
assert!(loss_val < 1e-6);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Test high variance inputs have low variance loss
|
||||
#[test]
|
||||
fn test_high_variance_low_loss() -> Result<()> {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let batch_size = 100;
|
||||
let feature_dim = 256;
|
||||
let gamma = 1.0;
|
||||
let epsilon = 1e-4;
|
||||
|
||||
// Create high variance input (scale up random values)
|
||||
let y_base = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;
|
||||
let scale = Tensor::full(y_base.shape(), 3.0, DType::F32, &device)?;
|
||||
let y = y_base.mul(&scale)?;
|
||||
|
||||
let var_loss = compute_variance_loss(&y, gamma, epsilon)?;
|
||||
let loss_val = var_loss.to_scalar::<f32>()?;
|
||||
|
||||
// High variance should result in low variance loss (hinge = max(0, γ - std))
|
||||
assert!(loss_val < 0.1);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Test uncorrelated features have low covariance loss
|
||||
#[test]
|
||||
fn test_uncorrelated_features_low_covariance_loss() -> Result<()> {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let batch_size = 100;
|
||||
let feature_dim = 10; // Small for controlled test
|
||||
let epsilon = 1e-4;
|
||||
|
||||
// Create orthogonal/uncorrelated features manually
|
||||
let mut y_data = vec![0.0; batch_size * feature_dim];
|
||||
|
||||
// Fill each feature dimension independently
|
||||
for i in 0..feature_dim {
|
||||
for b in 0..batch_size {
|
||||
y_data[b * feature_dim + i] = (b as f32 + i as f32 * 100.0).sin();
|
||||
}
|
||||
}
|
||||
|
||||
let y = Tensor::from_slice(&y_data, vec![batch_size, feature_dim], DType::F32, &device)?;
|
||||
|
||||
let cov_loss = compute_covariance_loss(&y, epsilon)?;
|
||||
let loss_val = cov_loss.to_scalar::<f32>()?;
|
||||
|
||||
// Uncorrelated features should have relatively low covariance loss
|
||||
assert!(loss_val >= 0.0); // Should be non-negative
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Test batch size consistency across loss computations
|
||||
#[test]
|
||||
fn test_batch_size_consistency() -> Result<()> {
|
||||
let device = Device::cuda(0).unwrap_or(Device::default());
|
||||
let feature_dim = 512;
|
||||
let epsilon = 1e-4;
|
||||
let gamma = 1.0;
|
||||
|
||||
// Test different batch sizes
|
||||
for batch_size in [1, 8, 32, 128] {
|
||||
let y1 = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;
|
||||
let y2 = Tensor::randn(vec![batch_size, feature_dim], DType::F32, &device)?;
|
||||
|
||||
// All loss computations should work regardless of batch size
|
||||
let inv_loss = compute_invariance_loss(&y1, &y2)?;
|
||||
let var_loss = compute_variance_loss(&y1, gamma, epsilon)?;
|
||||
let cov_loss = compute_covariance_loss(&y1, epsilon)?;
|
||||
|
||||
assert_eq!(inv_loss.shape(), &[]);
|
||||
assert_eq!(var_loss.shape(), &[]);
|
||||
assert_eq!(cov_loss.shape(), &[]);
|
||||
|
||||
assert!(inv_loss.to_scalar::<f32>()? >= 0.0);
|
||||
assert!(var_loss.to_scalar::<f32>()? >= 0.0);
|
||||
assert!(cov_loss.to_scalar::<f32>()? >= 0.0);
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user