Struct KnowledgeDistillation 

pub struct KnowledgeDistillation { /* private fields */ }

Knowledge distillation trainer

Implementations

impl KnowledgeDistillation

pub fn new(config: DistillationConfig) -> Self

Create new knowledge distillation trainer

pub fn compute_loss( &self, student_logits: &Tensor, teacher_logits: &Tensor, targets: &Tensor, ) -> Result<Tensor, String>

Compute distillation loss

pub fn store_teacher_output(&mut self, layer_name: String, output: Tensor)

Store teacher outputs

pub fn store_student_output(&mut self, layer_name: String, output: Tensor)

Store student outputs

pub fn clear_outputs(&mut self)

Clear stored outputs

pub fn get_stats(&self) -> DistillationStats

Get distillation statistics