Transformers documentation
Метод дистилляции знаний для компьютерного зрения
Метод дистилляции знаний для компьютерного зрения
Метод дистилляции знаний (knowledge distillation) используется для передачи знаний от более крупной и сложной модели (учителя) к более малой и простой модели (ученика). Чтобы дистиллировать знания из одной модели в другую, мы берём предобученную модель-учителя, обученную на определённой задаче (в данном случае — классификация изображений), и случайно инициализируем модель-ученика, которая будет обучаться на классификации изображений. Далее мы обучаем модель-ученика минимизировать разницу между её выходными данными и выходными данными модели-учителя, тем самым заставляя её имитировать поведение учителя. Этот метод был впервые предложен в работе Distilling the Knowledge in a Neural Network by Hinton et al. В этом руководстве мы выполним дистилляцию знаний, специфичную для конкретной задачи. Для этого мы будем использовать beans dataset.
В этом руководстве показано, как с помощью Trainer API из 🤗 Transformers дистиллировать дообученную модель ViT (учительская модель) в MobileNet (студенческая модель).
Установим библиотеки, необходимые для дистилляции и оценки процесса.
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("AI-Lab-Makerere/beans")Мы можем использовать процессор изображений из любой из моделей, поскольку в данном случае они возвращают одинаковый результат с одинаковым разрешением. Мы будем использовать метод map() объекта dataset, чтобы применить предобработку к каждому разделу набора данных.
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. Для вычисления расхождения между ученической и учительской моделями мы будем использовать потерю на основе дивергенции Кульбака-Лейблера. Дана две распределения P и Q, дивергенция Кульбака-Лейблера показывает, сколько дополнительной информации необходимо для представления P с помощью Q. Если два распределения идентичны, их дивергенция Кульбака-Лейблера равна нулю, так как для объяснения P через Q не требуется никакой дополнительной информации. Таким образом, в контексте дистилляции знаний дивергенция Кульбака-Лейблера оказывается полезной.
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, чтобы загрузить нашу модель на Hugging Face Hub с помощью Trainer.
from huggingface_hub import notebook_login
notebook_login()Определим TrainingArguments, модель-учителя и модель-ученика.
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 с аргументами обучения, которые мы определили. Также инициализируем наш data collator.
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 процента. Чтобы убедиться в эффективности дистилляции, мы также обучили MobileNet на наборе данных beans с нуля с теми же гиперпараметрами и получили точность 63 процента на тестовой выборке. Мы приглашаем читателей попробовать различные предобученные модели-учителя, архитектуры ученика и параметры дистилляции и поделиться своими результатами. Журналы обучения и контрольные точки для дистиллированной модели можно найти в этом репозитории, а MobileNetV2, обученный с нуля, — в этом репозитории.
Update on GitHub