Destilación de conocimiento para visión por computadora

La destilación de conocimiento es una técnica que se utiliza para transferir conocimiento de un modelo más grande y complejo (el maestro) a un modelo más pequeño y simple (el estudiante). Para destilar conocimiento de un modelo a otro, tomamos un modelo maestro preentrenado en una tarea específica (clasificación de imágenes, en este caso) e inicializamos aleatoriamente un modelo estudiante que se entrenará para la clasificación de imágenes. A continuación, entrenamos el modelo estudiante para minimizar la diferencia entre sus salidas y las del maestro, de modo que imite su comportamiento. Esta técnica fue introducida por primera vez en Distilling the Knowledge in a Neural Network, de Hinton et al. En esta guía, realizaremos destilación de conocimiento específica de la tarea. Utilizaremos el conjunto de datos beans para ello.

Esta guía demuestra cómo puedes destilar un modelo ViT afinado (modelo maestro) a un MobileNet (modelo estudiante) utilizando la API de Trainer de 🤗 Transformers.

Instalemos las bibliotecas necesarias para la destilación y la evaluación del proceso.

pip install transformers datasets accelerate tensorboard evaluate trackio --upgrade

En este ejemplo, utilizamos el modelo merve/beans-vit-224 como modelo maestro. Es un modelo de clasificación de imágenes, basado en google/vit-base-patch16-224-in21k y ajustado en el conjunto de datos beans. Destilaremos este modelo en un MobileNetV2 inicializado aleatoriamente.

Ahora cargaremos el conjunto de datos.

from datasets import load_dataset

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

Podemos usar un procesador de imágenes de cualquiera de los modelos, ya que en este caso devuelven la misma salida con la misma resolución. Utilizaremos el método map() de dataset para aplicar el preprocesamiento a cada división del conjunto de datos.

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)

En esencia, queremos que el modelo estudiante (un MobileNet inicializado aleatoriamente) imite al modelo maestro (un vision transformer afinado). Para lograrlo, primero obtenemos los logits de salida del maestro y del estudiante. Luego, dividimos cada uno de ellos por el parámetro temperature, que controla la importancia de cada objetivo suave. Un parámetro llamado lambda pondera la importancia de la pérdida de destilación. En este ejemplo, utilizaremos temperature=5 y lambda=0.5. Utilizaremos la pérdida de Divergencia de Kullback-Leibler para calcular la divergencia entre el estudiante y el maestro. Dadas dos distribuciones P y Q, la Divergencia de KL explica cuánta información adicional necesitamos para representar P usando Q. Si ambas son idénticas, su divergencia de KL es cero, ya que no se necesita otra información para explicar P a partir de Q. Por lo tanto, en el contexto de la destilación de conocimiento, la divergencia de KL resulta útil.

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

Ahora iniciaremos sesión en Hugging Face Hub para poder subir nuestro modelo al Hugging Face Hub a través de Trainer.

from huggingface_hub import notebook_login

notebook_login()

Establezcamos los TrainingArguments, el modelo maestro y el modelo estudiante.

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)

Podemos usar la función compute_metrics para evaluar nuestro modelo en el conjunto de prueba. Esta función se utilizará durante el proceso de entrenamiento para calcular la accuracy y la f1 de nuestro modelo.

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

Iniciemos el Trainer con los argumentos de entrenamiento que definimos. También inicializaremos nuestro 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
)

Ahora podemos entrenar nuestro modelo.

trainer.train()

Podemos evaluar el modelo en el conjunto de prueba.

trainer.evaluate(processed_datasets["test"])

En el conjunto de prueba, nuestro modelo alcanza una precisión del 72 por ciento. Para verificar la eficiencia de la destilación, también entrenamos MobileNet desde cero en el conjunto de datos beans con los mismos hiperparámetros y observamos una precisión del 63 por ciento en el conjunto de prueba. Invitamos a los lectores a probar diferentes modelos de maestro preentrenados, arquitecturas de estudiante y parámetros de destilación, y a reportar sus hallazgos. Los registros de entrenamiento y los puntos de control del modelo destilado se pueden encontrar en este repositorio, y MobileNetV2 entrenado desde cero se puede encontrar en este repositorio.