docs/guide/43_Knowledge_Distillation_for_Computer_Vision.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5e7e876a",
   "metadata": {},
   "source": [
    "# 知识蒸馏在计算机视觉中的应用\n",
    "\n",
    "知识蒸馏是一种将大型复杂模型（教师模型）的知识转移到小型简单模型（学生模型）的技术。为了从一个模型中提取知识并传递给另一个模型，我们首先使用一个预训练的教师模型（例如图像分类任务），然后随机初始化一个学生模型来学习相同的任务。接下来，我们训练学生模型，使其输出与教师模型的输出尽可能接近，从而模拟教师模型的行为。这一技术最早由 Hinton 等人在论文 [《神经网络中的知识蒸馏》](https://arxiv.org/abs/1503.02531) 中提出。在本指南中，我们将进行特定任务的知识蒸馏。我们将使用 [beans 数据集](https://huggingface.co/datasets/beans)。\n",
    "\n",
    "本指南展示了如何使用 🤗 Transformers 的 [Trainer API](https://huggingface.co/docs/transformers/en/main_classes/trainer#trainer) 将一个经过微调的 [ViT 模型](https://huggingface.co/merve/vit-mobilenet-beans-224)（教师模型）蒸馏到一个 [MobileNet](https://huggingface.co/google/mobilenet_v2_1.4_224)（学生模型）。\n",
    "\n",
    "让我们安装蒸馏和评估过程中需要的库。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "df364fdc",
   "metadata": {
    "vscode": {
     "languageId": "shellscript"
    }
   },
   "outputs": [],
   "source": [
    "pip install transformers datasets accelerate tensorboard evaluate --upgrade"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ef5e2458",
   "metadata": {},
   "source": [
    "\n",
    "在这个例子中，我们将使用 `merve/beans-vit-224` 模型作为教师模型。这是一个基于 `google/vit-base-patch16-224-in21k` 在 beans 数据集上微调的图像分类模型。我们将把这个模型蒸馏到一个随机初始化的 MobileNetV2。\n",
    "\n",
    "现在我们加载数据集。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dcbb1023",
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "\n",
    "dataset = load_dataset(\"beans\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4564ea5",
   "metadata": {},
   "source": [
    "\n",
    "我们可以使用任一模型的图像处理器，因为在这种情况下它们返回相同分辨率的相同输出。我们将使用 `dataset` 的 `map()` 方法对数据集的每个分片进行预处理。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "74f3ec30",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoImageProcessor\n",
    "teacher_processor = AutoImageProcessor.from_pretrained(\"merve/beans-vit-224\")\n",
    "\n",
    "def process(examples):\n",
    "    processed_inputs = teacher_processor(examples[\"image\"])\n",
    "    return processed_inputs\n",
    "\n",
    "processed_datasets = dataset.map(process, batched=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0afce494",
   "metadata": {},
   "source": [
    "\n",
    "我们的目标是让学生模型（随机初始化的 MobileNet）模仿教师模型（微调后的视觉变换器）。为了实现这一点，我们首先获取教师模型和学生模型的 logits 输出。然后，我们将每种输出除以参数 `temperature`，该参数控制每个软目标的重要性。参数 `lambda` 用于权衡蒸馏损失的重要性。在本例中，我们将使用 `temperature=5` 和 `lambda=0.5`。我们将使用 Kullback-Leibler 散度损失来计算学生模型和教师模型之间的差异。给定两个数据 P 和 Q，KL 散度解释了用 Q 表示 P 所需的额外信息量。如果两者完全相同，它们的 KL 散度为零，因为不需要其他信息来解释 P。因此，在知识蒸馏的背景下，KL 散度是有用的。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6bc16ad",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import TrainingArguments, Trainer\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "class ImageDistilTrainer(Trainer):\n",
    "    def __init__(self, teacher_model=None, student_model=None, temperature=None, lambda_param=None, *args, **kwargs):\n",
    "        super().__init__(model=student_model, *args, **kwargs)\n",
    "        self.teacher = teacher_model\n",
    "        self.student = student_model\n",
    "        self.loss_function = nn.KLDivLoss(reduction=\"batchmean\")\n",
    "        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
    "        self.teacher.to(device)\n",
    "        self.teacher.eval()\n",
    "        self.temperature = temperature\n",
    "        self.lambda_param = lambda_param\n",
    "\n",
    "    def compute_loss(self, student, inputs, return_outputs=False):\n",
    "        student_output = self.student(**inputs)\n",
    "\n",
    "        with torch.no_grad():\n",
    "            teacher_output = self.teacher(**inputs)\n",
    "\n",
    "        # 计算教师和学生的软目标\n",
    "        soft_teacher = F.softmax(teacher_output.logits / self.temperature, dim=-1)\n",
    "        soft_student = F.log_softmax(student_output.logits / self.temperature, dim=-1)\n",
    "\n",
    "        # 计算损失\n",
    "        distillation_loss = self.loss_function(soft_student, soft_teacher) * (self.temperature ** 2)\n",
    "\n",
    "        # 计算真实标签的损失\n",
    "        student_target_loss = student_output.loss\n",
    "\n",
    "        # 计算最终损失\n",
    "        loss = (1. - self.lambda_param) * student_target_loss + self.lambda_param * distillation_loss\n",
    "        return (loss, student_output) if return_outputs else loss"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b65504da",
   "metadata": {},
   "source": [
    "\n",
    "现在我们登录 Hugging Face Hub，以便通过 `Trainer` 将模型推送到 Hugging Face Hub。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "485a7672",
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import notebook_login\n",
    "\n",
    "notebook_login()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8156c6f8",
   "metadata": {},
   "source": [
    "\n",
    "设置 `TrainingArguments`、教师模型和学生模型。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b8a79d36",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForImageClassification, MobileNetV2Config, MobileNetV2ForImageClassification\n",
    "\n",
    "training_args = TrainingArguments(\n",
    "    output_dir=\"my-awesome-model\",\n",
    "    num_train_epochs=30,\n",
    "    fp16=True,\n",
    "    logging_dir=f\"{repo_name}/logs\",\n",
    "    logging_strategy=\"epoch\",\n",
    "    eval_strategy=\"epoch\",\n",
    "    save_strategy=\"epoch\",\n",
    "    load_best_model_at_end=True,\n",
    "    metric_for_best_model=\"accuracy\",\n",
    "    report_to=\"tensorboard\",\n",
    "    push_to_hub=True,\n",
    "    hub_strategy=\"every_save\",\n",
    "    hub_model_id=repo_name,\n",
    ")\n",
    "\n",
    "num_labels = len(processed_datasets[\"train\"].features[\"labels\"].names)\n",
    "\n",
    "# 初始化模型\n",
    "teacher_model = AutoModelForImageClassification.from_pretrained(\n",
    "    \"merve/beans-vit-224\",\n",
    "    num_labels=num_labels,\n",
    "    ignore_mismatched_sizes=True\n",
    ")\n",
    "\n",
    "# 从头开始训练 MobileNetV2\n",
    "student_config = MobileNetV2Config()\n",
    "student_config.num_labels = num_labels\n",
    "student_model = MobileNetV2ForImageClassification(student_config)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e4a733ad",
   "metadata": {},
   "source": [
    "\n",
    "我们可以使用 `compute_metrics` 函数在测试集上评估模型。此函数将在训练过程中计算模型的 `accuracy` 和 `f1`。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bd718c3f",
   "metadata": {},
   "outputs": [],
   "source": [
    "import evaluate\n",
    "import numpy as np\n",
    "\n",
    "accuracy = evaluate.load(\"accuracy\")\n",
    "\n",
    "def compute_metrics(eval_pred):\n",
    "    predictions, labels = eval_pred\n",
    "    acc = accuracy.compute(references=labels, predictions=np.argmax(predictions, axis=1))\n",
    "    return {\"accuracy\": acc[\"accuracy\"]}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c38d2c2b",
   "metadata": {},
   "source": [
    "\n",
    "现在我们使用定义的训练参数初始化 `Trainer`。我们还将初始化数据收集器。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b50b2dd4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import DefaultDataCollator\n",
    "\n",
    "data_collator = DefaultDataCollator()\n",
    "trainer = ImageDistilTrainer(\n",
    "    student_model=student_model,\n",
    "    teacher_model=teacher_model,\n",
    "    training_args=training_args,\n",
    "    train_dataset=processed_datasets[\"train\"],\n",
    "    eval_dataset=processed_datasets[\"validation\"],\n",
    "    data_collator=data_collator,\n",
    "    processing_class=teacher_processor,\n",
    "    compute_metrics=compute_metrics,\n",
    "    temperature=5,\n",
    "    lambda_param=0.5\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "424aa589",
   "metadata": {},
   "source": [
    "\n",
    "现在我们可以训练模型。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2fd45d48",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fc132630",
   "metadata": {},
   "source": [
    "\n",
    "我们可以在测试集上评估模型。\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fe9486ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.evaluate(processed_datasets[\"test\"])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c36ff428",
   "metadata": {},
   "source": [
    "\n",
    "在测试集上，我们的模型达到了 72% 的准确率。为了验证蒸馏的有效性，我们还使用相同的超参数从头开始训练了一个 MobileNet，结果在测试集上的准确率为 63%。我们鼓励读者尝试不同的预训练教师模型、学生架构、蒸馏参数，并报告他们的发现。蒸馏模型的训练日志和检查点可以在 [这个仓库](https://huggingface.co/merve/vit-mobilenet-beans-224) 中找到，从头开始训练的 MobileNetV2 可以在 [这个仓库](https://huggingface.co/merve/resnet-mobilenet-beans-5) 中找到。\n"
   ]
  }
 ],
 "metadata": {
  "language_info": {
   "name": "python"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}