증류

# 지식 증류 (Knowledge Distillation) 기본 구현
import torch
import torch.nn as nn
import torch.nn.functional as F

class DistillationLoss(nn.Module):
    """
    Knowledge Distillation Loss
    - alpha: soft/hard loss 비율
    - temperature: 확률 분포 smoothing
    """
    def __init__(self, alpha=0.5, temperature=4.0):
        super().__init__()
        self.alpha = alpha
        self.temperature = temperature
        self.ce_loss = nn.CrossEntropyLoss()
        self.kl_loss = nn.KLDivLoss(reduction='batchmean')

    def forward(self, student_logits, teacher_logits, labels):
        # Hard Loss: Student vs 실제 정답
        hard_loss = self.ce_loss(student_logits, labels)

        # Soft Loss: Student vs Teacher (temperature scaling)
        T = self.temperature
        soft_student = F.log_softmax(student_logits / T, dim=1)
        soft_teacher = F.softmax(teacher_logits / T, dim=1)
        soft_loss = self.kl_loss(soft_student, soft_teacher) * (T * T)

        # 총 손실: alpha로 가중 평균
        total_loss = self.alpha * soft_loss + (1 - self.alpha) * hard_loss
        return total_loss

# 사용 예시
def train_with_distillation(student, teacher, dataloader, optimizer):
    teacher.eval()  # Teacher는 학습하지 않음
    student.train()
    criterion = DistillationLoss(alpha=0.7, temperature=4.0)

    for inputs, labels in dataloader:
        optimizer.zero_grad()

        # Teacher 추론 (gradient 계산 안 함)
        with torch.no_grad():
            teacher_logits = teacher(inputs)

        # Student 학습
        student_logits = student(inputs)
        loss = criterion(student_logits, teacher_logits, labels)

        loss.backward()
        optimizer.step()

    print(f"Distillation 학습 완료!")

# Hugging Face 기반 LLM 증류 (DistilBERT 스타일)
from transformers import (
    AutoModelForSequenceClassification,
    AutoTokenizer,
    Trainer,
    TrainingArguments
)
from datasets import load_dataset

# Teacher: 대형 모델 (예: BERT-large)
teacher_name = "klue/roberta-large"
teacher = AutoModelForSequenceClassification.from_pretrained(
    teacher_name, num_labels=2
)

# Student: 소형 모델 (예: DistilBERT)
student_name = "monologg/distilkobert"
student = AutoModelForSequenceClassification.from_pretrained(
    student_name, num_labels=2
)
tokenizer = AutoTokenizer.from_pretrained(student_name)

# 커스텀 Trainer for Distillation
class DistillationTrainer(Trainer):
    def __init__(self, teacher_model, alpha=0.5, temperature=4.0, **kwargs):
        super().__init__(**kwargs)
        self.teacher = teacher_model
        self.teacher.eval()
        self.alpha = alpha
        self.temperature = temperature

    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.pop("labels")

        # Student forward
        student_outputs = model(**inputs)
        student_logits = student_outputs.logits

        # Teacher forward
        with torch.no_grad():
            teacher_outputs = self.teacher(**inputs)
            teacher_logits = teacher_outputs.logits

        # Distillation loss 계산
        loss = distillation_loss(
            student_logits, teacher_logits, labels,
            self.alpha, self.temperature
        )

        return (loss, student_outputs) if return_outputs else loss

print("Student 모델 파라미터:", sum(p.numel() for p in student.parameters()))
print("Teacher 모델 파라미터:", sum(p.numel() for p in teacher.parameters()))
# 출력: Student가 Teacher 대비 약 40-60% 크기