Deit:知识蒸馏与vit的结合 学习笔记（附代码）

?论文地址：https://arxiv.org/abs/2012.12877

代码地址：GitHub - facebookresearch/deit: Official DeiT repository

1.是什么？

DeiT（Data-efficient Image Transformer）是一种用于图像分类任务的神经网络模型，它基于Transformer架构。这个模型的主要目标是在参数较少的情况下实现高效的图像分类。相比于传统的卷积神经网络（CNN），DeiT采用了Transformer的注意力机制，使其能够更好地捕捉图像中的全局关系。
以下是DeiT模型的一些关键特点和组成部分：

1. 1.Transformer 架构： DeiT采用了Transformer的架构，这是一种自注意力机制的模型。这种架构在自然语言处理任务中取得了显著的成功，DeiT将其成功地应用于图像分类领域。
2. 2.小模型参数： 为了提高数据效率，DeiT设计为具有相对较少的参数。这使得模型在训练和推理时需要更少的计算资源。
3. 3.Knowledge Distillation： DeiT使用知识蒸馏（Knowledge Distillation）的方法进行训练。这意味着它通过从一个大型预训练模型中传递知识来训练，而不是从头开始训练。这有助于在资源受限的情况下实现更好的性能。
4. 4.Patch Embedding： 与传统的卷积层不同，DeiT使用了补丁嵌入（Patch Embedding）来将图像分割成小块，然后对这些块进行变换。
5. 5.Positional Embeddings： 由于Transformer不涉及卷积层，它需要一种处理输入序列的方式。在DeiT中，位置嵌入（Positional Embeddings）用于为模型提供输入中元素的相对位置信息。

总体而言，DeiT是一个旨在通过Transformer的优势实现图像分类的轻量级模型，适用于数据受限的情况。通过知识蒸馏和小模型参数，它在参数较少的情况下达到了令人满意的性能。

2.为什么？

Transformer的输入是一个序列（Sequence），ViT 所采用的思路是把图像分块（patches），然后把每一块视为一个向量（vector），所有的向量并在一起就成为了一个序列（Sequence），ViT 使用的数据集包括了一个巨大的包含了 300 million images的 JFT-300，这个数据集是私有的，即外部研究者无法复现实验。而且在ViT的实验中作者明确地提到：

意思是当不使用 JFT-300 大数据集时，效果不如CNN模型。也就反映出Transformer结构若想取得理想的性能和泛化能力就需要这样大的数据集。DeiT 作者通过所提出的蒸馏的训练方案，只在 Imagenet 上进行训练，就产生了一个有竞争力的无卷积 Transformer。

3.怎么样？

在 DeiT 模型中，首先需要一个强力的图像分类模型作为teacher model。然后，引入了一个 Distillation Token，然后在 self-attention layers 中跟 class token，patch token 在 Transformer 结构中不断学习。Class token的目标是跟真实的label一致，而Distillation Token是要跟teacher model预测的label一致。蒸馏过程如下图所示。

3.1知识蒸馏

知识蒸馏（Knowledge Distillation）是一种模型训练的技术，旨在通过传递一个大型教师模型的知识来训练一个小型学生模型。这个方法的目标是使得学生模型能够获得与教师模型相似的性能，同时减少学生模型的复杂性和计算成本。
以下是知识蒸馏的关键思想和步骤：

1. 1.教师模型： 首先，有一个在任务上表现良好的大型教师模型。这个模型通常拥有更多的参数和计算能力，以便更好地捕捉任务的复杂性和结构。
2. 2.软目标（Soft Targets）： 在传统的监督学习中，模型通常以硬标签（one-hot编码的标签）作为目标进行训练。而在知识蒸馏中，使用了软目标，这是由教师模型输出的概率分布。这样的软目标包含了关于样本的更丰富信息，使得学生模型可以学到更多的任务相关知识。
3. 3.温度参数： 软目标的概率分布可以通过温度参数进行调节。较高的温度使概率分布更平滑，有助于学生模型更好地学到教师模型的知识。
4. 4.学生模型： 有了教师模型和软目标，接下来就是训练学生模型。学生模型通常是一个比教师模型简化的小型模型，可以在资源受限的环境中更轻松地进行推理。
5. 5.蒸馏损失： 为了引导学生模型学习教师模型的知识，引入了蒸馏损失。这个损失函数用于比较学生模型的输出概率分布和教师模型的输出概率分布，促使学生模型模仿教师模型的行为。

知识蒸馏的优势在于，通过传递教师模型的知识，可以在小型模型上实现接近教师模型性能的效果。这对于移动设备、嵌入式系统或其他计算资源受限的环境中的部署非常有用。

具体方法：

第一步是训练Net-T；第二步是在高温 T 下，蒸馏 Net-T 的知识到 Net-S。

训练 Net-T 的过程很简单，而高温蒸馏过程的目标函数由distill loss（对应soft target）和student loss（对应hard target）加权得到：

Deit 中使用 Conv-Based 架构作为教师网络，以 soft 的方式将归纳偏置传递给学生模型，将局部性的假设通过蒸馏方式引入 Transformer 中，取得了不错的效果。

3.2Distillation Token

Distillation Token 和 ViT 中的 class token 一起加入 Transformer 中，和class token 一样通过 self-attention 与其它的 embedding 一起计算，并且在最后一层之后由网络输出。

而 Distillation Token 对应的这个输出的目标函数就是蒸馏损失。Distillation Token 允许模型从教师网络的输出中学习，就像在常规的蒸馏中一样，同时也作为一种对class token的补充。

3.3代码实现

class DistilledVisionTransformer(VisionTransformer):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.dist_token = nn.Parameter(torch.zeros(1, 1, self.embed_dim))
        num_patches = self.patch_embed.num_patches
        self.pos_embed = nn.Parameter(torch.zeros(1, num_patches + 2, self.embed_dim))
        self.head_dist = nn.Linear(self.embed_dim, self.num_classes) if self.num_classes > 0 else nn.Identity()

        trunc_normal_(self.dist_token, std=.02)
        trunc_normal_(self.pos_embed, std=.02)
        self.head_dist.apply(self._init_weights)

    def forward_features(self, x):
        # taken from https://github.com/rwightman/pytorch-image-models/blob/master/timm/models/vision_transformer.py
        # with slight modifications to add the dist_token
        B = x.shape[0]
        x = self.patch_embed(x)

        cls_tokens = self.cls_token.expand(B, -1, -1)  # stole cls_tokens impl from Phil Wang, thanks
        dist_token = self.dist_token.expand(B, -1, -1)
        x = torch.cat((cls_tokens, dist_token, x), dim=1)

        x = x + self.pos_embed
        x = self.pos_drop(x)

        for blk in self.blocks:
            x = blk(x)

        x = self.norm(x)
        return x[:, 0], x[:, 1]

    def forward(self, x):
        x, x_dist = self.forward_features(x)
        x = self.head(x)
        x_dist = self.head_dist(x_dist)
        if self.training:
            return x, x_dist
        else:
            # during inference, return the average of both classifier predictions
            return (x + x_dist) / 2

参考：ViT、Deit这类视觉transformer是如何处理变长序列输入的?

DeiT：使用Attention蒸馏Transformer