Struct MultiTaskLearningTrainer 

pub struct MultiTaskLearningTrainer<M: Model> {
    pub base_model: M,
    pub task_heads: HashMap<String, TaskHead>,
    pub config: MTLConfig,
    pub task_weights: HashMap<String, f32>,
    pub task_performance: HashMap<String, Vec<f32>>,
    pub step_counter: usize,
    pub scheduler_state: TaskSchedulerState,
    pub gradient_stats: HashMap<String, GradientStats>,
}

Multi-task learning trainer

Fields

base_model: M

Base model (shared layers)

task_heads: HashMap<String, TaskHead>

Task-specific heads

config: MTLConfig

Configuration

task_weights: HashMap<String, f32>

Task losses and weights

task_performance: HashMap<String, Vec<f32>>

Task performance history

step_counter: usize

Current training step

scheduler_state: TaskSchedulerState

Task scheduling state

gradient_stats: HashMap<String, GradientStats>

Gradient statistics for balancing

Implementations

impl<M: Model<Input = Tensor, Output = Tensor>> MultiTaskLearningTrainer<M>

pub fn new(base_model: M, config: MTLConfig) -> Result<Self>

Create a new multi-task learning trainer

pub fn train_multi_task_step( &mut self, task_data: &HashMap<String, TaskBatch>, ) -> Result<MultiTaskOutput>

Train on multiple tasks for one step

pub fn evaluate_all_tasks( &self, test_data: &HashMap<String, TaskBatch>, ) -> Result<MultiTaskEvaluation>

Evaluate all tasks

pub fn get_mtl_stats(&self) -> MTLStats

Get multi-task learning statistics