Destilação de conhecimento para visão computacional

A destilação de conhecimento é uma técnica usada para transferir conhecimento de um modelo maior e mais complexo (professor) para um modelo menor e mais simples (aluno). Para destilar conhecimento de um modelo para outro, pegamos um modelo professor pré-treinado em uma determinada tarefa (classificação de imagens, neste caso) e inicializamos aleatoriamente um modelo aluno para ser treinado em classificação de imagens. Em seguida, treinamos o modelo aluno para minimizar a diferença entre suas saídas e as saídas do professor, fazendo com que ele imite o comportamento. A técnica foi introduzida pela primeira vez em Distilling the Knowledge in a Neural Network, de Hinton et al. Neste guia, realizaremos destilação de conhecimento específica para a tarefa. Usaremos o conjunto de dados beans para isso.

Este guia demonstra como você pode destilar um modelo ViT ajustado (modelo professor) para um MobileNet (modelo aluno) usando a API Trainer do 🤗 Transformers.

Vamos instalar as bibliotecas necessárias para a destilação e a avaliação do processo.

pip install transformers datasets accelerate tensorboard evaluate trackio --upgrade

Neste exemplo, estamos usando o modelo merve/beans-vit-224 como modelo professor. Trata-se de um modelo de classificação de imagens, baseado no google/vit-base-patch16-224-in21k ajustado no conjunto de dados beans. Vamos destilar esse modelo para um MobileNetV2 inicializado aleatoriamente.

Agora, carregaremos o conjunto de dados.

from datasets import load_dataset

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

Podemos usar um processador de imagem de qualquer um dos modelos, pois, neste caso, eles retornam a mesma saída com a mesma resolução. Usaremos o método map() de dataset para aplicar o pré-processamento a cada divisão do conjunto de dados.

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)

Em essência, queremos que o modelo aluno (um MobileNet inicializado aleatoriamente) imite o modelo professor (um vision transformer ajustado). Para isso, primeiro obtemos os logits de saída do professor e do aluno. Em seguida, dividimos cada um deles pelo parâmetro temperature, que controla a importância de cada alvo suave. Um parâmetro chamado lambda pondera a importância da perda de destilação. Neste exemplo, usaremos temperature=5 e lambda=0.5. Usaremos a perda de Divergência de Kullback-Leibler para calcular a divergência entre o aluno e o professor. Dadas duas distribuições P e Q, a Divergência de KL explica quanta informação extra precisamos para representar P usando Q. Se as duas forem idênticas, sua divergência de KL é zero, pois não há outra informação necessária para explicar P a partir de Q. Assim, no contexto de destilação de conhecimento, a divergência de KL é ú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

Agora, faremos login no Hugging Face Hub para que possamos enviar nosso modelo ao Hugging Face Hub por meio do Trainer.

from huggingface_hub import notebook_login

notebook_login()

Vamos definir os TrainingArguments, o modelo professor e o modelo aluno.

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 a função compute_metrics para avaliar nosso modelo no conjunto de teste. Essa função será usada durante o processo de treinamento para calcular a accuracy e o f1 do nosso 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"]}

Vamos inicializar o Trainer com os argumentos de treinamento que definimos. Também vamos inicializar nosso 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
)

Agora podemos treinar nosso modelo.

trainer.train()

Podemos avaliar o modelo no conjunto de teste.

trainer.evaluate(processed_datasets["test"])

No conjunto de teste, nosso modelo atinge 72 por cento de precisão. Para verificar a eficiência do destilamento, também treinamos o MobileNet no conjunto de dados beans do zero, com os mesmos hiperparâmetros, e observamos 63 por cento de precisão no conjunto de teste. Convidamos os leitores a experimentar diferentes modelos de professor pré-treinados, arquiteturas de aluno e parâmetros de destilamento, e a relatar seus resultados. Os logs de treinamento e os checkpoints do modelo destilado podem ser encontrados neste repositório, e o MobileNetV2 treinado do zero pode ser encontrado neste repositório.

Update on GitHub