컴퓨터 비전을 위한 지식 증류

Knowledge distillation은 더 크고 복잡한 모델(teacher)에서 더 작고 단순한 모델(student)로 지식을 전달하는 기법입니다. 한 모델에서 다른 모델로 지식을 증류하려면, 특정 작업(이 경우 이미지 분류)으로 학습된 사전 학습된 teacher 모델을 가져오고, 이미지 분류로 학습될 student 모델을 무작위로 초기화합니다. 그런 다음 student 모델의 출력과 teacher의 출력 사이의 차이를 최소화하도록 student 모델을 학습시켜 teacher의 동작을 모방하게 만듭니다. 이 기법은 Hinton et al의 Distilling the Knowledge in a Neural Network에서 처음 소개되었습니다. 이 가이드에서는 작업별(task-specific) knowledge distillation을 수행합니다. 이를 위해 beans dataset을 사용합니다.

이 가이드에서는 🤗 Transformers의 Trainer API를 사용하여 fine-tuned ViT 모델(teacher model)을 MobileNet(student model)로 증류하는 방법을 보여줍니다.

디스틸레이션 및 프로세스 평가에 필요한 라이브러리를 설치해 보겠습니다.

pip install transformers datasets accelerate tensorboard evaluate trackio --upgrade

이 예제에서는 merve/beans-vit-224 모델을 teacher model로 사용합니다. 이는 beans dataset에서 fine-tuning된 google/vit-base-patch16-224-in21k 기반의 이미지 분류 모델입니다. 이 모델을 무작위로 초기화된 MobileNetV2로 distill합니다.

이제 데이터셋을 로드하겠습니다.

from datasets import load_dataset

dataset = load_dataset("AI-Lab-Makerere/beans")

두 모델 중 어느 쪽의 이미지 프로세서든 사용할 수 있습니다. 이 경우 두 프로세서 모두 동일한 해상도로 동일한 출력을 반환하기 때문입니다. dataset의 map() 메서드를 사용하여 데이터셋의 모든 split에 전처리를 적용합니다.

from transformers import AutoImageProcessor
teacher_processor = AutoImageProcessor.from_pretrained("merve/beans-vit-224")

def process(examples):
    processed_inputs = teacher_processor(examples["image"])
    return processed_inputs

processed_datasets = dataset.map(process, batched=True)

구체적으로, 학생 모델(무작위로 초기화된 MobileNet)이 교사 모델(파인튜닝된 vision transformer)을 모방하도록 하는 것이 목표입니다. 이를 위해 먼저 교사와 학생의 logits 출력을 가져옵니다. 그런 다음, 각 소프트 타겟의 중요도를 조절하는 매개변수 temperature로 각각을 나눕니다. lambda라는 매개변수는 증류 손실의 중요도를 가중합니다. 이 예제에서는 temperature=5와 lambda=0.5를 사용합니다. 학생과 교사 사이의 발산을 계산하기 위해 Kullback-Leibler Divergence 손실을 사용합니다. 두 데이터 P와 Q가 주어졌을 때, KL Divergence는 Q를 사용하여 P를 표현하는 데 필요한 추가 정보의 양을 설명합니다. 두 데이터가 동일하면 Q로부터 P를 설명하는 데 다른 정보가 필요하지 않으므로 KL 발산은 0이 됩니다. 따라서 지식 증류 맥락에서 KL 발산은 유용합니다.

from transformers import TrainingArguments, Trainer
from accelerate import Accelerator
import torch
import torch.nn as nn
import torch.nn.functional as F

