335 lines
9.7 KiB
Rust
335 lines
9.7 KiB
Rust
//! Tests for knowledge distillation functionality
|
|
//!
|
|
//! TDD: Define expected behavior for teacher-student model compression
|
|
#![cfg(feature = "disabled_tests")]
|
|
|
|
use rtx_compress::{
|
|
Result,
|
|
distillation::{DistillationConfig, DistillationLoss, DistillationMethod, KnowledgeDistiller},
|
|
};
|
|
use rtx_tensor::{Device, Tensor};
|
|
|
|
#[test]
|
|
fn test_distillation_config_creation() {
|
|
let config = DistillationConfig::new(
|
|
DistillationMethod::ResponseBased {
|
|
temperature: 4.0,
|
|
alpha: 0.7,
|
|
},
|
|
DistillationLoss::KullbackLeibler,
|
|
);
|
|
|
|
assert!(
|
|
matches!(config.method, DistillationMethod::ResponseBased { temperature, alpha } if temperature == 4.0 && alpha == 0.7)
|
|
);
|
|
assert_eq!(config.loss_function, DistillationLoss::KullbackLeibler);
|
|
}
|
|
|
|
#[test]
|
|
fn test_basic_knowledge_distillation() {
|
|
let device = Device::cpu();
|
|
|
|
let config = DistillationConfig::new(
|
|
DistillationMethod::ResponseBased {
|
|
temperature: 3.0,
|
|
alpha: 0.9,
|
|
},
|
|
DistillationLoss::KullbackLeibler,
|
|
);
|
|
|
|
let distiller = KnowledgeDistiller::new(config).unwrap();
|
|
|
|
// Teacher logits (from larger model)
|
|
let teacher_logits = Tensor::from_data(
|
|
vec![2.0, 1.0, 0.1, -1.0], // Confident predictions
|
|
vec![1, 4],
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
// Student logits (from smaller model)
|
|
let student_logits = Tensor::from_data(
|
|
vec![1.5, 0.8, 0.2, -0.5], // Less confident
|
|
vec![1, 4],
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
// Compute distillation loss
|
|
let loss = distiller
|
|
.compute_loss(&teacher_logits, &student_logits)
|
|
.unwrap();
|
|
|
|
// Loss should be a scalar
|
|
assert_eq!(loss.shape().dims(), &[1]);
|
|
|
|
// Loss should be positive
|
|
let loss_val = loss.to_scalar::<f32>().unwrap();
|
|
assert!(loss_val > 0.0, "Distillation loss should be positive");
|
|
}
|
|
|
|
#[test]
|
|
fn test_temperature_scaling() {
|
|
let device = Device::cpu();
|
|
|
|
// Test different temperature values
|
|
for temp in [1.0, 3.0, 5.0, 10.0] {
|
|
let config = DistillationConfig::new(
|
|
DistillationMethod::ResponseBased {
|
|
temperature: temp,
|
|
alpha: 1.0,
|
|
},
|
|
DistillationLoss::KullbackLeibler,
|
|
);
|
|
|
|
let distiller = KnowledgeDistiller::new(config).unwrap();
|
|
|
|
let logits = Tensor::from_data(vec![10.0, 5.0, 1.0, 0.0], vec![1, 4], &device).unwrap();
|
|
|
|
// Apply temperature scaling
|
|
let scaled = distiller.apply_temperature(&logits).unwrap();
|
|
let scaled_data = scaled.to_vec().unwrap();
|
|
|
|
// Higher temperature should make distribution more uniform
|
|
if temp > 1.0 {
|
|
// Check that values are scaled down
|
|
assert!(scaled_data[0] < 10.0 / temp + 0.1);
|
|
}
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_soft_target_generation() {
|
|
let device = Device::cpu();
|
|
|
|
let config = DistillationConfig::default();
|
|
let distiller = KnowledgeDistiller::new(config).unwrap();
|
|
|
|
// Teacher logits
|
|
let teacher_logits = Tensor::from_data(
|
|
vec![
|
|
3.0, 1.0, -1.0, -2.0, // Clear preference for first class
|
|
-1.0, 2.0, 1.0, -2.0, // Preference for second class
|
|
],
|
|
vec![2, 4],
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
// Generate soft targets
|
|
let soft_targets = distiller.generate_soft_targets(&teacher_logits).unwrap();
|
|
|
|
// Check shape
|
|
assert_eq!(soft_targets.shape().dims(), teacher_logits.shape().dims());
|
|
|
|
// Check that soft targets sum to 1 (probability distribution)
|
|
let soft_data = soft_targets.to_vec().unwrap();
|
|
for batch in 0..2 {
|
|
let start = batch * 4;
|
|
let end = start + 4;
|
|
let sum: f32 = soft_data[start..end].iter().sum();
|
|
assert!((sum - 1.0).abs() < 1e-5, "Sum should be 1.0, got {}", sum);
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn test_attention_transfer() {
|
|
let device = Device::cpu();
|
|
|
|
use rtx_compress::distillation::AttentionTransferConfig;
|
|
let config = DistillationConfig::new(
|
|
DistillationMethod::AttentionTransfer {
|
|
attention_config: AttentionTransferConfig::new(8, 16, false, 1.0),
|
|
transfer_heads_individually: false,
|
|
},
|
|
DistillationLoss::MeanSquaredError,
|
|
);
|
|
|
|
let distiller = KnowledgeDistiller::new(config).unwrap();
|
|
|
|
// Teacher attention maps [batch, heads, seq_len, seq_len]
|
|
let teacher_attention = Tensor::randn(&[2, 8, 16, 16], &device).unwrap();
|
|
|
|
// Student attention maps (fewer heads)
|
|
let student_attention = Tensor::randn(&[2, 4, 16, 16], &device).unwrap();
|
|
|
|
// Compute attention transfer loss
|
|
let loss = distiller
|
|
.attention_transfer_loss(&teacher_attention, &student_attention)
|
|
.unwrap();
|
|
|
|
// Loss should be scalar
|
|
assert_eq!(loss.shape().dims(), &[1]);
|
|
|
|
let loss_val = loss.to_scalar::<f32>().unwrap();
|
|
assert!(
|
|
loss_val >= 0.0,
|
|
"Attention transfer loss should be non-negative"
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_feature_matching() {
|
|
let device = Device::cpu();
|
|
|
|
use rtx_compress::distillation::FeatureMatchingConfig;
|
|
let config = DistillationConfig::new(
|
|
DistillationMethod::FeatureBased {
|
|
feature_config: FeatureMatchingConfig::new(
|
|
vec![256, 512, 1024],
|
|
vec![0.3, 0.3, 0.4],
|
|
"mse".to_string(),
|
|
),
|
|
intermediate_matching: true,
|
|
},
|
|
DistillationLoss::MeanSquaredError,
|
|
);
|
|
|
|
let distiller = KnowledgeDistiller::new(config).unwrap();
|
|
|
|
// Teacher features from intermediate layers
|
|
let teacher_features = vec![
|
|
Tensor::randn(&[2, 256, 32, 32], &device).unwrap(), // Layer 1
|
|
Tensor::randn(&[2, 512, 16, 16], &device).unwrap(), // Layer 2
|
|
Tensor::randn(&[2, 1024, 8, 8], &device).unwrap(), // Layer 3
|
|
];
|
|
|
|
// Student features (potentially different dimensions)
|
|
let student_features = vec![
|
|
Tensor::randn(&[2, 128, 32, 32], &device).unwrap(), // Layer 1
|
|
Tensor::randn(&[2, 256, 16, 16], &device).unwrap(), // Layer 2
|
|
Tensor::randn(&[2, 512, 8, 8], &device).unwrap(), // Layer 3
|
|
];
|
|
|
|
// Compute feature matching loss
|
|
let loss = distiller
|
|
.feature_matching_loss(&teacher_features, &student_features)
|
|
.unwrap();
|
|
|
|
assert_eq!(loss.shape().dims(), &[1]);
|
|
let loss_val = loss.to_scalar::<f32>().unwrap();
|
|
assert!(loss_val >= 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_distillation_with_labels() {
|
|
let device = Device::cpu();
|
|
|
|
let config = DistillationConfig::new(
|
|
DistillationMethod::ResponseBased {
|
|
temperature: 4.0,
|
|
alpha: 0.7,
|
|
},
|
|
DistillationLoss::KullbackLeibler,
|
|
);
|
|
|
|
let distiller = KnowledgeDistiller::new(config).unwrap();
|
|
|
|
// Logits
|
|
let teacher_logits = Tensor::randn(&[4, 10], &device).unwrap();
|
|
let student_logits = Tensor::randn(&[4, 10], &device).unwrap();
|
|
|
|
// Ground truth labels
|
|
let labels = Tensor::from_data(vec![0.0, 1.0, 2.0, 3.0], vec![4], &device).unwrap();
|
|
|
|
// Compute combined loss (distillation + hard targets)
|
|
let loss = distiller
|
|
.compute_combined_loss(&teacher_logits, &student_logits, &labels)
|
|
.unwrap();
|
|
|
|
assert_eq!(loss.shape().dims(), &[1]);
|
|
let loss_val = loss.to_scalar::<f32>().unwrap();
|
|
assert!(loss_val > 0.0);
|
|
}
|
|
|
|
#[test]
|
|
fn test_student_compression_ratio() {
|
|
let device = Device::cpu();
|
|
|
|
let config = DistillationConfig::default();
|
|
let distiller = KnowledgeDistiller::new(config).unwrap();
|
|
|
|
// Teacher model size (parameters)
|
|
let teacher_params = 100_000_000; // 100M parameters
|
|
|
|
// Student model size
|
|
let student_params = 10_000_000; // 10M parameters
|
|
|
|
let compression_ratio = distiller.compute_compression_ratio(teacher_params, student_params);
|
|
|
|
assert!(
|
|
(compression_ratio - 10.0).abs() < 1e-6,
|
|
"Expected 10.0, got {}",
|
|
compression_ratio
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_gradual_distillation() {
|
|
let device = Device::cpu();
|
|
|
|
// Test temperature annealing during training
|
|
let config = DistillationConfig::new_with_adaptive_temperature(
|
|
DistillationMethod::ResponseBased {
|
|
temperature: 4.0,
|
|
alpha: 0.7,
|
|
},
|
|
DistillationLoss::KullbackLeibler,
|
|
true, // adaptive
|
|
0.001, // rate
|
|
);
|
|
let mut distiller = KnowledgeDistiller::new(config).unwrap();
|
|
|
|
// Check temperature at different steps
|
|
let temp_0 = distiller.get_temperature_at_step(0);
|
|
assert!((temp_0 - 4.0).abs() < 1e-6, "Expected 4.0, got {}", temp_0);
|
|
|
|
let temp_500 = distiller.get_temperature_at_step(500);
|
|
assert!(temp_500 < 4.0 && temp_500 > 1.0);
|
|
|
|
let temp_999 = distiller.get_temperature_at_step(999);
|
|
assert!(
|
|
(temp_999 - 1.0).abs() < 0.1,
|
|
"Expected ~1.0, got {}",
|
|
temp_999
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn test_distillation_metrics() {
|
|
let device = Device::cpu();
|
|
|
|
let config = DistillationConfig::default();
|
|
let distiller = KnowledgeDistiller::new(config).unwrap();
|
|
|
|
// Simulate teacher and student predictions
|
|
let teacher_preds = Tensor::from_data(
|
|
vec![0.0, 1.0, 2.0, 1.0], // Teacher predictions
|
|
vec![4],
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
let student_preds = Tensor::from_data(
|
|
vec![0.0, 1.0, 1.0, 1.0], // Student predictions (one mistake)
|
|
vec![4],
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
let labels = Tensor::from_data(
|
|
vec![0.0, 1.0, 2.0, 1.0], // Ground truth
|
|
vec![4],
|
|
&device,
|
|
)
|
|
.unwrap();
|
|
|
|
let metrics = distiller
|
|
.compute_metrics(&teacher_preds, &student_preds, &labels)
|
|
.unwrap();
|
|
|
|
assert_eq!(metrics.teacher_accuracy, 1.0); // Teacher is perfect
|
|
assert_eq!(metrics.student_accuracy, 0.75); // Student has 75% accuracy
|
|
assert_eq!(metrics.agreement_rate, 0.75); // They agree on 3/4 samples
|
|
}
|