コンピュータビジョンにおける知識蒸留

知識蒸留（Knowledge distillation）は、より大きく複雑なモデル（教師）から、より小さくシンプルなモデル（生徒）へ知識を移転するための技術です。あるモデルから別のモデルへ知識を蒸留するには、特定のタスク（ここでは画像分類）で訓練済みの教師モデルを取得し、画像分類で訓練する生徒モデルをランダムに初期化します。次に、生徒モデルの出力と教師モデルの出力との差を最小化するように生徒モデルを訓練し、教師の挙動を模倣させます。この手法は、Hinton らによる論文「Distilling the Knowledge in a Neural Network」で初めて提案されました。本ガイドでは、タスク固有の知識蒸留を行います。これには beans データセットを使用します。

このガイドでは、🤗 Transformers の Trainer API を使用して、ファインチューニング済みの ViT モデル（教師モデル）を MobileNet（生徒モデル）に蒸留する方法を示します。

蒸留とプロセスの評価に必要なライブラリをインストールしましょう。

pip install transformers datasets accelerate tensorboard evaluate trackio --upgrade

この例では、merve/beans-vit-224 モデルを教師モデルとして使用しています。これは画像分類モデルであり、google/vit-base-patch16-224-in21k をベースに、beans データセットでファインチューニングされています。このモデルを、ランダムに初期化された MobileNetV2 に蒸留します。

次に、データセットを読み込みます。

from datasets import load_dataset

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

この場合、両方のモデルは同じ解像度で同じ出力を返すため、どちらかのモデルの画像プロセッサを使用できます。datasetのmap()メソッドを使用して、データセットのすべての分割に前処理を適用します。

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)

本質的には、教師モデル（ファインチューニング済みの vision transformer）の挙動を、生徒モデル（ランダム初期化された MobileNet）に模倣させたいと考えています。そのために、まず教師モデルと生徒モデルそれぞれから logits の出力を取得します。次に、それらをパラメータ temperature で除算します。temperature は各ソフトターゲットの重要度を制御します。パラメータ lambda は、蒸留損失の重要度を重み付けします。この例では、temperature=5 と lambda=0.5 を使用します。生徒と教師の間の乖離を計算するために、Kullback-Leibler 発散損失を使用します。2つのデータ P と Q が与えられたとき、KL 発散は、Q を用いて P を表現するために必要な追加情報の量を説明します。2つが完全に一致する場合、Q から P を説明するために他の情報は不要であるため、KL 発散はゼロになります。したがって、知識蒸留の文脈において、KL 発散は有用です。

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

次に、Trainer を介して Hugging Face Hub にモデルをプッシュできるよう、Hugging Face Hub にログインします。

from huggingface_hub import notebook_login

notebook_login()

TrainingArguments、教師モデル、生徒モデルを設定しましょう。

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 を計算するために使用されます。

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を初期化しましょう。また、データコラテーターも初期化します。

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
)

これでモデルを学習できます。

trainer.train()

モデルをテストセットで評価できます。

trainer.evaluate(processed_datasets["test"])

テストセットにおいて、当社のモデルは72パーセントの精度に達しました。蒸留の効率性を健全性チェックするために、同じハイパーパラメータでMobileNetをbeansデータセットからゼロから訓練し、テストセットで63パーセントの精度を観察しました。読者には、異なる事前学習済み教師モデル、生徒アーキテクチャ、蒸留パラメータを試して、その結果を報告していただくことをお勧めします。蒸留モデルの訓練ログとチェックポイントはこのリポジトリで、ゼロから訓練したMobileNetV2はこのリポジトリで確認できます。

Update on GitHub