Bilgisayar görüşü için bilgi damıtma

Bilgi damıtma, daha büyük ve karmaşık bir modelden (öğretmen) daha küçük ve basit bir modele (öğrenci) bilgi aktarmak için kullanılan bir tekniktir. Bir modelden diğerine bilgi damıtmak için, belirli bir görevde (bu durumda görüntü sınıflandırma) eğitilmiş önceden eğitilmiş bir öğretmen modelini alır ve görüntü sınıflandırmada eğitilecek bir öğrenci modelini rastgele başlatır. Ardından, öğrenci modelinin çıktılarını öğretmen modelinin çıktılarını taklit edecek şekilde, bu çıktılar arasındaki farkı en aza indirecek biçimde eğitiriz. Bu yöntem ilk kez Hinton ve arkadaşlarının Distilling the Knowledge in a Neural Network makalesinde tanıtılmıştır. Bu kılavuzda, görev bazlı bilgi damıtma yapacağız. Bunun için beans veri setini kullanacağız.

Bu kılavuz, 🤗 Transformers'ın Trainer API kullanarak ince ayarlanmış bir ViT modelini (öğretmen modeli) bir MobileNet'e (öğrenci modeli) nasıl damıtacağınızı gösterir.

Damıtma ve süreci değerlendirme için gerekli kütüphaneleri yükleyelim.

pip install transformers datasets accelerate tensorboard evaluate trackio --upgrade

Bu örnekte, öğretmen modeli olarak merve/beans-vit-224 modelini kullanıyoruz. Bu, beans veri setinde ince ayar (fine-tune) edilmiş, google/vit-base-patch16-224-in21k tabanlı bir görüntü sınıflandırma modelidir. Bu modeli, rastgele başlatılmış bir MobileNetV2'ye damıtacağız.

Şimdi veri setini yükleyeceğiz.

from datasets import load_dataset

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

Bu durumda her iki model de aynı çözünürlükte aynı çıktıyı döndürdüğü için, görüntü işleyicisini modellerden herhangi birinden kullanabiliriz. Ön işleme adımını veri setinin her bölmesine uygulamak için dataset nesnesinin map() yöntemini kullanacağız.

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)

Temelde, öğrenci modelinin (rastgele başlatılmış bir MobileNet) öğretmen modelini (ince ayarlanmış vision transformer) taklit etmesini istiyoruz. Bunu yapmak için önce öğretmen ve öğrenciden gelen logits çıktılarını alırız. Ardından, her birini, her yumuşak hedefin önemini kontrol eden temperature parametresine böleriz. lambda adlı bir parametre, damıtma kaybının önemini tartar. Bu örnekte temperature=5 ve lambda=0.5 değerlerini kullanacağız. Öğrenci ile öğretmen arasındaki sapmayı hesaplamak için Kullback-Leibler Divergence kaybını kullanacağız. P ve Q olmak üzere iki veri verildiğinde, KL Divergence, P'yi Q kullanarak temsil etmek için ne kadar ek bilgiye ihtiyaç duyduğumuzu açıklar. İki dağılım aynıysa, Q'den P'yi açıklamak için başka bir bilgiye gerek olmadığından KL sapmaları sıfırdır. Dolayısıyla, bilgi damıtma bağlamında KL sapması faydalıdır.

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

Şimdi Trainer aracılığıyla modelimizi Hugging Face Hub'a yükleyebilmemiz için Hugging Face Hub'a giriş yapacağız.

from huggingface_hub import notebook_login

notebook_login()

TrainingArguments, öğretmen modelini ve öğrenci modelini ayarlayalım.

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)

Modelimizi test kümesinde değerlendirmek için compute_metrics fonksiyonunu kullanabiliriz. Bu fonksiyon, eğitim sürecinde modelimizin accuracy ve f1 değerlerini hesaplamak için kullanılacaktır.

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"]}

Eğitim argümanlarını tanımladığımız Trainer nesnesini başlatalım. Ayrıca veri toplayıcımızı da başlatacağız.

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
)

Artık modelimizi eğitebiliriz.

trainer.train()

Modeli test kümesinde değerlendirebiliriz.

trainer.evaluate(processed_datasets["test"])

Test kümesinde modelimiz %72 doğruluğa ulaşır. Damıtmanın verimliliği üzerinde bir sağlamlık kontrolü yapmak için, MobileNet'i aynı hiperparametrelerle beans veri setinden sıfırdan eğittik ve test kümesinde %63 doğruluk gözlemledik. Okuyucuları farklı önceden eğitilmiş öğretmen modelleri, öğrenci mimarileri ve damıtma parametreleri denemeye ve bulgularını raporlamaya davet ediyoruz. Damıtılmış modelin eğitim günlükleri ve denetim noktaları bu depoda, sıfırdan eğitilen MobileNetV2 ise bu depoda bulunabilir.

Update on GitHub