# 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. Essa técnica foi introduzida pela primeira vez em [Distilling the Knowledge in a Neural Network, de Hinton et al](https://huggingface.co/papers/1503.02531). Neste guia, faremos destilação de conhecimento específica para a tarefa. Usaremos o [conjunto de dados beans](https://huggingface.co/datasets/AI-Lab-Makerere/beans) para isso.

Este guia demonstra como você pode destilar um [modelo ViT ajustado](https://huggingface.co/merve/vit-mobilenet-beans-224) (modelo professor) para um [MobileNet](https://huggingface.co/google/mobilenet_v2_1.4_224) (modelo aluno) usando a [API Trainer](../main_classes/trainer#trainer) do 🤗 Transformers.

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

```bash
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 este modelo para um MobileNetV2 inicializado aleatoriamente.

Agora, carregaremos o conjunto de dados.

```python
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.

```python
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 a saída de logits do professor e do aluno. Em seguida, dividimos cada uma delas 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.

```python
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`.

```python
from huggingface_hub import notebook_login

notebook_login()
```

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

```python
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.

```python
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.

```python
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.

```python
trainer.train()
```

Podemos avaliar o modelo no conjunto de teste.

```python
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](https://huggingface.co/merve/vit-mobilenet-beans-224), e o MobileNetV2 treinado do zero pode ser encontrado neste [repositório](https://huggingface.co/merve/resnet-mobilenet-beans-5).