# 计算机视觉中的知识蒸馏

知识蒸馏是一种将知识从更大、更复杂的模型（教师模型）迁移到更小、更简单的模型（学生模型）的技术。为了将一个模型的知识蒸馏到另一个模型中，我们采用一个在特定任务（本例中为图像分类）上训练好的预训练教师模型，并随机初始化一个学生模型，使其在图像分类任务上进行训练。接下来，我们训练学生模型以最小化其输出与教师模型输出之间的差异，从而使其模仿教师模型的行为。该技术最早由 Hinton 等人在 [Distilling the Knowledge in a Neural Network](https://huggingface.co/papers/1503.02531) 中提出。在本指南中，我们将进行任务特定的知识蒸馏。我们将使用 [beans 数据集](https://huggingface.co/datasets/AI-Lab-Makerere/beans) 来完成此任务。

本指南演示如何使用 🤗 Transformers 的 [Trainer API](../main_classes/trainer#trainer)，将一个 [微调过的 ViT 模型](https://huggingface.co/merve/vit-mobilenet-beans-224)（教师模型）蒸馏到一个 [MobileNet](https://huggingface.co/google/mobilenet_v2_1.4_224)（学生模型）。

让我们安装蒸馏和评估流程所需的库。

```bash
pip install transformers datasets accelerate tensorboard evaluate trackio --upgrade
```

在此示例中，我们使用 `merve/beans-vit-224` 模型作为教师模型。它是一个图像分类模型，基于 `google/vit-base-patch16-224-in21k` 并在 beans 数据集上进行了微调。我们将对该模型进行蒸馏，目标是一个随机初始化的 MobileNetV2。

现在我们将加载数据集。

```python
from datasets import load_dataset

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

我们可以使用这两个模型中任意一个的图像处理器，因为在这种情况下，它们以相同的分辨率返回相同的输出。我们将使用 `dataset` 的 `map()` 方法，将预处理应用于数据集的每个划分。

```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)
```

本质上，我们希望学生模型（随机初始化的 MobileNet）模仿教师模型（经过微调的视觉 Transformer）。为此，我们首先获取教师模型和学生模型的 logits 输出。然后，我们将它们各自除以参数 `temperature`，该参数控制每个软目标的重要性。另一个参数 `lambda` 用于衡量蒸馏损失的重要性。在本示例中，我们将使用 `temperature=5` 和 `lambda=0.5`。我们将使用 Kullback-Leibler 散度损失来计算学生模型与教师模型之间的散度。给定两个数据 P 和 Q，KL 散度说明了使用 Q 来表示 P 需要多少额外信息。如果两者完全相同，它们的 KL 散度为零，因为不需要其他信息就能从 Q 解释 P。因此，在知识蒸馏的语境下，KL 散度非常有用。

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

接下来，我们将登录 Hugging Face Hub，以便通过 `Trainer` 将模型推送到 Hugging Face Hub。

```python
from huggingface_hub import notebook_login

notebook_login()
```

让我们设置 `TrainingArguments`、教师模型和学生模型。

```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)
```

我们可以使用 `compute_metrics` 函数在测试集上评估模型。该函数将在训练过程中用于计算模型的 `accuracy` 和 `f1`。

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

让我们使用之前定义的训练参数来初始化 `Trainer`。我们还将初始化数据整理器（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
)
```

现在我们可以训练模型了。

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

我们可以在测试集上评估模型。

```python
trainer.evaluate(processed_datasets["test"])
```

在测试集上，我们的模型达到了 72% 的准确率。为了对蒸馏的效率进行合理性检查，我们还使用相同的超参数从头训练了 MobileNet 以在 beans 数据集上进行训练，并在测试集上观察到 63% 的准确率。我们邀请读者尝试不同的预训练教师模型、学生架构和蒸馏参数，并报告他们的发现。蒸馏模型的训练日志和检查点可以在 [此仓库](https://huggingface.co/merve/vit-mobilenet-beans-224) 中找到，从头训练的 MobileNetV2 可以在 [此仓库](https://huggingface.co/merve/resnet-mobilenet-beans-5) 中找到。