461 lines
13 KiB
Rust
461 lines
13 KiB
Rust
//! Standalone Mean Teacher test
|
|
//!
|
|
//! Tests the Mean Teacher implementation independently
|
|
|
|
use std::sync::Arc;
|
|
use parking_lot::RwLock;
|
|
|
|
// 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();
|
|
let data: Vec<f32> = (0..size).map(|_| rand::random::<f32>() - 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, Box<dyn std::error::Error>> {
|
|
if self.shape != other.shape {
|
|
return Err("Shape mismatch".into());
|
|
}
|
|
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, Box<dyn std::error::Error>> {
|
|
if self.shape != other.shape {
|
|
return Err("Shape mismatch".into());
|
|
}
|
|
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, Box<dyn std::error::Error>> {
|
|
if self.shape != other.shape {
|
|
return Err("Shape mismatch".into());
|
|
}
|
|
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, Box<dyn std::error::Error>> {
|
|
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, Box<dyn std::error::Error>> {
|
|
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, Box<dyn std::error::Error>> {
|
|
if self.shape.len() != 2 || other.shape.len() != 2 {
|
|
return Err("Only 2D tensors supported for matmul".into());
|
|
}
|
|
if self.shape[1] != other.shape[0] {
|
|
return Err("Matrix dimension mismatch".into());
|
|
}
|
|
|
|
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, Box<dyn std::error::Error>> {
|
|
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, Box<dyn std::error::Error>>;
|
|
|
|
/// 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: Arc<RwLock<SimpleTensor>>,
|
|
bias: Arc<RwLock<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: Arc::new(RwLock::new(weight)),
|
|
bias: Arc::new(RwLock::new(bias)),
|
|
})
|
|
}
|
|
|
|
fn forward(&self, input: &SimpleTensor) -> Result<SimpleTensor> {
|
|
let weight = self.weight.read();
|
|
let bias = self.bias.read();
|
|
|
|
let output = input.matmul(&*weight)?;
|
|
output.add(&*bias)
|
|
}
|
|
|
|
fn get_weight(&self) -> SimpleTensor {
|
|
self.weight.read().clone()
|
|
}
|
|
|
|
fn get_bias(&self) -> SimpleTensor {
|
|
self.bias.read().clone()
|
|
}
|
|
|
|
fn set_weight(&self, new_weight: SimpleTensor) -> Result<()> {
|
|
*self.weight.write() = new_weight;
|
|
Ok(())
|
|
}
|
|
|
|
fn set_bias(&self, new_bias: SimpleTensor) -> Result<()> {
|
|
*self.bias.write() = new_bias;
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
/// 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
|
|
}
|
|
|
|
/// Set training mode
|
|
pub fn train(&mut self) {
|
|
self.training = true;
|
|
}
|
|
|
|
/// Set evaluation mode
|
|
pub fn eval(&mut self) {
|
|
self.training = false;
|
|
}
|
|
|
|
/// 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)
|
|
}
|
|
|
|
/// Get all parameters
|
|
pub fn parameters(&self) -> Vec<SimpleTensor> {
|
|
self.layers.iter().flat_map(|layer| {
|
|
vec![layer.get_weight(), layer.get_bias()]
|
|
}).collect()
|
|
}
|
|
|
|
/// Set parameters
|
|
pub fn set_parameters(&self, params: Vec<SimpleTensor>) -> Result<()> {
|
|
assert_eq!(params.len(), self.layers.len() * 2);
|
|
|
|
for (i, layer) in self.layers.iter().enumerate() {
|
|
layer.set_weight(params[i * 2].clone())?;
|
|
layer.set_bias(params[i * 2 + 1].clone())?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
}
|
|
|
|
/// 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> {
|
|
// EMA update: teacher = decay * teacher + (1 - decay) * student
|
|
let updated = teacher_param.mul_scalar(self.decay)?
|
|
.add(&student_param.mul_scalar(1.0 - self.decay)?)?;
|
|
|
|
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> {
|
|
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 {
|
|
// Sigmoid ramp-up function
|
|
let progress = epoch as f32 / self.rampup_epochs as f32;
|
|
let sigmoid = 1.0 / (1.0 + (-5.0 * (progress - 0.5)).exp());
|
|
sigmoid * 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]);
|
|
let output = student.forward(&input)?;
|
|
assert_eq!(output.shape(), &[2, 3]);
|
|
println!("✅ Forward pass test passed");
|
|
|
|
// 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]);
|
|
let updated = ema.update(&teacher_param, &student_param)?;
|
|
assert_eq!(updated.shape(), &[5, 5]);
|
|
println!("✅ EMA updater test passed");
|
|
|
|
// Test 5: Consistency loss
|
|
let pred1 = SimpleTensor::randn(vec![4, 3]);
|
|
let pred2 = SimpleTensor::randn(vec![4, 3]);
|
|
let loss = compute_consistency_loss(&pred1, &pred2)?;
|
|
assert_eq!(loss.shape(), &[]);
|
|
println!("✅ Consistency loss test passed");
|
|
|
|
// Test 6: Ramp-up scheduler
|
|
let scheduler = ConsistencyRampUp::new(5, 100.0);
|
|
assert_eq!(scheduler.get_weight(0), 0.0);
|
|
assert!((scheduler.get_weight(5) - 100.0).abs() < 1e-6);
|
|
println!("✅ Ramp-up scheduler test passed");
|
|
|
|
println!("🎉 All tests passed! Mean Teacher implementation is working correctly.");
|
|
|
|
Ok(())
|
|
} |