[特殊字符] Transformers 图像分类知识蒸馏实战：用 Trainer 将 ViT 教师模型蒸馏到 MobileNetV2

Transformers 图像分类知识蒸馏实战用 Trainer 将 ViT 教师模型蒸馏到 MobileNetV2【免费下载链接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers本指南演示如何基于 Transformers 的TrainerAPI将一个在 Beans 数据集上微调好的 ViT 图像分类模型教师蒸馏到一个随机初始化的 MobileNetV2学生模型。你将学会如何准备与预处理数据集、如何通过重写Trainer.compute_loss()实现基于 KL 散度的蒸馏损失、如何配置TrainingArguments并完成训练、评估与推送到 Hugging Face Hub以及如何用基线对比验证蒸馏的真实收益。知识蒸馏Knowledge Distillation核心概念知识蒸馏是一种模型压缩与知识迁移技术从更大、更复杂的模型教师teacher向更小、更简单的模型学生student传递知识。该思想最早由 Hinton 等人在论文Distilling the Knowledge in a Neural Network中提出。其核心直觉是教师模型输出的**软目标soft target即经过温度缩放的概率分布**比硬标签携带更丰富的类间相似性信息——例如一张豆子图片在健康类上概率高、在锈病类上概率中等这种分布形态正是学生模型需要模仿的知识。本指南执行的是任务特定的知识蒸馏task-specific distillation教师与学生在同一任务图像分类上对齐。流程为取一个在某任务上完成训练的预训练教师模型本例为merve/beans-vit-224基于google/vit-base-patch16-224-in21k在 Beans 数据集上微调而来随机初始化一个学生模型本例为 MobileNetV2同样面向图像分类训练学生模型使其输出分布与教师输出分布的差异最小化从而模仿教师的行为。蒸馏核心损失由两项加权合成学生与教师分布之间的KL 散度蒸馏损失与学生的真实标签交叉熵损失前者传递教师知识后者保证学生不偏离真实任务。说明原指南发布于 2023 年文档中的部分 API如image_processor、eval_strategy、report_to参数名在仓库当前版本中已有演进。本文以当前仓库源码为准对示例做了适配如使用processing_class替代旧式image_processor传参、统一使用eval_strategy与report_totensorboard并标注了文档原文与当前实现的差异便于读者对照。环境准备与依赖安装蒸馏与评估过程所需库如下pip install transformers datasets accelerate tensorboard evaluate --upgradetransformers模型、Trainer、TrainingArguments与图像处理器等核心组件datasets加载与处理 Beans 数据集accelerateTrainer的底层训练加速后端设备放置、混合精度等tensorboard训练日志可视化对应TrainingArguments中的report_totensorboardevaluate加载accuracy等评估指标。仓库中examples/pytorch/image-classification目录run_image_classification.py提供了可直接运行的图像分类训练脚本可结合本指南对照学习。本文示例为单机场景如需分布式训练背景知识可参考仓库中的 分布式训练示例。加载数据集与图像预处理使用datasets加载 Beans 数据集from datasets import load_dataset dataset load_dataset(beans)Beans 是一个植物病害图像分类数据集包含train、validation、test三个划分标签为三类叶片病害如角斑病、锈病、健康。图像预处理直接复用教师模型的处理器即可。本例中教师ViT与学生MobileNetV2的处理器在相同输入分辨率下返回相同输出因此任选其一均可这里使用教师的处理器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, batchedTrue)dataset.map(process, batchedTrue)会对数据集的每个划分train/validation/test批量应用预处理AutoImageProcessor内部完成图像缩放、归一化ViT 默认使用 ImageNet 均值/方差并输出模型所需的pixel_values张量。在当前仓库中AutoImageProcessor已被逐步统一到AutoProcessor见 processing_utils.py 中的处理器自动加载逻辑但图像分类场景下两者均可使用。设计蒸馏训练器重写 Trainer 的 compute_loss我们的目标让随机初始化的 MobileNet 模仿微调后的 ViT。实现方式是继承Trainer并重写compute_loss()在每一步训练中分别获取教师与学生的 logits 输出用temperature温度缩放 logits 得到软目标soft target温度控制各软目标的重要程度用lambda蒸馏损失权重衡量蒸馏损失在总损失中的占比用KL 散度损失计算学生与教师分布的差异。关于 KL 散度给定两个分布 P 与 QKL 散度描述用 Q 表示 P 额外需要多少信息。若两者完全相同则 KL 散度为 0。在蒸馏语境下我们最小化学生分布表示教师分布所需的额外信息量从而让学生分布逼近教师分布。示例实现如下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_modelNone, student_modelNone, temperatureNone, lambda_paramNone, *args, **kwargs): super().__init__(modelstudent_model, *args, **kwargs) self.teacher teacher_model self.student student_model self.loss_function nn.KLDivLoss(reductionbatchmean) 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_outputsFalse): student_output self.student(**inputs) with torch.no_grad(): teacher_output self.teacher(**inputs) # 计算教师与学生的软目标 soft_teacher F.softmax(teacher_output.logits / self.temperature, dim-1) soft_student F.log_softmax(student_output.logits / self.temperature, dim-1) # 蒸馏损失乘以 temperature 的平方以补偿缩放 distillation_loss self.loss_function(soft_student, soft_teacher) * (self.temperature ** 2) # 真实标签损失由学生模型前向返回 student_target_loss student_output.loss # 最终损失 加权组合 loss (1. - self.lambda_param) * student_target_loss self.lambda_param * distillation_loss return (loss, student_output) if return_outputs else loss实现要点nn.KLDivLoss(reductionbatchmean)PyTorch 要求输入为 log 概率学生对数软目标log_softmax目标为概率教师软目标softmaxbatchmean按 batch 求平均配合温度平方缩放后数值稳定温度缩放与补偿logits / temperature使分布更平滑温度越高越均匀突出类间软信息同时蒸馏损失乘以temperature ** 2以抵消温度缩放带来的梯度幅度变化这是 Hinton 论文中的标准做法torch.no_grad()教师模型仅用于前向推理不参与反向传播因此置于no_grad上下文并设为eval()模式关闭 dropout 等总损失(1 - lambda) * 交叉熵 lambda * KLlambda0.5表示两者等权本示例中取temperature5、lambda0.5读者可自行调参对比效果。底层原理Trainer 如何调用 compute_loss在 trainer.py 中基类Trainer.compute_loss()默认将模型前向返回的lossoutputs[loss]作为训练损失直接返回。而在训练主循环training_steptrainer.py与评估循环evaluation_looptrainer.py中均通过self.compute_loss(...)获取损失当return_outputsTrue时返回(loss, outputs)元组其中outputs[1:]即 logits供评估阶段计算指标。因此子类只需重写compute_loss即可在不改动训练/评估循环的前提下自定义损失函数——这正是ImageDistilTrainer能无缝接入Trainer的原因。教师侧损失来源ViTForImageClassification.forward()modeling_vit.py在传入labels时通过self.loss_function(labels, logits, self.config)计算交叉熵并放进输出对象的loss字段MobileNetV2ForImageClassification.forward()modeling_mobilenet_v2.py同样在labels存在时返回loss。这就是student_output.loss真实标签损失的直接来源。登录 Hub 并配置训练参数登录 Hugging Face Hub以便通过Trainer将模型推送到 Hubfrom huggingface_hub import notebook_login notebook_login()配置TrainingArguments、教师模型与学生模型from transformers import AutoModelForImageClassification, MobileNetV2Config, MobileNetV2ForImageClassification training_args TrainingArguments( output_dirmy-awesome-model, num_train_epochs30, fp16True, logging_strategyepoch, eval_strategyepoch, save_strategyepoch, load_best_model_at_endTrue, metric_for_best_modelaccuracy, report_totensorboard, push_to_hubTrue, hub_strategyevery_save, hub_model_idrepo_name, ) num_labels len(processed_datasets[train].features[labels].names) # 初始化教师模型在 beans 上微调过的 ViT teacher_model AutoModelForImageClassification.from_pretrained( merve/beans-vit-224, num_labelsnum_labels, ignore_mismatched_sizesTrue ) # 从零训练学生模型随机初始化 MobileNetV2 student_config MobileNetV2Config() student_config.num_labels num_labels student_model MobileNetV2ForImageClassification(student_config)TrainingArguments关键参数说明参数值作用output_dirmy-awesome-model检查点与模型保存目录num_train_epochs30训练轮数fp16True启用半精度混合精度训练需 GPU 支持logging_strategy/eval_strategy/save_strategyepoch每个 epoch 记录日志、评估、保存检查点load_best_model_at_endTrue训练结束时加载验证集最优检查点metric_for_best_modelaccuracy以 accuracy 作为选优指标report_totensorboard日志上报到 TensorBoardpush_to_hub/hub_strategy/hub_model_idTrue/every_save/repo_name每次保存检查点即推送 Hub需已登录且repo_name已定义版本适配提示原文档使用的eval_strategy参数在当前仓库为推荐名称report_to原文为trackio本文统一采用生态更通用的tensorboard。若你的环境支持 trackio 等第三方跟踪器可按需替换。MobileNetV2Config 关键字段仓库源码从 configuration_mobilenet_v2.py 可见MobileNetV2Config()默认即对应google/mobilenet_v2_1.0_224结构核心字段包括num_channels默认3输入图像通道数image_size默认224输入分辨率与教师 ViT 的 224 一致因此两模型处理器可互换depth_multiplier默认1.0通道数缩放系数宽度因子expand_ratio默认6.0倒残差块首层输出通道 输入通道 × 扩展比output_stride默认32输入/输出特征图空间分辨率比值设为 8 或 16 时使用空洞卷积hidden_act默认relu6激活函数classifier_dropout_prob默认0.8分类头 dropout 概率num_labels分类类别数需手动设置为数据集的标签数本示例代码中通过student_config.num_labels num_labels覆盖。MobileNetV2ForImageClassification在 modeling_mobilenet_v2.py 中实现主干网络输出池化特征后经 dropout 与线性分类头得到 logits当config.num_labels 1时计算回归损失MSE大于 1 时计算分类交叉熵损失。教师模型加载时使用ignore_mismatched_sizesTrue用于忽略分类头尺寸与 Hub 上原模型不一致的层从而适配 Beans 的 3 类标签。定义评估指标compute_metrics函数用于在训练过程中计算模型的accuracyimport evaluate import numpy as np accuracy evaluate.load(accuracy) def compute_metrics(eval_pred): predictions, labels eval_pred acc accuracy.compute(referenceslabels, predictionsnp.argmax(predictions, axis1)) return {accuracy: acc[accuracy]}eval_pred来自评估循环预测 logits 通过np.argmax(..., axis1)转为类别索引再与真实标签比对得到准确率。该指标同时被metric_for_best_modelaccuracy用于挑选最优检查点。初始化蒸馏训练器并开始训练使用上面定义的训练参数初始化Trainer同时初始化数据整理器data collatorfrom transformers import DefaultDataCollator data_collator DefaultDataCollator() trainer ImageDistilTrainer( student_modelstudent_model, teacher_modelteacher_model, training_argstraining_args, train_datasetprocessed_datasets[train], eval_datasetprocessed_datasets[validation], data_collatordata_collator, processing_classteacher_processor, compute_metricscompute_metrics, temperature5, lambda_param0.5 )版本适配提示Trainer构造时传递处理器的参数在当前仓库中为processing_class旧版本为image_processor或tokenizer。ImageDistilTrainer.__init__通过**kwargs透传给基类Trainer因此processing_class、compute_metrics等参数会正常生效。开始训练trainer.train()在测试集上评估trainer.evaluate(processed_datasets[test])实验结果与蒸馏效率验证据原指南报告蒸馏后的 MobileNet 在测试集上达到72% 准确率作为蒸馏有效性的健康性检查sanity check使用相同超参数在 Beans 数据集上从零训练 MobileNet无教师监督测试集准确率仅为63%。两者对比说明蒸馏带来的约 9 个百分点的提升直接归因于教师软目标传递的知识蒸馏是在相同训练预算epoch、batch、学习率一致下实现的排除了超参数差异的干扰。蒸馏后的训练日志与检查点可参考 Hub 仓库教师-学生对merve/vit-mobilenet-beans-224从零训练的 MobileNetV2 见merve/resnet-mobilenet-beans-5系列。请注意上述数字为原指南作者在特定数据集划分与超参数下的实测结果仅供参考在不同环境复现时数值可能略有波动。建议读者自行尝试不同的预训练教师、学生架构与蒸馏参数temperature、lambda并对比从零训练基线以验证蒸馏收益。进阶方向进一步探索更换教师/学生架构教师可替换为其他在 Hub 上微调过的 ViT 变体或更大的 CNN学生可尝试 MobileNetV3、EfficientNet 等轻量架构验证压缩比与精度权衡调整蒸馏超参数增大temperature使软目标更平滑、更强调类间关系调节lambda平衡蒸馏损失与标签损失结合硬标签与软目标本示例已同时使用两者还可引入特征层蒸馏对齐中间特征图、attention 蒸馏等变体生产部署蒸馏后的轻量学生模型更易部署。仓库提供了模型导出支持见 exporters与 ONNX 等格式转换能力可将训练好的 MobileNetV2 导出用于推理服务。总结本文以完整可运行的代码路径演示了基于 TransformersTrainer的图像分类知识蒸馏通过继承Trainer并重写compute_loss()将 ViT 教师的软目标知识以 KL 散度形式注入 MobileNetV2 学生模型并结合真实标签损失进行联合优化。文中同时给出了数据预处理、训练配置、评估指标、结果对比等完整闭环并补充了Trainer.compute_loss调用链、ViTForImageClassification与MobileNetV2ForImageClassification损失计算实现、MobileNetV2Config关键字段等仓库源码级依据。掌握该方法后你可以将其推广到任意图像分类任务与任意教师-学生模型组合实现高效、低成本的模型压缩与知识迁移。核心文件索引蒸馏训练器基类trainer.pycompute_loss默认实现与重写约定教师模型实现modeling_vit.pyViTForImageClassification学生模型实现modeling_mobilenet_v2.pyMobileNetV2ForImageClassification学生模型配置configuration_mobilenet_v2.pyMobileNetV2Config图像分类参考脚本run_image_classification.py【免费下载链接】transformers Transformers: the model-definition framework for state-of-the-art machine learning models in text, vision, audio, and multimodal models, for both inference and training.项目地址: https://gitcode.com/GitHub_Trending/tra/transformers创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考