Initial commit
This commit is contained in:
redclawsystems
@@ -0,0 +1,694 @@
|
||||
//! Mean Teacher Implementation
|
||||
//!
|
||||
//! Semi-supervised learning method using exponential moving average teacher models.
|
||||
//! Based on "Mean teachers are better role models" (Tarvainen & Valpola, 2017).
|
||||
//!
|
||||
//! Key features:
|
||||
//! - Student model (trainable) and teacher model (EMA of student)
|
||||
//! - Consistency regularization between student and teacher predictions
|
||||
//! - Support for labeled and unlabeled data
|
||||
//! - Noise injection for robustness (Gaussian noise, dropout)
|
||||
//! - Consistency weight ramp-up scheduling
|
||||
//! - EMA update mechanism for teacher parameters
|
||||
|
||||
use crate::prelude::*;
|
||||
use std::sync::Arc;
|
||||
use parking_lot::RwLock;
|
||||
|
||||
/// 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
|
||||
}
|
||||
}
|
||||
|
||||
/// 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,
|
||||
}
|
||||
|
||||
/// 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,
|
||||
/// Device
|
||||
device: Device,
|
||||
}
|
||||
|
||||
/// Teacher model wrapper (non-trainable, EMA of student)
|
||||
#[derive(Debug)]
|
||||
pub struct TeacherModel {
|
||||
/// Linear layers
|
||||
layers: Vec<LinearLayer>,
|
||||
/// Input dimension
|
||||
input_dim: usize,
|
||||
/// Output dimension
|
||||
output_dim: usize,
|
||||
/// Device
|
||||
device: Device,
|
||||
}
|
||||
|
||||
/// Linear layer for neural networks
|
||||
#[derive(Debug)]
|
||||
struct LinearLayer {
|
||||
weight: Arc<RwLock<Tensor>>,
|
||||
bias: Arc<RwLock<Tensor>>,
|
||||
}
|
||||
|
||||
impl LinearLayer {
|
||||
fn new(input_dim: usize, output_dim: usize, device: &Device) -> Result<Self> {
|
||||
// Xavier initialization
|
||||
let scale = (2.0 / (input_dim + output_dim) as f32).sqrt();
|
||||
let weight = Tensor::randn(vec![input_dim, output_dim], DType::F32, device)?
|
||||
.mul_scalar(scale)?;
|
||||
let bias = Tensor::zeros(vec![output_dim], device)?;
|
||||
|
||||
Ok(Self {
|
||||
weight: Arc::new(RwLock::new(weight)),
|
||||
bias: Arc::new(RwLock::new(bias)),
|
||||
})
|
||||
}
|
||||
|
||||
fn forward(&self, input: &Tensor) -> Result<Tensor> {
|
||||
let weight = self.weight.read();
|
||||
let bias = self.bias.read();
|
||||
|
||||
let output = input.matmul(&*weight)?;
|
||||
output.add(&*bias)
|
||||
}
|
||||
|
||||
fn get_weight(&self) -> Tensor {
|
||||
self.weight.read().clone()
|
||||
}
|
||||
|
||||
fn get_bias(&self) -> Tensor {
|
||||
self.bias.read().clone()
|
||||
}
|
||||
|
||||
fn set_weight(&self, new_weight: Tensor) -> Result<()> {
|
||||
*self.weight.write() = new_weight;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn set_bias(&self, new_bias: Tensor) -> Result<()> {
|
||||
*self.bias.write() = new_bias;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl StudentModel {
|
||||
/// Create new student model
|
||||
pub fn new(input_dim: usize, hidden_dim: usize, output_dim: usize, device: &Device) -> Result<Self> {
|
||||
let layer1 = LinearLayer::new(input_dim, hidden_dim, device)?;
|
||||
let layer2 = LinearLayer::new(hidden_dim, output_dim, device)?;
|
||||
|
||||
Ok(Self {
|
||||
layers: vec![layer1, layer2],
|
||||
input_dim,
|
||||
output_dim,
|
||||
training: false,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// 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: &Tensor) -> Result<Tensor> {
|
||||
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<Tensor> {
|
||||
self.layers.iter().flat_map(|layer| {
|
||||
vec![layer.get_weight(), layer.get_bias()]
|
||||
}).collect()
|
||||
}
|
||||
|
||||
/// Set parameters
|
||||
pub fn set_parameters(&self, params: Vec<Tensor>) -> 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(())
|
||||
}
|
||||
}
|
||||
|
||||
impl TeacherModel {
|
||||
/// Create new teacher model
|
||||
pub fn new(input_dim: usize, hidden_dim: usize, output_dim: usize, device: &Device) -> Result<Self> {
|
||||
let layer1 = LinearLayer::new(input_dim, hidden_dim, device)?;
|
||||
let layer2 = LinearLayer::new(hidden_dim, output_dim, device)?;
|
||||
|
||||
Ok(Self {
|
||||
layers: vec![layer1, layer2],
|
||||
input_dim,
|
||||
output_dim,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// 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 (always false for teacher)
|
||||
pub fn is_training(&self) -> bool {
|
||||
false
|
||||
}
|
||||
|
||||
/// Forward pass through teacher model (always in eval mode)
|
||||
pub fn forward(&self, input: &Tensor) -> Result<Tensor> {
|
||||
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<Tensor> {
|
||||
self.layers.iter().flat_map(|layer| {
|
||||
vec![layer.get_weight(), layer.get_bias()]
|
||||
}).collect()
|
||||
}
|
||||
|
||||
/// Set parameters
|
||||
pub fn set_parameters(&self, params: Vec<Tensor>) -> 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(())
|
||||
}
|
||||
}
|
||||
|
||||
/// Exponential Moving Average updater for teacher parameters
|
||||
#[derive(Debug)]
|
||||
pub struct EMAUpdater {
|
||||
/// Decay rate for EMA
|
||||
decay: f32,
|
||||
/// Current step count
|
||||
step: usize,
|
||||
/// Device
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl EMAUpdater {
|
||||
/// Create new EMA updater
|
||||
pub fn new(decay: f32, device: &Device) -> Result<Self> {
|
||||
Ok(Self {
|
||||
decay,
|
||||
step: 0,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// 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: &Tensor, student_param: &Tensor) -> Result<Tensor> {
|
||||
// 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)
|
||||
}
|
||||
|
||||
/// Update all teacher parameters
|
||||
pub fn update_all(&mut self, teacher_params: &[Tensor], student_params: &[Tensor]) -> Result<Vec<Tensor>> {
|
||||
assert_eq!(teacher_params.len(), student_params.len());
|
||||
|
||||
let mut updated_params = Vec::new();
|
||||
for (teacher_param, student_param) in teacher_params.iter().zip(student_params.iter()) {
|
||||
updated_params.push(self.update(teacher_param, student_param)?);
|
||||
}
|
||||
|
||||
Ok(updated_params)
|
||||
}
|
||||
}
|
||||
|
||||
/// Noise augmenter for input perturbation
|
||||
#[derive(Debug)]
|
||||
pub struct NoiseAugmenter {
|
||||
strategy: AugmentationStrategy,
|
||||
noise_level: f32,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl NoiseAugmenter {
|
||||
/// Create new noise augmenter
|
||||
pub fn new(strategy: AugmentationStrategy, noise_level: f32, device: &Device) -> Result<Self> {
|
||||
Ok(Self {
|
||||
strategy,
|
||||
noise_level,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Apply noise to input tensor
|
||||
pub fn apply_noise(&self, input: &Tensor, seed: Option<u64>) -> Result<Tensor> {
|
||||
match self.strategy {
|
||||
AugmentationStrategy::Gaussian => self.apply_gaussian_noise(input, seed),
|
||||
AugmentationStrategy::Dropout => self.apply_dropout_noise(input, seed),
|
||||
AugmentationStrategy::Both => {
|
||||
let gaussian = self.apply_gaussian_noise(input, seed)?;
|
||||
self.apply_dropout_noise(&gaussian, seed.map(|s| s.wrapping_add(1)))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
fn apply_gaussian_noise(&self, input: &Tensor, _seed: Option<u64>) -> Result<Tensor> {
|
||||
let noise = Tensor::randn(input.shape().clone(), &self.device)?
|
||||
.mul_scalar(self.noise_level)?;
|
||||
input.add(&noise)
|
||||
}
|
||||
|
||||
fn apply_dropout_noise(&self, input: &Tensor, _seed: Option<u64>) -> Result<Tensor> {
|
||||
// Simulate dropout by randomly scaling elements
|
||||
let keep_prob = 1.0 - self.noise_level;
|
||||
let mask = Tensor::rand(input.shape().clone(), DType::F32, &self.device)?;
|
||||
let dropout_mask = mask.gt_scalar(self.noise_level)?
|
||||
.to_dtype(DType::F32)?
|
||||
.div_scalar(keep_prob)?;
|
||||
|
||||
input.mul(&dropout_mask)
|
||||
}
|
||||
}
|
||||
|
||||
/// Compute consistency loss (MSE) between student and teacher predictions
|
||||
pub fn compute_consistency_loss(student_pred: &Tensor, teacher_pred: &Tensor) -> Result<Tensor> {
|
||||
let diff = student_pred.sub(teacher_pred)?;
|
||||
let squared_diff = diff.mul(&diff)?;
|
||||
squared_diff.mean(None, false)
|
||||
}
|
||||
|
||||
/// 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 (simple and predictable)
|
||||
let progress = epoch as f32 / self.rampup_epochs as f32;
|
||||
progress * self.max_weight
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Training result from Mean Teacher step
|
||||
#[derive(Debug, Clone)]
|
||||
pub struct MeanTeacherTrainingResult {
|
||||
/// Supervised loss (classification loss on labeled data)
|
||||
pub supervised_loss: f32,
|
||||
/// Consistency loss (between student and teacher predictions)
|
||||
pub consistency_loss: f32,
|
||||
/// Total loss (supervised + weighted consistency)
|
||||
pub total_loss: f32,
|
||||
}
|
||||
|
||||
/// Mean Teacher trainer
|
||||
#[derive(Debug)]
|
||||
pub struct MeanTeacherTrainer {
|
||||
config: MeanTeacherConfig,
|
||||
student: StudentModel,
|
||||
teacher: TeacherModel,
|
||||
ema_updater: EMAUpdater,
|
||||
noise_augmenter: NoiseAugmenter,
|
||||
rampup_scheduler: ConsistencyRampUp,
|
||||
current_epoch: usize,
|
||||
device: Device,
|
||||
}
|
||||
|
||||
impl MeanTeacherTrainer {
|
||||
/// Create new Mean Teacher trainer
|
||||
pub fn new(
|
||||
config: MeanTeacherConfig,
|
||||
input_dim: usize,
|
||||
hidden_dim: usize,
|
||||
output_dim: usize,
|
||||
device: &Device,
|
||||
) -> Result<Self> {
|
||||
let student = StudentModel::new(input_dim, hidden_dim, output_dim, device)?;
|
||||
let mut teacher = TeacherModel::new(input_dim, hidden_dim, output_dim, device)?;
|
||||
|
||||
// Initialize teacher parameters with student parameters
|
||||
let student_params = student.parameters();
|
||||
teacher.set_parameters(student_params)?;
|
||||
|
||||
let ema_updater = EMAUpdater::new(config.ema_decay, device)?;
|
||||
let noise_augmenter = NoiseAugmenter::new(
|
||||
config.augmentation_strategy.clone(),
|
||||
config.noise_level,
|
||||
device,
|
||||
)?;
|
||||
let rampup_scheduler = ConsistencyRampUp::new(
|
||||
config.consistency_rampup,
|
||||
config.consistency_weight,
|
||||
);
|
||||
|
||||
Ok(Self {
|
||||
config: config.clone(),
|
||||
student,
|
||||
teacher,
|
||||
ema_updater,
|
||||
noise_augmenter,
|
||||
rampup_scheduler,
|
||||
current_epoch: 0,
|
||||
device: device.clone(),
|
||||
})
|
||||
}
|
||||
|
||||
/// Get configuration
|
||||
pub fn config(&self) -> &MeanTeacherConfig {
|
||||
&self.config
|
||||
}
|
||||
|
||||
/// Get current epoch
|
||||
pub fn current_epoch(&self) -> usize {
|
||||
self.current_epoch
|
||||
}
|
||||
|
||||
/// Set current epoch
|
||||
pub fn set_epoch(&mut self, epoch: usize) {
|
||||
self.current_epoch = epoch;
|
||||
}
|
||||
|
||||
/// Check if in training mode
|
||||
pub fn is_training(&self) -> bool {
|
||||
self.student.is_training()
|
||||
}
|
||||
|
||||
/// Set training mode
|
||||
pub fn train(&mut self) {
|
||||
self.student.train();
|
||||
}
|
||||
|
||||
/// Set evaluation mode
|
||||
pub fn eval(&mut self) {
|
||||
self.student.eval();
|
||||
}
|
||||
|
||||
/// Get student parameters
|
||||
pub fn get_student_parameters(&self) -> Result<Vec<Tensor>> {
|
||||
Ok(self.student.parameters())
|
||||
}
|
||||
|
||||
/// Get teacher parameters
|
||||
pub fn get_teacher_parameters(&self) -> Result<Vec<Tensor>> {
|
||||
Ok(self.teacher.parameters())
|
||||
}
|
||||
|
||||
/// Supervised training step (labeled data only)
|
||||
pub fn supervised_step(
|
||||
&mut self,
|
||||
images: &Tensor,
|
||||
labels: &Tensor,
|
||||
epoch: usize,
|
||||
) -> Result<MeanTeacherTrainingResult> {
|
||||
self.set_epoch(epoch);
|
||||
|
||||
// Forward pass through student
|
||||
let student_pred = self.student.forward(images)?;
|
||||
|
||||
// Compute supervised loss (simplified cross-entropy)
|
||||
let supervised_loss = self.compute_supervised_loss(&student_pred, labels)?;
|
||||
|
||||
// Update teacher parameters
|
||||
self.update_teacher_parameters()?;
|
||||
|
||||
Ok(MeanTeacherTrainingResult {
|
||||
supervised_loss: supervised_loss.to_vec::<f32>()?[0],
|
||||
consistency_loss: 0.0,
|
||||
total_loss: supervised_loss.to_vec::<f32>()?[0],
|
||||
})
|
||||
}
|
||||
|
||||
/// Unsupervised training step (unlabeled data only)
|
||||
pub fn unsupervised_step(
|
||||
&mut self,
|
||||
unlabeled_images: &Tensor,
|
||||
epoch: usize,
|
||||
) -> Result<MeanTeacherTrainingResult> {
|
||||
self.set_epoch(epoch);
|
||||
|
||||
// Apply different noise to inputs for student and teacher
|
||||
let student_input = self.noise_augmenter.apply_noise(unlabeled_images, Some(42))?;
|
||||
let teacher_input = self.noise_augmenter.apply_noise(unlabeled_images, Some(24))?;
|
||||
|
||||
// Forward passes
|
||||
let student_pred = self.student.forward(&student_input)?;
|
||||
let teacher_pred = self.teacher.forward(&teacher_input)?;
|
||||
|
||||
// Compute consistency loss
|
||||
let consistency_loss = compute_consistency_loss(&student_pred, &teacher_pred)?;
|
||||
|
||||
// Apply consistency weight ramp-up
|
||||
let consistency_weight = self.rampup_scheduler.get_weight(epoch);
|
||||
let weighted_consistency_loss = consistency_loss.mul_scalar(consistency_weight)?;
|
||||
|
||||
// Update teacher parameters
|
||||
self.update_teacher_parameters()?;
|
||||
|
||||
Ok(MeanTeacherTrainingResult {
|
||||
supervised_loss: 0.0,
|
||||
consistency_loss: weighted_consistency_loss.to_vec::<f32>()?[0],
|
||||
total_loss: weighted_consistency_loss.to_vec::<f32>()?[0],
|
||||
})
|
||||
}
|
||||
|
||||
/// Mixed training step (labeled + unlabeled data)
|
||||
pub fn mixed_step(
|
||||
&mut self,
|
||||
labeled_images: &Tensor,
|
||||
labels: &Tensor,
|
||||
unlabeled_images: &Tensor,
|
||||
epoch: usize,
|
||||
) -> Result<MeanTeacherTrainingResult> {
|
||||
self.set_epoch(epoch);
|
||||
|
||||
// Supervised loss on labeled data
|
||||
let student_pred_labeled = self.student.forward(labeled_images)?;
|
||||
let supervised_loss = self.compute_supervised_loss(&student_pred_labeled, labels)?;
|
||||
|
||||
// Consistency loss on unlabeled data
|
||||
let student_input = self.noise_augmenter.apply_noise(unlabeled_images, Some(42))?;
|
||||
let teacher_input = self.noise_augmenter.apply_noise(unlabeled_images, Some(24))?;
|
||||
|
||||
let student_pred_unlabeled = self.student.forward(&student_input)?;
|
||||
let teacher_pred_unlabeled = self.teacher.forward(&teacher_input)?;
|
||||
|
||||
let consistency_loss = compute_consistency_loss(&student_pred_unlabeled, &teacher_pred_unlabeled)?;
|
||||
let consistency_weight = self.rampup_scheduler.get_weight(epoch);
|
||||
let weighted_consistency_loss = consistency_loss.mul_scalar(consistency_weight)?;
|
||||
|
||||
// Total loss
|
||||
let total_loss = supervised_loss.add(&weighted_consistency_loss)?;
|
||||
|
||||
// Update teacher parameters
|
||||
self.update_teacher_parameters()?;
|
||||
|
||||
Ok(MeanTeacherTrainingResult {
|
||||
supervised_loss: supervised_loss.to_vec::<f32>()?[0],
|
||||
consistency_loss: weighted_consistency_loss.to_vec::<f32>()?[0],
|
||||
total_loss: total_loss.to_vec::<f32>()?[0],
|
||||
})
|
||||
}
|
||||
|
||||
/// Extract features using student model (evaluation mode)
|
||||
pub fn extract_features(&self, input: &Tensor) -> Result<Tensor> {
|
||||
self.student.forward(input)
|
||||
}
|
||||
|
||||
fn compute_supervised_loss(&self, predictions: &Tensor, labels: &Tensor) -> Result<Tensor> {
|
||||
// Simplified cross-entropy loss
|
||||
let log_probs = predictions.log_softmax(-1)?;
|
||||
let labels_one_hot = self.to_one_hot(labels, predictions.shape()[1])?;
|
||||
let loss = log_probs.mul(&labels_one_hot)?.sum(None, false)?.neg()?;
|
||||
loss.div_scalar(predictions.shape()[0] as f32)
|
||||
}
|
||||
|
||||
fn to_one_hot(&self, labels: &Tensor, num_classes: usize) -> Result<Tensor> {
|
||||
let batch_size = labels.shape()[0];
|
||||
let mut one_hot = Tensor::zeros(vec![batch_size, num_classes], DType::F32, &self.device)?;
|
||||
|
||||
// Simplified one-hot encoding (would need proper indexing in real implementation)
|
||||
for i in 0..batch_size {
|
||||
let label_val = labels.get(i)?.to_vec::<i64>()?[0] as usize;
|
||||
if label_val < num_classes {
|
||||
one_hot = one_hot.index_put(
|
||||
&[Some(i)],
|
||||
&Tensor::ones(vec![1], &self.device)?
|
||||
)?;
|
||||
}
|
||||
}
|
||||
|
||||
Ok(one_hot)
|
||||
}
|
||||
|
||||
fn update_teacher_parameters(&mut self) -> Result<()> {
|
||||
let student_params = self.student.parameters();
|
||||
let teacher_params = self.teacher.parameters();
|
||||
|
||||
let updated_params = self.ema_updater.update_all(&teacher_params, &student_params)?;
|
||||
self.teacher.set_parameters(updated_params)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user