class ImageDistilTrainer(Trainer):
    def __init__(self, teacher_model=None, student_model=None, temperature=None, lambda_param=None,  *args, **kwargs):
        super().__init__(model=student_model, *args, **kwargs)
        self.teacher = teacher_model
        self.student = student_model
        self.loss_function = nn.KLDivLoss(reduction="batchmean")
        device = Accelerator().device
        self.teacher.to(device)
        self.teacher.eval()
        self.temperature = temperature
        self.lambda_param = lambda_param

    def compute_loss(self, student, inputs, return_outputs=False):
        student_output = self.student(**inputs)

        with torch.no_grad():
          teacher_output = self.teacher(**inputs)

        # Compute soft targets for teacher and student
        soft_teacher = F.softmax(teacher_output.logits / self.temperature, dim=-1)
        soft_student = F.log_softmax(student_output.logits / self.temperature, dim=-1)

        # Compute the loss
        distillation_loss = self.loss_function(soft_student, soft_teacher) * (self.temperature ** 2)

        # Compute the true label loss
        student_target_loss = student_output.loss

        # Calculate final loss
        loss = (1. - self.lambda_param) * student_target_loss + self.lambda_param * distillation_loss
        return (loss, student_output) if return_outputs else loss

이제 Trainer를 통해 Hugging Face Hub에 모델을 push할 수 있도록 Hugging Face Hub에 로그인합니다.

from huggingface_hub import notebook_login

notebook_login()

TrainingArguments, teacher 모델, student 모델을 설정해 보겠습니다.

from transformers import AutoModelForImageClassification, MobileNetV2Config, MobileNetV2ForImageClassification

training_args = TrainingArguments(
    output_dir="my-awesome-model",
    num_train_epochs=30,
    fp16=True,
    logging_strategy="epoch",
    eval_strategy="epoch",
    save_strategy="epoch",
    load_best_model_at_end=True,
    metric_for_best_model="accuracy",
    push_to_hub=True,
    hub_strategy="every_save",
    hub_model_id=repo_name,
    report_to="trackio",
    run_name="distillation",
    )

num_labels = len(processed_datasets["train"].features["labels"].names)

# initialize models
teacher_model = AutoModelForImageClassification.from_pretrained(
    "merve/beans-vit-224",
    num_labels=num_labels,
    ignore_mismatched_sizes=True
)

# training MobileNetV2 from scratch
student_config = MobileNetV2Config()
student_config.num_labels = num_labels
student_model = MobileNetV2ForImageClassification(student_config)

compute_metrics 함수를 사용하여 테스트 세트에서 모델을 평가할 수 있습니다. 이 함수는 학습 과정에서 모델의 accuracy 및 f1을 계산하는 데 사용됩니다.

import evaluate
import numpy as np

accuracy = evaluate.load("accuracy")

def compute_metrics(eval_pred):
    predictions, labels = eval_pred
    acc = accuracy.compute(references=labels, predictions=np.argmax(predictions, axis=1))
    return {"accuracy": acc["accuracy"]}

정의한 학습 인수를 사용하여 Trainer를 초기화해 보겠습니다. 데이터 콜레이터도 함께 초기화합니다.

from transformers import DefaultDataCollator

data_collator = DefaultDataCollator()
trainer = ImageDistilTrainer(
    student_model=student_model,
    teacher_model=teacher_model,
    training_args=training_args,
    train_dataset=processed_datasets["train"],
    eval_dataset=processed_datasets["validation"],
    data_collator=data_collator,
    processing_class=teacher_processor,
    compute_metrics=compute_metrics,
    temperature=5,
    lambda_param=0.5
)

이제 모델을 학습할 수 있습니다.

trainer.train()

모델을 테스트 세트에서 평가할 수 있습니다.

trainer.evaluate(processed_datasets["test"])

테스트 세트에서 우리 모델은 72%의 정확도를 달성했습니다. 증류의 효율성에 대한 sanity check를 위해, 동일한 하이퍼파라미터로 MobileNet을 beans 데이터셋에서 처음부터 학습시켜 보았으며, 테스트 세트에서 63%의 정확도를 확인했습니다. 독자분들께는 다양한 사전 학습된 teacher 모델, student 아키텍처, 증류 파라미터를 시도해 보고 그 결과를 공유해 주시기를 권합니다. 증류된 모델의 학습 로그와 체크포인트는 이 저장소에서, 처음부터 학습된 MobileNetV2는 이 저장소에서 확인하실 수 있습니다.

Update on GitHub