knowledge_distillation.py

import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from tqdm import tqdm
import logging
import os

logger = logging.getLogger(__name__)

class DistillationTrainer:
    """
    Trainer for knowledge distillation from teacher model (BERT) to student model (LSTM)
    """
    def __init__(
        self, 
        teacher_model, 
        student_model,
        train_loader, 
        val_loader, 
        test_loader=None,
        temperature=2.0,
        alpha=0.5,  # Weight for distillation loss vs. regular loss
        lr=0.001,
        weight_decay=1e-5,
        device=None
    ):
        self.teacher_model = teacher_model
        self.student_model = student_model
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.test_loader = test_loader
        self.temperature = temperature
        self.alpha = alpha
        
        self.device = device if device else torch.device('cuda' if torch.cuda.is_available() else 'cpu')
        logger.info(f"Using device: {self.device}")
        
        # Move models to device
        self.teacher_model.to(self.device)
        self.student_model.to(self.device)
        
        # Set teacher model to evaluation mode
        self.teacher_model.eval()
        
        # Optimizer for student model
        self.optimizer = torch.optim.Adam(
            self.student_model.parameters(), 
            lr=lr, 
            weight_decay=weight_decay
        )
        
        # Learning rate scheduler
        self.scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(
            self.optimizer, mode='max', factor=0.5, patience=2, verbose=True
        )
        
        # Loss functions
        self.ce_loss = nn.CrossEntropyLoss()  # For hard targets
        
        # Tracking metrics
        self.best_val_f1 = 0.0
        self.best_model_state = None
    
    def distillation_loss(self, student_logits, teacher_logits, labels, temperature, alpha):
        """
        Compute the knowledge distillation loss
        
        Args:
            student_logits: Output from student model
            teacher_logits: Output from teacher model
            labels: Ground truth labels
            temperature: Temperature for softening probability distributions
            alpha: Weight for distillation loss vs. cross-entropy loss
            
        Returns:
            Combined loss
        """
        # Softmax with temperature for soft targets
        soft_targets = F.softmax(teacher_logits / temperature, dim=1)
        soft_prob = F.log_softmax(student_logits / temperature, dim=1)
        
        # Distillation loss (KL divergence)
        distill_loss = F.kl_div(soft_prob, soft_targets, reduction='batchmean') * (temperature ** 2)
        
        # Standard cross entropy with hard targets
        ce_loss = self.ce_loss(student_logits, labels)
        
        # Weighted combination of the two losses
        loss = alpha * distill_loss + (1 - alpha) * ce_loss
        
        return loss
    
    def train(self, epochs, save_path='best_distilled_model.pth'):
        """
        Train student model with knowledge distillation
        """
        logger.info(f"Starting distillation training for {epochs} epochs")
        logger.info(f"Temperature: {self.temperature}, Alpha: {self.alpha}")
        
        for epoch in range(epochs):
            self.student_model.train()
            train_loss = 0.0
            all_preds = []
            all_labels = []
            
            # Training loop
            train_iterator = tqdm(self.train_loader, desc=f"Epoch {epoch+1}/{epochs} [Train]")
            for batch in train_iterator:
                # Move batch to device
                input_ids = batch['input_ids'].to(self.device)
                attention_mask = batch['attention_mask'].to(self.device)
                labels = batch['label'].to(self.device)
                
                # Get teacher predictions (no grad needed for teacher)
                with torch.no_grad():
                    teacher_logits = self.teacher_model(
                        input_ids=input_ids,
                        attention_mask=attention_mask
                    )
                
                # Forward pass through student model
                student_logits = self.student_model(
                    input_ids=input_ids,
                    attention_mask=attention_mask
                )
                
                # Calculate distillation loss
                loss = self.distillation_loss(
                    student_logits, 
                    teacher_logits, 
                    labels, 
                    self.temperature, 
                    self.alpha
                )
                
                # Backward and optimize
                self.optimizer.zero_grad()
                loss.backward()
                torch.nn.utils.clip_grad_norm_(self.student_model.parameters(), 1.0)
                self.optimizer.step()
                
                train_loss += loss.item()
                
                # Calculate accuracy for progress tracking
                _, preds = torch.max(student_logits, 1)
                all_preds.extend(preds.cpu().tolist())
                all_labels.extend(labels.cpu().tolist())
                
                # Update progress bar
                train_iterator.set_postfix({'loss': f"{loss.item():.4f}"})
            
            # Calculate training metrics
            train_loss = train_loss / len(self.train_loader)
            train_acc = sum(1 for p, l in zip(all_preds, all_labels) if p == l) / len(all_preds)
            
            # Evaluate on validation set
            val_loss, val_acc, val_f1 = self.evaluate()
            
            # Update learning rate based on validation performance
            self.scheduler.step(val_f1)
            
            # Save best model
            if val_f1 > self.best_val_f1:
                self.best_val_f1 = val_f1
                self.best_model_state = self.student_model.state_dict().copy()
                torch.save({
                    'epoch': epoch,
                    'model_state_dict': self.student_model.state_dict(),
                    'optimizer_state_dict': self.optimizer.state_dict(),
                    'val_f1': val_f1,
                }, save_path)
                logger.info(f"New best model saved with validation F1: {val_f1:.4f}")
            
            logger.info(f"Epoch {epoch+1}/{epochs}: "
                      f"Train Loss: {train_loss:.4f}, Train Acc: {train_acc:.4f}, "
                      f"Val Loss: {val_loss:.4f}, Val Acc: {val_acc:.4f}, Val F1: {val_f1:.4f}")
        
        # Load best model for final evaluation
        if self.best_model_state is not None:
            self.student_model.load_state_dict(self.best_model_state)
            logger.info(f"Loaded best model with validation F1: {self.best_val_f1:.4f}")
        
        # Final evaluation on test set if provided
        if self.test_loader:
            test_loss, test_acc, test_f1 = self.evaluate(self.test_loader, "Test")
            logger.info(f"Test Loss: {test_loss:.4f}, Test Acc: {test_acc:.4f}, Test F1: {test_f1:.4f}")
    
    def evaluate(self, data_loader=None, phase="Validation"):
        """
        Evaluate the student model
        """
        if data_loader is None:
            data_loader = self.val_loader
        
        self.student_model.eval()
        eval_loss = 0.0
        all_preds = []
        all_labels = []
        
        with torch.no_grad():
            for batch in tqdm(data_loader, desc=f"[{phase}]"):
                input_ids = batch['input_ids'].to(self.device)
                attention_mask = batch['attention_mask'].to(self.device)
                labels = batch['label'].to(self.device)
                
                # Forward pass through student
                student_logits = self.student_model(
                    input_ids=input_ids,
                    attention_mask=attention_mask
                )
                
                # Calculate regular CE loss (no distillation during evaluation)
                loss = self.ce_loss(student_logits, labels)
                eval_loss += loss.item()
                
                # Get predictions
                _, preds = torch.max(student_logits, 1)
                all_preds.extend(preds.cpu().tolist())
                all_labels.extend(labels.cpu().tolist())
        
        # Calculate metrics
        eval_loss = eval_loss / len(data_loader)
        
        # Accuracy
        accuracy = sum(1 for p, l in zip(all_preds, all_labels) if p == l) / len(all_preds)
        
        # F1 score (macro-averaged)
        from sklearn.metrics import f1_score
        f1 = f1_score(all_labels, all_preds, average='macro')
        
        return eval_loss, accuracy, f1