465 lines
13 KiB
Rust
465 lines
13 KiB
Rust
//! Simple Mean Teacher test
|
|
//!
|
|
//! Tests the Mean Teacher implementation independently
|
|
|
|
// Minimal tensor implementation for testing
|
|
#[derive(Debug, Clone)]
|
|
pub struct SimpleTensor {
|
|
data: Vec<f32>,
|
|
shape: Vec<usize>,
|
|
}
|
|
|
|
impl SimpleTensor {
|
|
pub fn randn(shape: Vec<usize>) -> Self {
|
|
let size: usize = shape.iter().product();
|
|
// Simple pseudo-random using linear congruential generator
|
|
let mut seed = 42u64;
|
|
let data: Vec<f32> = (0..size).map(|_| {
|
|
seed = seed.wrapping_mul(1103515245).wrapping_add(12345);
|
|
((seed >> 16) as f32 / 65536.0) - 0.5
|
|
}).collect();
|
|
Self { data, shape }
|
|
}
|
|
|
|
pub fn zeros(shape: Vec<usize>) -> Self {
|
|
let size: usize = shape.iter().product();
|
|
let data = vec![0.0; size];
|
|
Self { data, shape }
|
|
}
|
|
|
|
pub fn ones(shape: Vec<usize>) -> Self {
|
|
let size: usize = shape.iter().product();
|
|
let data = vec![1.0; size];
|
|
Self { data, shape }
|
|
}
|
|
|
|
pub fn shape(&self) -> &[usize] {
|
|
&self.shape
|
|
}
|
|
|
|
pub fn add(&self, other: &Self) -> Result<Self> {
|
|
if self.shape != other.shape {
|
|
return Err("Shape mismatch");
|
|
}
|
|
let data: Vec<f32> = self.data.iter().zip(&other.data)
|
|
.map(|(a, b)| a + b)
|
|
.collect();
|
|
Ok(Self { data, shape: self.shape.clone() })
|
|
}
|
|
|
|
pub fn sub(&self, other: &Self) -> Result<Self> {
|
|
if self.shape != other.shape {
|
|
return Err("Shape mismatch");
|
|
}
|
|
let data: Vec<f32> = self.data.iter().zip(&other.data)
|
|
.map(|(a, b)| a - b)
|
|
.collect();
|
|
Ok(Self { data, shape: self.shape.clone() })
|
|
}
|
|
|
|
pub fn mul(&self, other: &Self) -> Result<Self> {
|
|
if self.shape != other.shape {
|
|
return Err("Shape mismatch");
|
|
}
|
|
let data: Vec<f32> = self.data.iter().zip(&other.data)
|
|
.map(|(a, b)| a * b)
|
|
.collect();
|
|
Ok(Self { data, shape: self.shape.clone() })
|
|
}
|
|
|
|
pub fn mul_scalar(&self, scalar: f32) -> Result<Self> {
|
|
let data: Vec<f32> = self.data.iter().map(|x| x * scalar).collect();
|
|
Ok(Self { data, shape: self.shape.clone() })
|
|
}
|
|
|
|
pub fn relu(&self) -> Result<Self> {
|
|
let data: Vec<f32> = self.data.iter().map(|x| x.max(0.0)).collect();
|
|
Ok(Self { data, shape: self.shape.clone() })
|
|
}
|
|
|
|
pub fn matmul(&self, other: &Self) -> Result<Self> {
|
|
if self.shape.len() != 2 || other.shape.len() != 2 {
|
|
return Err("Only 2D tensors supported for matmul");
|
|
}
|
|
if self.shape[1] != other.shape[0] {
|
|
return Err("Matrix dimension mismatch");
|
|
}
|
|
|
|
let m = self.shape[0];
|
|
let n = other.shape[1];
|
|
let k = self.shape[1];
|
|
|
|
let mut result = vec![0.0; m * n];
|
|
for i in 0..m {
|
|
for j in 0..n {
|
|
for l in 0..k {
|
|
result[i * n + j] += self.data[i * k + l] * other.data[l * n + j];
|
|
}
|
|
}
|
|
}
|
|
|
|
Ok(Self {
|
|
data: result,
|
|
shape: vec![m, n],
|
|
})
|
|
}
|
|
|
|
pub fn mean(&self) -> Result<Self> {
|
|
let sum: f32 = self.data.iter().sum();
|
|
let mean_val = sum / self.data.len() as f32;
|
|
Ok(Self {
|
|
data: vec![mean_val],
|
|
shape: vec![],
|
|
})
|
|
}
|
|
|
|
pub fn to_vec(&self) -> Vec<f32> {
|
|
self.data.clone()
|
|
}
|
|
}
|
|
|
|
type Result<T> = std::result::Result<T, &'static str>;
|
|
|
|
/// Augmentation strategy for noise injection
|
|
#[derive(Debug, Clone, PartialEq, Eq)]
|
|
pub enum AugmentationStrategy {
|
|
/// Gaussian noise only
|
|
Gaussian,
|
|
/// Dropout noise only
|
|
Dropout,
|
|
/// Both Gaussian noise and dropout
|
|
Both,
|
|
}
|
|
|
|
/// Mean Teacher configuration parameters
|
|
#[derive(Debug, Clone)]
|
|
pub struct MeanTeacherConfig {
|
|
/// Exponential moving average decay rate for teacher updates
|
|
pub ema_decay: f32,
|
|
/// Maximum weight for consistency loss
|
|
pub consistency_weight: f32,
|
|
/// Number of epochs for consistency weight ramp-up
|
|
pub consistency_rampup: usize,
|
|
/// Noise level for input augmentation
|
|
pub noise_level: f32,
|
|
/// Augmentation strategy (Gaussian, Dropout, or Both)
|
|
pub augmentation_strategy: AugmentationStrategy,
|
|
}
|
|
|
|
impl Default for MeanTeacherConfig {
|
|
fn default() -> Self {
|
|
Self {
|
|
ema_decay: 0.999,
|
|
consistency_weight: 100.0,
|
|
consistency_rampup: 5,
|
|
noise_level: 0.15,
|
|
augmentation_strategy: AugmentationStrategy::Gaussian,
|
|
}
|
|
}
|
|
}
|
|
|
|
impl MeanTeacherConfig {
|
|
/// Create new Mean Teacher configuration
|
|
pub fn new() -> Self {
|
|
Self::default()
|
|
}
|
|
|
|
/// Set EMA decay rate
|
|
pub fn with_ema_decay(mut self, ema_decay: f32) -> Self {
|
|
self.ema_decay = ema_decay;
|
|
self
|
|
}
|
|
|
|
/// Set consistency weight
|
|
pub fn with_consistency_weight(mut self, consistency_weight: f32) -> Self {
|
|
self.consistency_weight = consistency_weight;
|
|
self
|
|
}
|
|
|
|
/// Set consistency ramp-up epochs
|
|
pub fn with_consistency_rampup(mut self, consistency_rampup: usize) -> Self {
|
|
self.consistency_rampup = consistency_rampup;
|
|
self
|
|
}
|
|
|
|
/// Set noise level
|
|
pub fn with_noise_level(mut self, noise_level: f32) -> Self {
|
|
self.noise_level = noise_level;
|
|
self
|
|
}
|
|
|
|
/// Set augmentation strategy
|
|
pub fn with_augmentation_strategy(mut self, augmentation_strategy: AugmentationStrategy) -> Self {
|
|
self.augmentation_strategy = augmentation_strategy;
|
|
self
|
|
}
|
|
}
|
|
|
|
/// Linear layer for neural networks
|
|
#[derive(Debug)]
|
|
struct LinearLayer {
|
|
weight: SimpleTensor,
|
|
bias: SimpleTensor,
|
|
}
|
|
|
|
impl LinearLayer {
|
|
fn new(input_dim: usize, output_dim: usize) -> Result<Self> {
|
|
// Xavier initialization
|
|
let scale = (2.0 / (input_dim + output_dim) as f32).sqrt();
|
|
let weight = SimpleTensor::randn(vec![input_dim, output_dim]).mul_scalar(scale)?;
|
|
let bias = SimpleTensor::zeros(vec![output_dim]);
|
|
|
|
Ok(Self { weight, bias })
|
|
}
|
|
|
|
fn forward(&self, input: &SimpleTensor) -> Result<SimpleTensor> {
|
|
let output = input.matmul(&self.weight)?;
|
|
// Bias broadcasting: add bias to each row
|
|
let mut result_data = output.data.clone();
|
|
let batch_size = output.shape[0];
|
|
let output_dim = output.shape[1];
|
|
|
|
for i in 0..batch_size {
|
|
for j in 0..output_dim {
|
|
result_data[i * output_dim + j] += self.bias.data[j];
|
|
}
|
|
}
|
|
|
|
Ok(SimpleTensor {
|
|
data: result_data,
|
|
shape: output.shape.clone(),
|
|
})
|
|
}
|
|
|
|
fn get_weight(&self) -> &SimpleTensor {
|
|
&self.weight
|
|
}
|
|
|
|
fn get_bias(&self) -> &SimpleTensor {
|
|
&self.bias
|
|
}
|
|
}
|
|
|
|
/// Student model wrapper (trainable)
|
|
#[derive(Debug)]
|
|
pub struct StudentModel {
|
|
/// Linear layers
|
|
layers: Vec<LinearLayer>,
|
|
/// Input dimension
|
|
input_dim: usize,
|
|
/// Output dimension
|
|
output_dim: usize,
|
|
/// Training mode flag
|
|
training: bool,
|
|
}
|
|
|
|
impl StudentModel {
|
|
/// Create new student model
|
|
pub fn new(input_dim: usize, hidden_dim: usize, output_dim: usize) -> Result<Self> {
|
|
let layer1 = LinearLayer::new(input_dim, hidden_dim)?;
|
|
let layer2 = LinearLayer::new(hidden_dim, output_dim)?;
|
|
|
|
Ok(Self {
|
|
layers: vec![layer1, layer2],
|
|
input_dim,
|
|
output_dim,
|
|
training: false,
|
|
})
|
|
}
|
|
|
|
/// Get input dimension
|
|
pub fn input_dim(&self) -> usize {
|
|
self.input_dim
|
|
}
|
|
|
|
/// Get output dimension
|
|
pub fn output_dim(&self) -> usize {
|
|
self.output_dim
|
|
}
|
|
|
|
/// Check if in training mode
|
|
pub fn is_training(&self) -> bool {
|
|
self.training
|
|
}
|
|
|
|
/// Forward pass through student model
|
|
pub fn forward(&self, input: &SimpleTensor) -> Result<SimpleTensor> {
|
|
let mut x = input.clone();
|
|
|
|
// First layer with ReLU activation
|
|
x = self.layers[0].forward(&x)?;
|
|
x = x.relu()?;
|
|
|
|
// Output layer
|
|
x = self.layers[1].forward(&x)?;
|
|
|
|
Ok(x)
|
|
}
|
|
}
|
|
|
|
/// EMA updater for teacher parameters
|
|
#[derive(Debug)]
|
|
pub struct EMAUpdater {
|
|
/// Decay rate for EMA
|
|
decay: f32,
|
|
/// Current step count
|
|
step: usize,
|
|
}
|
|
|
|
impl EMAUpdater {
|
|
/// Create new EMA updater
|
|
pub fn new(decay: f32) -> Result<Self> {
|
|
Ok(Self {
|
|
decay,
|
|
step: 0,
|
|
})
|
|
}
|
|
|
|
/// Get decay rate
|
|
pub fn decay(&self) -> f32 {
|
|
self.decay
|
|
}
|
|
|
|
/// Get current step
|
|
pub fn step(&self) -> usize {
|
|
self.step
|
|
}
|
|
|
|
/// Update teacher parameter with student parameter
|
|
pub fn update(&mut self, teacher_param: &SimpleTensor, student_param: &SimpleTensor) -> Result<SimpleTensor> {
|
|
if teacher_param.shape() != student_param.shape() {
|
|
return Err("Parameter shape mismatch");
|
|
}
|
|
|
|
// EMA update: teacher = decay * teacher + (1 - decay) * student
|
|
let teacher_scaled = teacher_param.mul_scalar(self.decay)?;
|
|
let student_scaled = student_param.mul_scalar(1.0 - self.decay)?;
|
|
let updated = teacher_scaled.add(&student_scaled)?;
|
|
|
|
self.step += 1;
|
|
Ok(updated)
|
|
}
|
|
}
|
|
|
|
/// Compute consistency loss (MSE) between student and teacher predictions
|
|
pub fn compute_consistency_loss(student_pred: &SimpleTensor, teacher_pred: &SimpleTensor) -> Result<SimpleTensor> {
|
|
if student_pred.shape() != teacher_pred.shape() {
|
|
return Err("Prediction shape mismatch");
|
|
}
|
|
let diff = student_pred.sub(teacher_pred)?;
|
|
let squared_diff = diff.mul(&diff)?;
|
|
squared_diff.mean()
|
|
}
|
|
|
|
/// Consistency weight ramp-up scheduler
|
|
#[derive(Debug)]
|
|
pub struct ConsistencyRampUp {
|
|
rampup_epochs: usize,
|
|
max_weight: f32,
|
|
}
|
|
|
|
impl ConsistencyRampUp {
|
|
/// Create new consistency ramp-up scheduler
|
|
pub fn new(rampup_epochs: usize, max_weight: f32) -> Self {
|
|
Self {
|
|
rampup_epochs,
|
|
max_weight,
|
|
}
|
|
}
|
|
|
|
/// Get rampup epochs
|
|
pub fn rampup_epochs(&self) -> usize {
|
|
self.rampup_epochs
|
|
}
|
|
|
|
/// Get maximum weight
|
|
pub fn max_weight(&self) -> f32 {
|
|
self.max_weight
|
|
}
|
|
|
|
/// Get current consistency weight based on epoch
|
|
pub fn get_weight(&self, epoch: usize) -> f32 {
|
|
if epoch >= self.rampup_epochs {
|
|
self.max_weight
|
|
} else if self.rampup_epochs == 0 {
|
|
self.max_weight
|
|
} else {
|
|
// Linear ramp-up function (simpler and more predictable)
|
|
let progress = epoch as f32 / self.rampup_epochs as f32;
|
|
progress * self.max_weight
|
|
}
|
|
}
|
|
}
|
|
|
|
fn main() -> Result<()> {
|
|
println!("🧪 Testing Mean Teacher Implementation");
|
|
|
|
// Test 1: Configuration
|
|
let config = MeanTeacherConfig::default();
|
|
assert_eq!(config.ema_decay, 0.999);
|
|
assert_eq!(config.consistency_weight, 100.0);
|
|
println!("✅ Configuration test passed");
|
|
|
|
// Test 2: Student model
|
|
let student = StudentModel::new(10, 5, 3)?;
|
|
assert_eq!(student.input_dim(), 10);
|
|
assert_eq!(student.output_dim(), 3);
|
|
println!("✅ Student model test passed");
|
|
|
|
// Test 3: Forward pass
|
|
let input = SimpleTensor::randn(vec![2, 10]);
|
|
match student.forward(&input) {
|
|
Ok(output) => {
|
|
assert_eq!(output.shape(), &[2, 3]);
|
|
println!("✅ Forward pass test passed");
|
|
},
|
|
Err(e) => {
|
|
println!("❌ Forward pass test failed: {}", e);
|
|
println!("input shape: {:?}", input.shape());
|
|
return Err(e);
|
|
}
|
|
}
|
|
|
|
// Test 4: EMA updater
|
|
let mut ema = EMAUpdater::new(0.9)?;
|
|
let teacher_param = SimpleTensor::zeros(vec![5, 5]);
|
|
let student_param = SimpleTensor::ones(vec![5, 5]);
|
|
match ema.update(&teacher_param, &student_param) {
|
|
Ok(updated) => {
|
|
assert_eq!(updated.shape(), &[5, 5]);
|
|
println!("✅ EMA updater test passed");
|
|
},
|
|
Err(e) => {
|
|
println!("❌ EMA updater test failed: {}", e);
|
|
println!("teacher shape: {:?}, student shape: {:?}", teacher_param.shape(), student_param.shape());
|
|
return Err(e);
|
|
}
|
|
}
|
|
|
|
// Test 5: Consistency loss
|
|
let pred1 = SimpleTensor::randn(vec![4, 3]);
|
|
let pred2 = SimpleTensor::randn(vec![4, 3]);
|
|
match compute_consistency_loss(&pred1, &pred2) {
|
|
Ok(loss) => {
|
|
assert_eq!(loss.shape(), &[]);
|
|
println!("✅ Consistency loss test passed");
|
|
},
|
|
Err(e) => {
|
|
println!("❌ Consistency loss test failed: {}", e);
|
|
println!("pred1 shape: {:?}, pred2 shape: {:?}", pred1.shape(), pred2.shape());
|
|
return Err(e);
|
|
}
|
|
}
|
|
|
|
// Test 6: Ramp-up scheduler
|
|
let scheduler = ConsistencyRampUp::new(5, 100.0);
|
|
let weight_0 = scheduler.get_weight(0);
|
|
let weight_5 = scheduler.get_weight(5);
|
|
assert_eq!(weight_0, 0.0); // Should be exactly 0 with linear ramp-up
|
|
assert!((weight_5 - 100.0).abs() < 1e-6);
|
|
println!("✅ Ramp-up scheduler test passed");
|
|
|
|
println!("🎉 All tests passed! Mean Teacher implementation is working correctly.");
|
|
|
|
Ok(())
|
|
} |