多彩编程 多彩编程MZPH · CODE BLOG
ARTICLE DETAIL

文章详情

深耕前端与后端开发技术的一线实战笔记与踩坑复盘。

DeiT:数据高效视觉Transformer训练全解析与实战指南

DeiT:数据高效视觉Transformer训练全解析与实战指南 1. 项目概述从ViT到DeiT数据效率的破局之路在计算机视觉领域Transformer架构的引入无疑是一场革命。Vision TransformerViT的出现打破了卷积神经网络CNN长期以来的统治地位证明了纯注意力机制在图像分类任务上同样能取得SOTA级别的性能。然而ViT的成功有一个非常苛刻的前提它需要海量的标注数据进行预训练比如在JFT-300M这种包含3亿张图像的超大规模私有数据集上。对于我们绝大多数研究者、工程师和爱好者来说这无异于一道难以逾越的鸿沟。我们既没有谷歌那样的算力集群更无法获取如此规模的私有数据。难道Transformer在视觉领域的应用注定只是大厂的“特权”吗Data-efficient Image TransformerDeiT的诞生正是为了回答这个问题。它的核心目标就是让Transformer模型能够在相对“平民化”的数据集如ImageNet-1K约130万张图像上从头开始训练并达到优异的性能。这不仅仅是模型架构的微调更是一套完整的训练策略革新。它通过引入“蒸馏令牌”Distillation Token和一系列精心设计的训练技巧让模型能够像学生一样从一位强大的“教师”模型通常是一个高性能的CNN那里高效地学习知识从而在数据有限的情况下实现性能的飞跃。简单来说DeiT解决了我们普通玩家“玩不起”Transformer的痛点让先进的视觉Transformer技术变得触手可及。无论你是想在自己的研究项目中尝试Transformer还是在业务场景中部署一个高效、强大的视觉模型DeiT都是一个极具吸引力的起点。2. 核心原理深度拆解知识蒸馏如何赋能Transformer要理解DeiT为何有效我们必须深入其核心机制基于令牌的知识蒸馏。这与我们熟知的ViT训练范式有根本性的不同。2.1 ViT的训练瓶颈与数据依赖标准的ViT模型将图像分割成固定大小的图像块Patches然后通过线性投影得到一系列令牌Tokens再加上一个可学习的分类令牌[CLS] token。模型通过多头自注意力机制MSA和MLP层处理这些令牌最终由[CLS]令牌的状态输出分类结果。这种架构的强大之处在于其全局建模能力但缺点也很明显Transformer缺乏CNN固有的归纳偏置如局部性、平移等变性。CNN的卷积核天生就假设相邻像素是相关的这个先验知识让CNN在中小规模数据上也能学得很好。而Transformer是一张“白纸”它需要从海量数据中自己学习到图像的所有空间结构规律因此对数据量的要求极高。2.2 DeiT的蒸馏令牌一个聪明的“学生”DeiT的创新点在于它在输入序列中额外引入了一个可学习的蒸馏令牌Distillation Token。这个令牌与[CLS]令牌地位平行一起参与Transformer编码器的前向传播。它的目标非常明确在整个训练过程中学习去匹配“教师”模型例如一个在ImageNet上预训练好的RegNet或EfficientNet的输出。这里的关键在于蒸馏信号不是简单地在最终损失函数上做文章。在传统的知识蒸馏中我们通常让学生模型的输出概率分布软标签去逼近教师模型的软标签同时结合真实标签的硬损失。而DeiT采用的是基于令牌的蒸馏。蒸馏令牌通过自注意力机制与图像块令牌、[CLS]令牌进行交互在整个网络的深层特征层面吸收教师模型的知识。最终这个蒸馏令牌的输出会与教师模型的预测进行计算产生蒸馏损失。注意DeiT论文中探讨了两种蒸馏方式硬标签蒸馏和软标签蒸馏。硬标签蒸馏是让蒸馏令牌去预测教师模型输出的硬标签即类别索引损失函数是交叉熵。软标签蒸馏则是让学生模型的输出分布经过温度系数调整的softmax去逼近教师模型的输出分布使用KL散度作为损失。实验表明硬标签蒸馏在ImageNet上表现更优且实现更简单。2.3 联合训练与损失函数在训练阶段DeiT模型同时接收两种监督信号真实标签监督由[CLS]令牌的输出与图像真实标签计算交叉熵损失。蒸馏监督由蒸馏令牌的输出与教师模型的预测硬标签或软标签计算损失交叉熵或KL散度。总损失函数是两者的加权和。在推理阶段我们可以选择使用[CLS]令牌的输出、蒸馏令牌的输出或者将两者平均作为最终预测。论文中发现直接使用蒸馏令牌的输出通常能获得最好的性能这印证了从强大教师那里学到的知识是非常有效的。这种设计的美妙之处在于它为Transformer模型注入了一种强先验知识。教师CNN例如RegNet在ImageNet上已经学会了非常鲁棒和通用的视觉特征。通过蒸馏令牌Transformer“学生”不必再从零开始摸索图像的基本规律而是可以直接学习这些高质量的、经过验证的特征表示从而极大地加速了收敛过程并提升了在有限数据下的最终性能。3. 训练策略与技巧全景解析除了核心的蒸馏架构DeiT的成功还离不开一整套“组合拳”式的训练策略优化。这些技巧很多并非DeiT首创但被系统地整合应用共同解决了Transformer在小数据上的训练难题。3.1 优化器与超参数调优Transformer模型对优化器和超参数非常敏感。DeiT采用了AdamW优化器它是在Adam的基础上加入了权重衰减Weight Decay的正确实现解耦了权重衰减和梯度更新对于Transformer这类模型训练更加稳定。几个关键超参数设置学习率与热身Warmup使用较大的批量大小如1024时需要线性的学习率热身例如前5个或10个epoch从0线性增长到基础学习率。这有助于在训练初期稳定优化过程。学习率调度采用余弦退火Cosine Annealing调度器。学习率随着训练过程从最大值平滑地衰减到0像余弦曲线的一半。这种方式比阶梯式下降Step Decay更平滑通常能帮助模型收敛到更优的局部最优点。权重衰减Weight Decay设置为一个相对较高的值如0.05。较强的权重衰减起到了正则化的作用防止模型在有限数据上过拟合这对于数据高效的训练至关重要。随机深度Stochastic Depth也称为层丢弃Layer Drop。在训练时以一定的概率随机“跳过”即直接使用恒等映射网络中的某些层。这相当于一种深度的数据增强和模型正则化能提高模型的泛化能力。对于DeiT随机深度的丢弃率随着网络深度增加而线性增加。3.2 强力的数据增强组合数据增强是弥补数据不足最直接有效的手段。DeiT的训练管道集成了当时最有效的几种增强方法RandAugment一种自动化的增强策略它从一系列基础的图像变换如旋转、剪切、颜色抖动等中随机选择几种并以随机的强度应用。它只有两个超参数变换的数量N和强度M调参简单且效果显著。MixUp将两张图像及其标签进行线性插值生成新的训练样本。这鼓励模型在类别之间进行线性行为提高了鲁棒性。CutMix将一张图像的部分区域裁剪掉并用另一张图像的对应区域补上标签则按区域面积比例进行混合。这比简单的裁剪能保留更多的信息同时提供了更强的正则化。重复增强Repeated Augmentation在一个大批量Batch中对同一张原始图像应用不同的随机增强产生多个变体然后放入同一个批次中进行训练。这相当于在批次级别增加了数据多样性对于小批量训练尤其有益。实操心得这些增强策略的强度需要仔细权衡。过强的增强如极高的颜色抖动强度、大幅度的裁剪可能会破坏图像语义反而损害模型学习。DeiT论文中给出的是一组在ImageNet上验证过的默认参数对于你自己的数据集可能需要适当调低强度尤其是当你的数据集与ImageNet差异较大时。3.3 模型结构微调DeiT的基础骨架与ViT保持一致但在一些细节上做了调整以优化性能和效率位置编码沿用ViT的可学习一维位置编码。尽管有研究认为视觉任务中相对位置编码或二维编码可能更好但可学习一维编码在DeiT的设定下已经足够且实现简单。分类头与ViT相同使用一个单独的线性层作为分类头。模型变体DeiT提供了几个不同规模的模型如DeiT-Tiny、DeiT-Small和DeiT-Base分别对应不同的参数数量和计算量方便在不同资源约束下使用。4. 从零开始训练DeiT完整实操指南理解了原理和策略后我们来看如何亲手训练一个DeiT模型。这里以在ImageNet-1K数据集上训练DeiT-Small为例使用PyTorch框架。4.1 环境准备与依赖安装首先确保你的Python环境建议3.8以上并安装核心库pip install torch torchvision pip install timm # PyTorch Image Models库包含了DeiT的官方实现 pip install tensorboard # 用于可视化训练过程可选 pip install pyyaml # 用于读取配置文件可选timm库是Ross Wightman维护的一个宝藏库集成了大量视觉模型的高效实现包括DeiT。我们将主要依赖它。4.2 数据准备假设你的ImageNet数据存放在如下目录结构/path/to/imagenet/ train/ n01440764/ n01440764_10026.JPEG ... n01443537/ ... val/ n01440764/ ...你需要准备一个指向该路径的环境变量或直接在代码中指定。4.3 训练脚本核心解析下面是一个基于timm库的简化版训练脚本核心部分。完整的训练脚本通常包含分布式训练、混合精度训练等这里为清晰起见先展示关键逻辑。import torch import torch.nn as nn import torch.optim as optim from torch.utils.data import DataLoader from torchvision import datasets, transforms import timm from timm.data import create_transform from timm.loss import LabelSmoothingCrossEntropy, SoftTargetCrossEntropy from timm.models import create_model # 1. 定义模型和教师模型 model_name deit_small_patch16_224 teacher_model_name regnety_160 # 示例教师模型需提前用timm下载权重 model create_model( model_name, pretrainedFalse, num_classes1000, drop_rate0.0, drop_path_rate0.1, # 随机深度 drop_block_rateNone, ) teacher_model create_model( teacher_model_name, pretrainedTrue, # 加载预训练权重 num_classes1000, ).eval() # 教师模型设为评估模式不更新参数 for param in teacher_model.parameters(): param.requires_grad False # 2. 定义数据增强和加载器 train_transform create_transform( input_size224, is_trainingTrue, color_jitter0.4, auto_augmentrand-m9-mstd0.5-inc1, # RandAugment interpolationbicubic, re_prob0.25, # Random Erasing概率 re_modepixel, re_count1, ) train_dataset datasets.ImageFolder(root/path/to/imagenet/train, transformtrain_transform) train_loader DataLoader(train_dataset, batch_size128, shuffleTrue, num_workers8, pin_memoryTrue) # 3. 定义损失函数和优化器 distillation_type hard # 硬标签蒸馏 if distillation_type soft: distillation_loss SoftTargetCrossEntropy() else: # hard distillation_loss nn.CrossEntropyLoss() cls_loss LabelSmoothingCrossEntropy(smoothing0.1) # 标签平滑的交叉熵用于[CLS]令牌 optimizer optim.AdamW(model.parameters(), lr5e-4, weight_decay0.05) scheduler optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max300) # 假设训练300个epoch # 4. 训练循环核心部分 model.train() for epoch in range(num_epochs): for images, targets in train_loader: images, targets images.cuda(), targets.cuda() # 前向传播 outputs model(images) # outputs 是一个元组或字典包含cls_logits和distill_logits # 假设timm的DeiT实现返回 (cls_logits, distill_logits) cls_logits, distill_logits outputs # 计算损失 loss_cls cls_loss(cls_logits, targets) with torch.no_grad(): teacher_outputs teacher_model(images) if distillation_type hard: teacher_labels teacher_outputs.argmax(dim1) # 取教师预测的硬标签 loss_distill distillation_loss(distill_logits, teacher_labels) else: # soft # 可能需要调整温度系数 loss_distill distillation_loss(distill_logits, teacher_outputs) # 联合损失alpha是平衡系数论文中通常设为0.5 alpha 0.5 loss alpha * loss_cls (1 - alpha) * loss_distill # 反向传播和优化 optimizer.zero_grad() loss.backward() optimizer.step() scheduler.step() # 每个epoch结束后可以在验证集上评估性能...4.4 关键参数与配置说明drop_path_rate随机深度的丢弃率。对于DeiT-Small0.1是一个典型值。模型越深这个值可以适当增大。auto_augmenttimm中集成了RandAugment的策略rand-m9-mstd0.5-inc1是DeiT论文中使用的配置。LabelSmoothingCrossEntropy标签平滑是一种正则化技术将真实标签的1稍微分散到其他类别上如0.9 vs 0.1/类别数-1防止模型对训练标签过于自信提升泛化能力。学习率与批量大小学习率5e-4是针对批量大小1024设置的。如果你使用的批量大小不同需要按线性缩放规则调整lr base_lr * batch_size / 1024。同时Warmup的epoch数也应考虑批量大小。教师模型选择论文中使用的是RegNetY-16GF。在timm中你也可以选择其他高性能的CNN作为教师如efficientnet_b3等。教师模型越强通常学生能学到的上限也越高。5. 实战中的常见问题与调优技巧在实际复现和运用DeiT时你可能会遇到以下几个典型问题。5.1 训练不稳定或发散现象训练初期损失值剧烈震荡甚至变成NaN。检查学习率与Warmup这是最常见的原因。确保使用了Warmup并且初始学习率设置正确。如果批量大小远小于1024学习率可能过高。检查梯度裁剪虽然AdamW相对稳定但对于非常深的Transformer在训练初期梯度爆炸仍有可能。可以添加梯度裁剪torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm1.0)作为保险。检查数据增强过强的颜色抖动或CutMix/MixUp比例可能导致个别样本的数值异常。可以暂时关闭这些增强观察训练是否稳定再逐步开启。混合精度训练使用AMPAutomatic Mixed Precision可以提升训练速度并有时能增加稳定性。但需注意在梯度裁剪时需要对未缩放的梯度进行操作。5.2 模型性能不及预期现象训练完成后在验证集上的准确率远低于论文报告的数据。数据预处理一致性这是最大的坑确保你的验证/测试数据预处理与训练时除了数据增强部分完全一致。包括相同的resize尺寸如256中心裁剪到224、相同的归一化均值标准差通常是ImageNet的[0.485, 0.456, 0.406]和[0.229, 0.224, 0.225]。一个像素值的偏差都可能导致精度显著下降。教师模型预测方式在计算蒸馏损失时确保教师模型处于eval()模式并且输入数据也经过了正确的预处理。教师模型本身的预处理逻辑可能与你的训练管道不同需要统一。训练周期不足DeiT需要较长的训练周期如300个epoch才能充分收敛。检查你的训练是否真的达到了足够的迭代次数。标签错误检查你的数据集标签是否正确。ImageNet的标签索引是固定的如果自定义数据集需确保类别顺序一致。5.3 推理速度慢或显存占用高现象模型部署时延迟高或训练时批量大小提不上去。使用更小的变体如果对精度要求不是极致可以尝试DeiT-Tiny或DeiT-Small它们在速度和精度间有更好的平衡。优化注意力计算对于大分辨率输入自注意力的计算复杂度是序列长度的平方。可以考虑使用torch.jit.script或torch.compilePyTorch 2.0对模型进行编译优化或者寻找支持高效注意力计算的库如xFormers。动态轴序列长度如果输入图像尺寸固定可以尝试将模型转换为TorchScript或ONNX并进行图优化可能获得推理加速。梯度检查点如果是为了训练更大的模型或更大的图像而显存不足可以使用梯度检查点技术用计算时间换显存空间。5.4 在自己的数据集上微调DeiT如果你有一个中等规模的自定义数据集几万到几十万张图想用DeiT作为骨干网络进行微调建议如下加载预训练权重使用timm.create_model(deit_small_patch16_224, pretrainedTrue, num_classes你的类别数)。timm会自动处理分类头的适配。调整数据增强降低增强强度。对于医学影像、卫星图像等专业领域RandAugment、ColorJitter等增强可能不适用需要设计领域特定的增强如弹性形变、随机噪声等。使用更小的学习率微调时学习率通常设置为初始训练时的1/10或1/100。可以使用分层学习率让靠近输出的层学习率大一些靠近输入的层学习率小一些。考虑只微调部分层如果数据量非常少可以冻结Transformer编码器的大部分层只训练最后的分类头以及最后几层的参数防止过拟合。DeiT的成功不仅仅是一个模型的成功更是一种范式的证明通过巧妙的训练策略和知识蒸馏我们可以让数据饥渴的Transformer模型变得“勤俭持家”。它极大地降低了视觉Transformer的应用门槛为后续更多高效Transformer模型如Swin Transformer、PVT等的研究铺平了道路。当你下次面对一个数据量有限但想尝试先进架构的视觉任务时不妨将DeiT及其背后的训练“配方”作为你的第一块试金石。
返回列表