Pengetahuan distilasi untuk visi komputer

Knowledge distillation adalah teknik yang digunakan untuk mentransfer pengetahuan dari model yang lebih besar dan kompleks (teacher) ke model yang lebih kecil dan sederhana (student). Untuk mendistilasi pengetahuan dari satu model ke model lain, kita mengambil model teacher yang telah dilatih sebelumnya pada tugas tertentu (klasifikasi gambar dalam kasus ini) dan menginisialisasi secara acak model student yang akan dilatih untuk klasifikasi gambar. Selanjutnya, kita melatih model student untuk meminimalkan perbedaan antara output-nya dengan output teacher, sehingga model tersebut meniru perilaku teacher. Teknik ini pertama kali diperkenalkan dalam Distilling the Knowledge in a Neural Network oleh Hinton et al. Dalam panduan ini, kita akan melakukan knowledge distilasi yang spesifik untuk tugas tertentu. Kita akan menggunakan beans dataset untuk ini.

Panduan ini menunjukkan cara Anda mendistilasi model ViT yang telah dihaluskan (model guru) menjadi MobileNet (model siswa) menggunakan API Trainer dari 🤗 Transformers.

Mari kita instal pustaka yang diperlukan untuk distilasi dan mengevaluasi prosesnya.

pip install transformers datasets accelerate tensorboard evaluate trackio --upgrade

Pada contoh ini, kami menggunakan model merve/beans-vit-224 sebagai model guru. Ini adalah model klasifikasi gambar, yang didasarkan pada google/vit-base-patch16-224-in21k yang telah dihaluskan (fine-tuned) pada dataset beans. Kami akan mendistilasi model ini ke MobileNetV2 yang diinisialisasi secara acak.

Kita akan memuat dataset sekarang.

from datasets import load_dataset

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

Kita dapat menggunakan image processor dari salah satu model, karena dalam kasus ini keduanya menghasilkan output yang sama dengan resolusi yang sama. Kita akan menggunakan metode map() dari dataset untuk menerapkan pra-pemrosesan ke setiap split dari 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)

Pada dasarnya, kami ingin model student (MobileNet yang diinisialisasi secara acak) meniru model teacher (vision transformer yang telah di-fine-tune). Untuk mencapainya, kami terlebih dahulu mengambil output logits dari model teacher dan model student. Kemudian, kami membagi masing-masing dari mereka dengan parameter temperature yang mengontrol pentingnya setiap soft target. Parameter bernama lambda menimbang pentingnya distillation loss. Dalam contoh ini, kami akan menggunakan temperature=5 dan lambda=0.5. Kami akan menggunakan Kullback-Leibler Divergence loss untuk menghitung divergensi antara model student dan model teacher. Diberikan dua data P dan Q, KL Divergence menjelaskan berapa banyak informasi tambahan yang kami butuhkan untuk merepresentasikan P menggunakan Q. Jika keduanya identik, divergensi KL mereka adalah nol, karena tidak ada informasi lain yang diperlukan untuk menjelaskan P dari Q. Dengan demikian, dalam konteks knowledge distillation, KL divergence sangat berguna.

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

Sekarang kita akan login ke Hugging Face Hub agar dapat mendorong model kita ke Hugging Face Hub melalui Trainer.

from huggingface_hub import notebook_login

notebook_login()

Mari kita atur TrainingArguments, model guru, dan model siswa.

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)

Kita dapat menggunakan fungsi compute_metrics untuk mengevaluasi model kita pada set pengujian. Fungsi ini akan digunakan selama proses pelatihan untuk menghitung accuracy & f1 dari model kita.

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

Mari kita inisialisasi Trainer dengan argumen pelatihan yang telah kita definisikan. Kita juga akan menginisialisasi 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
)

Sekarang kita dapat melatih model kita.

trainer.train()

Kita dapat mengevaluasi model pada set pengujian.

trainer.evaluate(processed_datasets["test"])

Pada set pengujian, model kami mencapai akurasi 72 persen. Untuk melakukan pemeriksaan kewajaran atas efisiensi distilasi, kami juga melatih MobileNet pada dataset beans dari awal dengan hiperparameter yang sama dan mengamati akurasi 63 persen pada set pengujian. Kami mengundang pembaca untuk mencoba berbagai model guru yang telah dilatih sebelumnya, arsitektur siswa, dan parameter distilasi, lalu melaporkan temuan mereka. Log pelatihan dan checkpoint untuk model hasil distilasi dapat ditemukan di repositori ini, dan MobileNetV2 yang dilatih dari awal dapat ditemukan di repositori ini.