컴퓨터 비전을 위한 지식 증류
컴퓨터 비전을 위한 지식 증류 (Knowledge distillation for computer vision)
지식 증류는 더 크고 복잡한 모델(teacher)에서 더 작고 단순한 모델(student)로 지식을 전달하는 기법입니다. 한 모델에서 다른 모델로 지식을 증류하기 위해, 특정 작업(이 경우 이미지 분류)으로 훈련된 사전훈련된 teacher 모델을 가져와 이미지 분류용으로 훈련할 student 모델을 무작위로 초기화합니다.
출처: 문서
본문
다음으로, student 모델의 출력과 teacher의 출력 사이의 차이를 최소화하도록 훈련해 그 동작을 모방하게 만듭니다. 이 기법은 Distilling the Knowledge in a Neural Network by Hinton et al에서 처음 소개되었습니다. 이 가이드에서는 작업별(task-specific) 지식 증류를 수행합니다. 이를 위해 beans dataset을 사용하겠습니다.
이 가이드는 🤗 Transformers의 Trainer API를 사용해 파인튜닝된 ViT 모델(teacher 모델)을 MobileNet(student 모델)으로 증류하는 방법을 보여줍니다.
증류와 평가 과정에 필요한 라이브러리를 설치하겠습니다.
pip install transformers datasets accelerate tensorboard evaluate trackio --upgrade
이 예시에서는 teacher 모델로 merve/beans-vit-224를 사용합니다. 이것은 google/vit-base-patch16-224-in21k를 기반으로 beans 데이터셋에서 파인튜닝된 이미지 분류 모델입니다. 이 모델을 무작위로 초기화된 MobileNetV2로 증류하겠습니다.
이제 데이터셋을 로드하겠습니다.
from datasets import load_dataset
dataset = load_dataset("beans")
두 모델 모두 같은 해상도로 같은 출력을 반환하므로 두 모델 중 아무 이미지 프로세서나 사용할 수 있습니다. dataset의 map() 메서드를 사용해 데이터셋의 모든 스플릿에 전처리를 적용하겠습니다.
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)
기본적으로 우리는 student 모델(무작위로 초기화된 MobileNet)이 teacher 모델(파인튜닝된 비전 트랜스포머)을 모방하도록 만들고자 합니다. 이를 위해 먼저 teacher와 student에서 logits 출력을 얻습니다. 그런 다음 각각을 각 소프트 타겟의 중요성을 제어하는 매개변수 temperature로 나눕니다. lambda라는 매개변수는 증류 loss의 중요성에 가중치를 둡니다. 이 예시에서는 temperature=5와 lambda=0.5를 사용하겠습니다. Kullback-Leibler Divergence loss를 사용해 student와 teacher 사이의 발산을 계산하겠습니다. 두 데이터 P와 Q가 주어졌을 때, KL Divergence는 Q를 사용해 P를 표현하는 데 필요한 추가 정보가 얼마나 되는지 설명합니다. 둘이 동일하면 KL 발산은 0입니다. Q에서 P를 설명하는 데 다른 정보가 필요 없기 때문입니다. 따라서 지식 증류의 맥락에서 KL divergence는 유용합니다.
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
이제 Hugging Face Hub에 로그인해 Trainer를 통해 모델을 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퍼센트 정확도에 도달합니다. 증류의 효율성을 검증하기 위해 동일한 하이퍼파라미터로 beans 데이터셋에서 MobileNet을 처음부터 훈련했고, 테스트 세트에서 63퍼센트 정확도를 관찰했습니다. 독자 여러분께 다양한 사전훈련된 teacher 모델, student 아키텍처, 증류 매개변수를 시도해 보시고 결과를 보고해 주시기를 권합니다. 증류된 모델의 훈련 로그와 체크포인트는 this repository에서, 처음부터 훈련된 MobileNetV2는 이 저장소에서 찾을 수 있습니다.