Chưng cất kiến thức cho thị giác máy tính

Knowledge distillation là một kỹ thuật được sử dụng để chuyển giao kiến thức từ một mô hình lớn, phức tạp hơn (teacher) sang một mô hình nhỏ, đơn giản hơn (student). Để chưng cất kiến thức từ mô hình này sang mô hình khác, chúng ta lấy một mô hình teacher đã được huấn luyện trước trên một tác vụ nhất định (phân loại ảnh trong trường hợp này) và khởi tạo ngẫu nhiên một mô hình student để huấn luyện trên phân loại ảnh. Tiếp theo, chúng ta huấn luyện mô hình student để giảm thiểu sự khác biệt giữa đầu ra của nó và đầu ra của teacher, từ đó khiến nó bắt chước hành vi của teacher. Kỹ thuật này lần đầu được giới thiệu trong Distilling the Knowledge in a Neural Network của Hinton et al. Trong hướng dẫn này, chúng ta sẽ thực hiện knowledge distillation theo tác vụ. Chúng ta sẽ sử dụng beans dataset cho việc này.

Hướng dẫn này cho thấy cách bạn có thể chưng cất một mô hình ViT đã tinh chỉnh (mô hình giáo viên) sang một MobileNet (mô hình học sinh) bằng cách sử dụng Trainer API của 🤗 Transformers.

Hãy cài đặt các thư viện cần thiết cho quá trình chưng cất và đánh giá quy trình.

pip install transformers datasets accelerate tensorboard evaluate trackio --upgrade

Trong ví dụ này, chúng ta sử dụng mô hình merve/beans-vit-224 làm mô hình teacher. Đây là một mô hình phân loại hình ảnh, dựa trên google/vit-base-patch16-224-in21k được fine-tune trên bộ dữ liệu beans. Chúng ta sẽ distill mô hình này sang một MobileNetV2 được khởi tạo ngẫu nhiên.

Chúng ta sẽ tải bộ dữ liệu.

from datasets import load_dataset

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

Chúng ta có thể sử dụng bộ xử lý ảnh từ bất kỳ mô hình nào, vì trong trường hợp này chúng trả về cùng một đầu ra với cùng độ phân giải. Chúng ta sẽ sử dụng phương thức map() của dataset để áp dụng tiền xử lý cho mọi phần chia của tập dữ liệu.

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)

Về cơ bản, chúng ta muốn mô hình học sinh (MobileNet được khởi tạo ngẫu nhiên) bắt chước mô hình giáo viên (vision transformer đã được tinh chỉnh). Để đạt được điều này, trước tiên chúng ta lấy đầu ra logits từ mô hình giáo viên và mô hình học sinh. Sau đó, chúng ta chia mỗi đầu ra cho tham số temperature, tham số này kiểm soát mức độ quan trọng của mỗi mục tiêu mềm. Một tham số gọi là lambda cân bằng mức độ quan trọng của tổn thất chưng cất. Trong ví dụ này, chúng ta sẽ sử dụng temperature=5 và lambda=0.5. Chúng ta sẽ sử dụng tổn thất Divergence Kullback-Leibler để tính độ phân kỳ giữa mô hình học sinh và mô hình giáo viên. Với hai dữ liệu P và Q, Divergence KL giải thích lượng thông tin bổ sung cần thiết để biểu diễn P bằng Q. Nếu hai dữ liệu giống hệt nhau, độ phân kỳ KL của chúng bằng không, vì không cần thông tin nào khác để giải thích P từ Q. Do đó, trong ngữ cảnh của chưng cất tri thức, độ phân kỳ KL rất hữu ích.

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

Chúng ta sẽ đăng nhập vào Hugging Face Hub để có thể đẩy mô hình của mình lên Hugging Face Hub thông qua Trainer.

from huggingface_hub import notebook_login

notebook_login()

Hãy thiết lập TrainingArguments, mô hình giáo viên và mô hình học sinh.

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)

Chúng ta có thể sử dụng hàm compute_metrics để đánh giá mô hình của chúng ta trên tập kiểm tra. Hàm này sẽ được sử dụng trong quá trình huấn luyện để tính toán accuracy & f1 của mô hình.

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

Hãy khởi tạo Trainer với các đối số huấn luyện mà chúng ta đã định nghĩa. Chúng ta cũng sẽ khởi tạo bộ sắp xếp dữ liệu (data collator) của mình.

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
)

Bây giờ chúng ta có thể huấn luyện mô hình của mình.

trainer.train()

Chúng ta có thể đánh giá mô hình trên tập kiểm tra.

trainer.evaluate(processed_datasets["test"])

Trên tập kiểm tra, mô hình của chúng ta đạt độ chính xác 72 phần trăm. Để kiểm tra tính hợp lý về hiệu quả của quá trình chưng cất, chúng tôi cũng đã huấn luyện MobileNet trên bộ dữ liệu beans từ đầu với cùng các siêu tham số và quan sát được độ chính xác 63 phần trăm trên tập kiểm tra. Chúng tôi mời độc giả thử nghiệm các mô hình giáo viên được huấn luyện trước khác, kiến trúc học sinh khác, các tham số chưng cất khác và báo cáo kết quả của họ. Nhật ký huấn luyện và checkpoint cho mô hình đã chưng cất có thể được tìm thấy trong bộ lưu trữ này, và MobileNetV2 được huấn luyện từ đầu có thể được tìm thấy trong bộ lưu trữ này.

Update on GitHub