Distillation de connaissances pour la vision par ordinateur

La distillation de connaissances est une technique permettant de transférer les connaissances d’un modèle plus grand et plus complexe (enseignant) vers un modèle plus petit et plus simple (étudiant). Pour distiller les connaissances d’un modèle vers un autre, nous prenons un modèle enseignant pré-entraîné sur une tâche donnée (la classification d’images dans ce cas) et nous initialisons aléatoirement un modèle étudiant destiné à être entraîné sur la classification d’images. Ensuite, nous entraînons le modèle étudiant afin de minimiser l’écart entre ses sorties et celles de l’enseignant, ce qui lui permet d’imiter son comportement. Cette technique a été introduite pour la première fois dans Distilling the Knowledge in a Neural Network par Hinton et al. Dans ce guide, nous réaliserons une distillation de connaissances spécifique à une tâche. Nous utiliserons le beans dataset à cet effet.

Ce guide montre comment distiller un modèle ViT ajusté finement (modèle enseignant) vers un MobileNet (modèle élève) à l’aide de l’API Trainer de 🤗 Transformers.

Installons les bibliothèques nécessaires pour la distillation et l'évaluation du processus.

pip install transformers datasets accelerate tensorboard evaluate trackio --upgrade

Dans cet exemple, nous utilisons le modèle merve/beans-vit-224 en tant que modèle enseignant. Il s’agit d’un modèle de classification d’images, basé sur google/vit-base-patch16-224-in21k et affiné sur le jeu de données beans. Nous allons distiller ce modèle vers un MobileNetV2 initialisé aléatoirement.

Nous allons maintenant charger le jeu de données.

from datasets import load_dataset

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

Nous pouvons utiliser un processeur d’images de l’un ou l’autre des modèles, car dans ce cas ils produisent la même sortie avec la même résolution. Nous utiliserons la méthode map() de dataset pour appliquer le prétraitement à chaque partie du jeu de données.

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)

Essentiellement, nous voulons que le modèle élève (un MobileNet initialisé aléatoirement) imite le modèle enseignant (un vision transformer affiné). Pour y parvenir, nous obtenons d’abord les logits en sortie du modèle enseignant et du modèle élève. Ensuite, nous divisons chacun d’eux par le paramètre temperature, qui contrôle l’importance de chaque cible douce. Un paramètre appelé lambda pondère l’importance de la perte de distillation. Dans cet exemple, nous utiliserons temperature=5 et lambda=0.5. Nous utiliserons la perte de divergence de Kullback-Leibler pour calculer la divergence entre l’élève et l’enseignant. Étant données deux données P et Q, la divergence de KL explique la quantité d’informations supplémentaires nécessaires pour représenter P à l’aide de Q. Si les deux sont identiques, leur divergence de KL est nulle, car aucune autre information n’est nécessaire pour expliquer P à partir de Q. Ainsi, dans le contexte de la distillation de connaissances, la divergence de KL est utile.

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

Nous allons maintenant nous connecter au Hugging Face Hub afin de pouvoir pousser notre modèle sur le Hugging Face Hub via le Trainer.

from huggingface_hub import notebook_login

notebook_login()

Définissons les TrainingArguments, le modèle enseignant et le modèle élève.

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)

Nous pouvons utiliser la fonction compute_metrics pour évaluer notre modèle sur l'ensemble de test. Cette fonction sera utilisée pendant le processus d'entraînement pour calculer la accuracy et le f1 de notre modèle.

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

Initialisons le Trainer avec les arguments d’entraînement que nous avons définis. Nous initialisons également notre collateur de données.

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
)

Nous pouvons maintenant entraîner notre modèle.

trainer.train()

Nous pouvons évaluer le modèle sur le jeu de test.

trainer.evaluate(processed_datasets["test"])

Sur l’ensemble de test, notre modèle atteint un taux de précision de 72 %. Pour vérifier l’efficacité de la distillation, nous avons également entraîné MobileNet sur le jeu de données beans à partir de zéro avec les mêmes hyperparamètres et observé un taux de précision de 63 % sur l’ensemble de test. Nous invitons les lecteurs à essayer différents modèles d’enseignant pré-entraînés, architectures d’élève et paramètres de distillation, et à partager leurs résultats. Les journaux d’entraînement et les points de contrôle du modèle distillé sont disponibles dans ce dépôt, et MobileNetV2 entraîné à partir de zéro peut être trouvé dans ce dépôt.