【读点论文】torchdistill: A Modular, Configuration-Driven Framework for Knowledge Distillation，文件配置完成知识蒸馏

羞儿

于 2024-07-17 18:34:39 发布

阅读量549

点赞数 13

分类专栏：论文笔记文章标签：知识蒸馏深度学习开源工具

本文链接：https://blog.csdn.net/weixin_43424450/article/details/140502497

版权

论文笔记专栏收录该内容

137 篇文章 31 订阅

订阅专栏

torchdistill: A Modular, Configuration-Driven Framework for Knowledge Distillation

Abstract

虽然知识蒸馏（转移）一直受到研究界的关注，但该领域的最新发展凸显了对可重复研究和高度通用框架的需求，以降低此类高质量、可重复的深度学习研究的门槛。一些研究人员自愿发布了他们在知识蒸馏研究中使用的框架，以帮助其他感兴趣的研究人员重现他们的原创工作。然而，这样的框架通常既不通用也不维护，因此研究人员仍然需要编写大量代码来重构/构建框架，以引入新方法、模型、数据集和设计实验。在本文中，我们介绍了我们开发的基于 PyTorch 的、专用于知识蒸馏研究的开源框架。该框架旨在使用户能够通过声明式 PyYAML 配置文件设计实验，并帮助研究人员完成最近提出的 ML 代码完整性检查表。使用开发的框架，我们展示了其各种有效的训练策略，并实现了各种知识蒸馏方法。我们还复现了他们在 ICLR、NeurIPS、CVPR 和 ECCV 等主要机器学习会议上展示的 ImageNet 和 COCO 数据集上的一些原始实验结果，包括最近的最先进方法。所有源代码、配置、日志文件和经过训练的模型权重均可在 https://github.com/yoshitomo-matsubara/torchdistill 上公开获取。
论文地址：[2011.12913] torchdistill: A Modular, Configuration-Driven Framework for Knowledge Distillation (arxiv.org)

Introduction

深度学习方法已经取得了最先进的性能，并促进了各种任务应用程序的快速发展，例如图像分类和对象检测。此类最先进模型的关键问题之一是其复杂性，因此复杂模型难以部署用于实际应用。通常，模型复杂性和推理性能（例如以准确率衡量）之间存在权衡，并且有三种不同类型的方法可以使模型可部署：1）设计轻量级模型，2）模型压缩/修剪，以及 3）知识提炼。与复杂模型（例如，ResNet 和 Mask R-CNN ）相比，轻量级模型（例如 MobileNet 、MnasNet 和 YOLO 系列）通常会牺牲推理性能来减少推理时间。模型压缩和修剪技术通过量化参数和修剪冗余神经元来减小模型大小，这类方法在模型压缩的开源库 Distiller [Neural Network Distiller: A Python Package for DNN Compression Research] 中有所介绍。
在本文中，我们重点关注最后一类知识蒸馏，即训练一个更简单的（学生）模型来模仿强大的（老师）模型的行为。知识蒸馏源自 Buciluaˇ 等人 [Model compression] 的研究，他们提出了一种将大型复杂集成压缩为较小模型的方法，同时推理性能损失很小。有趣的是， [Do deep nets really need to be deep?] 报告称，训练模仿老师模型行为的学生模型（软标签）明显优于在原始（硬标签）数据集上训练的学生模型。在这些研究之后，知识蒸馏和转移引起了计算机视觉和自然语言处理等研究界的关注。
如下表所示，一些研究人员自愿发布他们的知识蒸馏框架，以帮助其他研究人员重现他们最初的研究。然而，这样的框架通常要么泛化能力不强，要么难以维护。此外，Distiller 仅支持一种知识蒸馏方法，而 Catalyst 是一个基于 PyTorch 构建的框架，专注于深度学习研究的可重现性。为了支持各种深度学习方法，这些框架具有很好的泛化能力，但却需要用户硬编码（重新实现）关键模块（如模型和数据集），即使这些实现在流行的库中是公开的，以设计复杂的知识蒸馏实验。正如 [AllenNLP] 指出的那样，参考方法和模型通常是从头开始重新实现的，这使得很难重现报告的结果。因此，为了进一步推进深度学习研究，需要一个新的通用框架，该框架应该能够让研究人员轻松尝试不同的模块（例如模型、数据集、损失配置），实施各种方法，并保证其工作的可重复性。
- 知识蒸馏框架。torchdistill 支持 PyTorch 和 torchvision 中的模块，例如损失、数据集和模型。ImageNet：ILSVRC 2012 、YT Faces：YouTube Faces DB 、MIT Scenes：室内场景数据集、CUB-2011：Caltech-UCSD Birds-200-2011 、Cars：汽车数据集、SOP：斯坦福在线产品。P：预训练模型，M：模块抽象，D：分布式训练。
我们的框架 torchdistill 的概念深受 AllenNLP 的启发，AllenNLP 是一个基于 PyTorch 构建的平台，用于研究自然语言处理中的深度学习方法。与 AllenNLP 类似，torchdistill 支持以下功能：
- 模块抽象使研究人员能够为实验编写更高级的代码，例如模型、数据集、优化器和损失；
- 声明性的 PyYAML 配置文件，可以看作是实验（训练和评估）的高级摘要，允许使用文件中的锚点和别名来引用同一个对象（例如文件路径）并简化自身，并可以轻松更改抽象的组件和超参数；
- 通用参考代码和配置，将知识蒸馏方法应用于在著名复杂基准数据集上预训练的 PyTorch 和 torchvision 模型：ImageNet（ILSVRC 2012）和 COCO 2017 。
此外，torchdistill 支持 1）无缝多阶段训练，2）缓存教师输出，3）无需硬编码（重新实现）即可重新设计（修剪）教师和学生模型。据我们所知，这是第一个高度通用的开源框架，可以支持各种知识提炼方法，并降低高质量、可重复的深度学习研究的门槛。研究人员可以探索方法并塑造新方法，在此通用框架的基础上，不仅可以轻松定制现有方法和模型，还可以引入全新的方法和模型。使用我们重新实现的一些方法，我们还重现了原始研究中报告的 ILSVRC 2012 和 COCO 2017 数据集上的实验结果。

Framework Design

我们开发的框架 torchdistill 是一个基于 PyTorch 构建的开源框架，专门用于知识蒸馏研究。对于图像分类和对象检测等视觉任务，该框架旨在支持 torchvision，它为数据集、模型架构和常见图像转换提供了大量选项。我们框架中支持的参考模型和数据集的集合取决于用户安装的 torchvision 版本。例如，当用户在最新的 torchvision 中找到新模型时，他们只需更新 torchvision 和配置文件即可使用我们的框架进行实验，从而快速尝试这些模型。

Module Abstractions

我们框架中模块抽象的一个目标是让研究人员只需更改第 2.3 节中描述的 PyYAML 配置文件即可试验各种模块。我们将抽象重点放在要进行实验的关键模块上，特别是模型架构、数据集、转换和训练期间要最小化的损失。这些模块通常在作者发布的框架中是硬编码的（参见附录 A），许多超参数也是硬编码的。
模型架构：torchvision 为视觉任务提供了各种模型系列，从 AlexNet 到 R-CNN ，其中许多模型都是在大型基准数据集上进行预训练的。具体来说，最新版本 (v0.8.2) 提供了大约 30 个在 ImageNet (ILSVRC 2012) 上预训练的图像分类模型和 4 个在 COCO 2017 上预训练的对象检测模型。由于我们的框架支持 torchvision 用于视觉任务，研究人员可以使用此类预训练模型作为教师和/或基线模型（例如，无需教师训练的学生）。除了 torchvision 中提供的预训练模型外，他们还可以使用自己的预训练模型权重和使用 PyTorch 实现的任何模型架构。此外，torchdistill 支持 PyTorch Hub【PyTorch Hub | PyTorch】，并允许用户通过在 PyYAML 配置文件中指定存储库名称来通过中心导入模块。
数据集：如上所述，torchvision 还支持各种数据集，先前的研究使用其中许多数据集来验证提出的蒸馏技术，例如 ImageNet 、COCO 、CIFAR-10 和 -100 和 Caltech101 。与模型架构类似，torchdistill 支持此类数据集，并且可以与使用 PyTorch 实现的任何数据集协作。
Transforms：在视觉任务中，存在事实上的标准图像变换技术。以 ImageNet 数据集上的图像分类为例，使用 torchvision 进行训练的标准变换管道包括：1）裁剪原始大小的随机大小和原始长宽比的随机长宽比，2）水平反射，以 50% 的概率进行数据增强，以降低过度拟合的风险，3）PIL 到张量转换，4）分别使用 (0.485, 0.456, 0.406) 和 (0.229, 0.224, 0.225) 作为均值和标准差进行通道规范化。在 torchdistill 中，用户可以在配置文件中定义自己的变换管道。
损失：在蒸馏过程中，学生模型使用教师模型的输出进行训练，研究界已经提出了许多独特的损失，包括/不包括特定于任务的损失，例如用于分类任务的交叉熵损失。PyTorch 支持各种损失类/函数，可以使用 torchdistill 的可定制损失模块组合这些支持的损失，在配置文件中定义简单的蒸馏损失（参见第 2.6 节）。

Registry

注册表是 torchdistill 中的一个重要组件，因为抽象模块是通过将配置文件中的字符串映射到代码中的对象来实例化的。此外，它还使用户可以轻松地将他们实现的模块/功能与此框架协作。与 AllenNLP 和 Catalyst 类似，这甚至可以在框架之外通过使用 Python 装饰器来完成。以下示例显示，只需使用 @register model（在框架中定义）即可将新模型类 MyModel 添加到框架中，并且可以通过在配置文件中的指定位置使用所需参数定义“MyModel”来实例化新类。
- ```
@register model
class MyModel(nn.Module):
    def init (self, *args, **kwargs):
        super(). init ()
        self.conv1 = nn.Conv2d(**kwargs[’conv1 kwargs’])
        ...
```

Configurations

实验可以通过 PyYAML 配置文件（参见附录 B）定义，用户无需硬编码即可调整超参数并更改方法 / 模型。借助 PyYAML 的功能，配置文件允许用户利用锚点和别名，这些功能有助于简化配置，以防用户想要重用配置文件中定义的参数（例如数据集的根目录路径、参数和模型名称作为检查点文件路径的一部分以更好地管理数据。在配置文件中，需要定义三个主要组件：数据集、教师和学生模型以及训练。每个关键组件都使用第 2.1 和 2.2 节中描述的抽象和注册模块来定义。配置文件为用户提供了实验摘要，并显示了重现实验结果的所有参数（实验所用的硬件规格等隐含因素除外）。
以下示例说明如何在 PyYAML 配置文件中定义全局教师模型。如前几节所述，我们的框架中抽象了各种类型的模块，并在用户安装的 torchvision 中注册了此类模块（类和函数）。在此示例中，’resnet34’ 函数用于通过使用关键字参数字典（**params）来实例化 ResNet 类型的对象。即 num classess = 1000 和 pretrained = True 作为 ’resnet34’ 函数的参数给出。对于在 torchvision 中实现的图像分类模型或用户添加到我们框架中的注册表中的模型，用户可以通过更改 ’resnet34’ 轻松尝试不同的模型，例如 ’densenet201’ 、’mnasnet1_0’ 。此外，ckpt 表示检查点的文件路径，即示例中的 ’./resnet34.pt’，它利用了 YAML 的一些特性：锚点 (&) 和别名 (*)。对于教师模型，如果检查点文件存在，则将使用用户自己的模型权重初始化模型。否则，本示例中的 ’resnet34’ 将使用 torchvision 为 ILSVRC 2012 预训练的权重进行初始化。
- ```
teacher_model:
	name: &teacher_model ’resnet34’
	params:
        num_classes: 1000
        pretrained: True
    ckpt: !join [ ’./’ , *teacher , ’.pt’ ]
```
此外，torchdistill 还提供了生成日志文件来监控实验的选项。例如，日志文件会显示使用了哪些参数、执行时间、配置文件中设置的频率下的训练行为趋势（例如训练损失、学习率和验证准确率）以及评估结果。
这些配置和日志文件还将帮助研究人员完成ML代码完整性检查表【GitHub - paperswithcode/releasing-research-code: Tips for releasing research code in Machine Learning (with official NeurIPS 2020 recommendations)】，该检查表最近被提出作为主要机器学习会议（如NeurIPS、ICML和CVPR）的官方代码提交过程的一部分，以促进研究社区的可重复性。

Dataset Wrappers

为了支持各种各样的知识蒸馏方法，数据集是一个需要推广的重要模块。通常，PyTorch 和 torchvision 中的数据集模块在每次迭代时都会返回一对输入批次（例如，整理后的图像张量）和目标（ground-truth），但现有的一些知识蒸馏方法需要批次的额外信息。例如，对比表示蒸馏 (CRD) 需要一种有效的策略来在训练会话中检索大量负样本，这需要数据集模块返回一个额外的对象（例如，负样本索引）。为了支持这样的扩展，我们设计了数据集包装器来返回输入批次、目标和补充词典，不使用时可以为空。对于上述情况，附加对象可以存储在补充词典中，并在计算对比损失时使用。这种设计还使我们能够支持根据原始数据集中的数据索引缓存教师模型的输出，以便可以通过在第一个时期给定数据索引缓存（序列化）教师模型的输出，并在接下来的时期给定一批数据索引读取和整理缓存的输出来跳过教师的推理。
为了证明缓存可以提高训练效率，我们进行了一个知识蒸馏实验，如下图a 所示，该实验缓存了在 ILSVRC 2012 数据集上训练 ResNet-18（学生）的第一个 epoch 的教师模型输出，并通过加载和将缓存在磁盘上的输出提供给损失模块来跳过教师模型的推理。
- （a）知识提炼，分别使用 ResNet34 和 ResNet-18 作为老师和学生模型。
- （b）使用辅助模块（卷积回归器）进行提示训练作为 FitNet 方法的第 1 阶段。其第 2 阶段是知识提炼，如图 1a 所示。
- 知识提炼和 FitNet 方法。黄色和蓝色模块分别表示其参数是冻结的和可训练的。
下表表明，在第 1 个 epoch 多花一分钟来序列化教师的输出，当使用 3 个 NVIDIA GeForce RTX 2080 Ti 且批次大小为 256 时，缓存策略使接下来的训练过程（即从第 2 个 epoch 开始）在 epoch 级别上快大约 1.23 – 2.11 倍。此外，当使用更大的教师模型（如 ResNet-152）时，这种改进变得更加显著（比没有缓存的训练快大约 2.11 倍）。 ILSVRC 2012 训练数据集包含大约 130 万张图片，缓存文件仅占用 10GB，而原始训练数据集占用约 140GB。请注意，如果要缓存的教师输出大得多，则缓存可能不会提高训练效率，例如基于提示的训练需要教师和学生模型的中间输出。此外，在应用数据增强策略时应关闭此模式
- 通过在第 1 个 epoch 缓存教师输出来提高 epoch 级别的训练速度，使用 ResNet-18 作为学生模型进行知识提炼。

Teacher and Student Models

师生对是知识蒸馏实验的关键，最近提出的方法引入了辅助模块，这些模块仅在训练过程中使用。此类辅助模块使用模型中间层的张量，将这些模块引入模型通常会导致其前馈路径出现分支，如上文图所示。然而，这种范式也是研究人员决定对模型进行硬编码（例如，每次在初步实验中更改辅助模块的位置时，都要修改 torchvision 中模型的原始实现）的背景之一，以引入用于他们所提出方法的此类辅助模块，并使其他研究人员难以在已发布的框架上进行构建。
torchdistill 利用 PyTorch 中的前向钩子范式，支持在不改变模型原始实现的情况下引入此类辅助模块。具体来说，用户可以将框架提供的前向钩子注册到特定模块，通过在配置文件中指定模块路径（例如，第 2.2 节中 MyModel 对象的“conv1”）将其输入和/或输出存储在 I/O 字典中。教师和学生模型的 I/O 字典将被输入到第 2.6 节中描述的通用、可定制的损失模块。
对于不仅需要提取中间输出（见上文图）而且还需要将提取的输出馈送到不同分支中可训练的辅助模块进行处理（见图 2b）的方法，我们在框架中定义了一个特殊模块，该模块旨在具有后前向功能。例如，在图 1 中，框架首先执行 ResNet-18 并通过注册的前向挂钩提取中间输出，然后存储在学生的 I/O 字典中的提取输出将作为后前向过程的一部分馈送到回归器。特殊模块的概念为用户在设计训练方法时提供了更大的灵活性，同时保持模型的原始实现（图 2 中的 ResNet-34 和 ResNet-18）不变。

Customizable Loss Module

利用包含已注册前向钩子的特定模块的输入/输出的 I/O 字典，torchdistill 提供了一个广义的可定制损失模块，允许用户通过配置文件（如图 2b 所示）轻松地将不同的损失模块与平衡因子组合起来。给定一对输入 x 和真实值 y，I/O 字典由一组键 J 以及分别从学生和教师模型中提取的值 $z^S_j$ 和 $z^T_j（j\in J）$ 组成。使用 I/O 字典和真实值，广义损失定义为
- $\sum_{j\in J} λ_j · L_j (z^S_j ; z^T_j ; y); (1)$
- 其中 λj 是 Lj 的平衡权重（超参数），它要么是 PyTorch 中实现的损失模块，要么是注册表中用户定义的损失模块。
例如，在因子转移第二阶段（图 2b）中，在 ILSVRC 2015 数据集上训练学生模型的损失函数可以定义为：
- 其中 $λ_{cls} = 1，λ_{FT} = 1000 且 p = 1$ ，遵循 [Paraphrasing complex network]。

Stage-wise Training Configuration

在前面的部分中，我们描述了 torchdistill 的主要功能，以及框架中可配置的模块。我们强调，上述所有训练配置都可以分阶段定义。无缝多阶段训练配置：具体来说，该框架旨在让用户配置关键组件，例如 1）epoch 数量、2）训练和验证数据集、3）教师和学生模型、4）要训练/冻结的模块（层）、5）优化器、6）学习率调度程序、7）损失模块。这些组件可以在每个训练阶段重新定义，否则框架将重用上一阶段的组件。请注意，这些训练配置可以在配置文件中声明，这种设计不仅可以支持两阶段训练策略，还可以支持更复杂的蒸馏方法，例如教师助理知识蒸馏 (TAKD) ，它可以训练教师来填补学生和教师模型之间的空白。迁移学习也可以通过从一个阶段更改模型和数据集来支持，用户只需使用配置文件执行一次代码即可。因此，他们不需要多次执行代码来执行多阶段训练，包括迁移学习。
重新设计模型以实现高效训练：此外，通过在配置文件中指定所需的模块，我们的框架为用户提供了在每个阶段重新设计教师和学生模型的选项。具体来说，用户可以通过使用辅助模块重新使用模型中的模块来重建模型。图 1 显示了一个例子，教师和学生模型的第 8 个和第 5 个块之后的模块可以分别被修剪，因为模块的输出不用于提示训练（第 1 阶段），因此不需要执行。在这个特定案例中，重新设计的学生模型将由可训练模块（蓝色）和回归器（辅助模块）组成，如下图所示，而第二阶段的教师和学生架构将恢复为原始架构（图 1a），其中包含在第一阶段学习到的参数。此外，重新设计的教师/学生模型可以是一个空模块，以节省执行时间。例如，在图 2a 中，不需要将输入批次提供给学生模型（因此可以为空），因为在因子转移的第一阶段，只执行教师模型来训练释义器。
- 与图 1 中的简单配置相比，只需在配置文件中指定所需模块即可对教师和学生模型进行提示训练，以便进行进一步有效的训练。
如第 2.4 节所述，当教师的输出可缓存时（例如，就可用磁盘空间而言），可以通过加载前一轮生成的缓存文件来跳过教师的推理。即使教师的输出不可缓存，重新设计模型也能帮助用户缩短训练时间。请注意，学生模型的输出无法缓存，因为模型的参数每次迭代都会更新。下表表明，仅使用要执行的训练模块来重新设计模型是节省训练时间的有效方法，而且这种改进对于在大型数据集和/或大量轮次上训练模型更为关键。我们强调，用户可以通过在配置文件中指定所需模块来重新设计（最小化）模型，而不是硬编码（重新实现）剪枝后的模型。
- 通过重新设计教师和学生（ResNet-18）模型（仅使用图 3 所示的提示训练所需模块），提高了 Epoch 级别的训练速度。

Reference Methods

在这里，我们描述了知识提炼方法的重新实现和实验，以重现 ImageNet 和 COCO 数据集上报告的结果。

Reimplementations

鉴于 torchvision 中的预训练模型是在大型基准数据集 ImageNet (ILSVRC 2012) 和 COCO 2017 上训练的，我们将实现重点放在这些数据集上，因为预训练模型可以用作教师模型和/或基线学生模型（在人工注释的数据集上进行简单训练）。请注意，其中一些方法在其原始工作中未在这些数据集上进行验证。
下表简要总结了使用 torchdistill 重新实现的参考蒸馏方法，并指出了为重新实现这些方法而实现并添加到注册表中的附加模块。我们强调，在“必需的附加模块”列中没有任何复选标记 (X) 的方法（例如 KD、AT、PKT、RKD、HND、SPKD、Tf-KD、GHND 和 L2）只需将新的损失模块添加到框架中的注册表即可重新实现（第 2.2 节）。
- 参考torchdistill中实现的知识提炼方法。
与现有框架不同，上表中的所有方法都是独立于 torchvision 中的模型重新实现的，因此用户可以通过在配置文件中指定模型名称及其参数来轻松切换模型。以图像分类为例，模型的输入和（中间）输出的形状通常是固定的（例如，对于在 ImageNet 数据集上训练的模型，输入和（中间）输出的形状分别为 3 × 224 × 224 和 1,000），这使得在计算要最小化的损失值时很容易将学生的输出形状与老师的输出形状相匹配。

Reproducing ImageNet experiments

在本节中，我们尝试使用他们提出的蒸馏方法重现一些实验结果。具体来说，我们选择了注意力转移（AT）、因子转移（FT）、对比表征蒸馏（CRD）、无教师知识蒸馏（Tf-KD）、自监督知识蒸馏（SSKD）、L2 和素数感知自适应蒸馏（PAD-L2）方法，原因如下：
- 这些方法在其原始工作中使用 ResNet-34 和 ResNet18 的 ImageNet 数据集作为老师和学生模型进行了验证（Tf-KD 的教师模型是预训练的 ResNet-18）；
- ImageNet 实验中使用的超参数在原始研究和/或其已发布的源代码中进行了描述；并且
- 我们没有时间调整原始论文中未在 ImageNet 数据集上验证的其他方法的超参数。
除了这些方法之外，我们还将知识蒸馏 (KD) 应用于同一对教师-学生。请注意，除了 KD （对于KD，我们设置超参数如下：温度T = 1和相对权重α = 0:5）之外，我们重用了他们原始工作中给出的 ImageNet 超参数（例如，epoch 数）来重现他们的实验结果，并且我们提供了配置和日志文件以及经过训练的模型权重。
我们还应该注意到，Zagoruyko 和 Komodakis 提出了注意力转移（AT），并为他们的 ImageNet 实验定义了以下总损失函数：
- 其中 $L(W_S; x)$ 是标准交叉熵损失， $Q^j_S$ 和 $Q^j_T$ 分别表示第 j 对学生和教师注意力图的矢量化形式。在他们发布的框架【GitHub - szagoruyko/attention-transfer: Improving Convolutional Networks via Attention Transfer (ICLR 2017)】中，他们将 β 和 p 分别设置为 1,000 和 2。然而，我们发现他们定义的损失函数（等式 (3)）与他们实现的损失函数（等式 (4)）之间存在差异，后者计算教师和学生注意力图之间的均方误差 (MSE)。
在我们使用作者提供的超参数进行的初步实验中，学生模型无法很好地使用基于公式 (3) 的损失模块进行训练。因此，我们在实验中改用公式 (4) 进行 AT。
下表总结了实验结果，这些实验使用了原始研究中描述的和/或作者验证过的训练配置（例如，教师-学生对、超参数）。除了使用单个 GPU 进行的实验外，我们还使用 PyTorch 支持的分布式训练策略进行了实验（用匕首标记 † 表示），以证明我们的框架支持节省训练时间的策略。至于 L2 和 PAD-L2 方法，原始研究在其 ImageNet 实验中使用了 512 的批处理大小，这不适合我们的单个 GPU。因此，我们将批处理大小拆分为每个 GPU 171，并仅报告使用分布式训练的结果（用 ‡ 表示）。
- † 在 3 个 GPU 上进行分布式训练，采用线性缩放规则：根据分布式训练过程的数量修改学习率。（即乘以 GPU 数量）。 ‡ 在 3 个 GPU 上进行分布式训练，总批次大小与原始工作中使用的相同。
- 使用八种不同的蒸馏方法，在 ILSVRC 2012 数据集上用 ResNet-34（老师）训练 ResNet-18（学生）的验证准确率。使用原始作品中描述或作者提供的超参数（例如，# Epochs），所有重新实现的方法都优于没有老师训练的学生模型。
同样的策略也适用于 SSKD（普通样本和增强样本的总批次大小分别为 256 和 768 ），因为与其他方法相比，由于其 4 倍增强训练数据，在 epoch 级别训练模型所需的时间至少是其他方法的 4 倍，并且我们每个 GPU 的批次大小为 85（普通样本 + 增强样本 255）。同样，由于时间有限，我们对 CRD 也采用了同样的策略。我们还注意到，Zhang 等人将他们提出的 PAD-L2 应用于使用他们提出的 L2 作为预训练模型训练的学生模型，并使用 PAD-L2 方法对学生模型进行 30 个 epoch（即 120 个 epoch）的训练。
基于我们用 torchdistill 重新实现的方法，我们成功地在 ILSVRC 2012 数据集上重现了 AT 、Tf-KD 、L2 和 PAD-L2 方法的原始论文中报告的师生对结果，并且 PAD-L2 的结果最近被报告为 ILSVRC 2012 数据集上师生对的最佳性能。所有结果都优于仅使用人工标签训练的基线性能 (S: ResNet-18)，预训练模型由 torchvision 提供。请注意，FT 已在他们最初的工作中在 ILSVRC 2015 数据集上进行了验证，并且我们确认 FT 比使用 ILSVRC 2012 数据集的基线有所改进，因为 torchvision 中的教师模型 (ResNet-34) 是在该数据集上进行预训练的。重新实现的 CRD 的结果几乎与原始研究中报告的准确率相当。在 CRD 中，正样本和负样本都用于学习表示，因此成为表 5 中耗时最多的方法。
重新实现的 SSKD 表现优于基线模型，尽管准确率与报告的结果不匹配。一个潜在因素可能是由于我们有限的计算资源而强制采用不同的训练配置（例如，每个 GPU 的批处理大小不同，而他们的工作中使用了 8 个并行 GPU），因为我们只是重构并使作者发布的代码与 ILSVRC 2012 数据集兼容。正如 Tian 等人指出的那样，KD 仍然是一种强大的方法。我们重新实现的 KD 表现优于他们提出的最先进的方法 CRD（71.17%），并且实现了与他们的 CRD+KD（71.38%）方法相当的准确率。

Reproducing COCO experiments

为了证明我们的框架 1）可以应用于不同的任务，以及 2）与 torchvision 中未实现的模型架构协作，我们将广义头部网络蒸馏 (GHND) 应用于瓶颈注入的 R-CNN 对象检测器以进行拆分计算，使用 COCO 2017 数据集。他们提出的带有 ResNet-50 和 FPN 的瓶颈注入的 Faster 和 Mask R-CNN 旨在划分为头部和尾部模型，分别部署在移动设备和边缘服务器上，以降低资源受限的边缘计算系统中的推理速度。遵循 GHND 的原始工作，我们将该方法分别应用于一对原始的和瓶颈注入的 Faster R-CNN 作为老师和学生，并对 Mask R-CNN 进行相同的实验。如下表所示，重现的平均精度 (mAP) 与原始研究中报告的相匹配。
- GHND 使用原始 Faster/Mask R-CNN 模型（老师）在 COCO 2017 数据集上训练的瓶颈注入 R-CNN 模型（学生）的拆分计算验证 mAP。复制结果与原始文献 [Neural Compression and Filtering for Edge-assisted Real-time Object Detection in Challenged Networks] 中报告的结果一致。

Conclusions

在这项工作中，我们提出了 torchdistill，这是一个专门用于知识蒸馏研究的开源框架，它支持高效的训练和配置系统，旨在为用户提供实验摘要。研究人员可以在框架的基础上（例如，通过分叉存储库）进行知识蒸馏研究，并且可以通过发送拉取请求将他们的研究集成到框架中。这将有助于研究界确保工作的可重复性，并推进深度学习研究，同时支持在基准上进行公平的方法比较。具体来说，研究人员可以发布冠军表现的日志、配置和预训练模型权重，这将有助于他们确保特定数据集和师生对的冠军表现。
此外，torchdistill 生成的配置文件和日志文件将帮助研究人员完成 ML 代码完整性检查表【github.com】，我们在代码存储库中提供完整的配置（超参数）、日志文件和检查点，包括表 5 和表 6 中显示的实验结果的模型权重。我们为图像分类和对象检测任务提供参考代码和配置，并计划使用流行的软件包（例如用于 NLP 任务的 Transformers ）将我们的框架扩展到不同的任务。我们的框架将随着 PyTorch 和 torchvision 的新版本一起维护和更新，以便用户可以节省编码时间并将其用作可重复知识蒸馏研究的标准框架。

羞儿

关注

13
点赞
踩
8

收藏

觉得还不错? 一键收藏
打赏
0
评论
【读点论文】torchdistill: A Modular, Configuration-Driven Framework for Knowledge Distillation，文件配置完成知识蒸馏

我们框架中模块抽象的一个目标是让研究人员只需更改第 2.3 节中描述的 PyYAML 配置文件即可试验各种模块。我们将抽象重点放在要进行实验的关键模块上，特别是模型架构、数据集、转换和训练期间要最小化的损失。这些模块通常在作者发布的框架中是硬编码的（参见附录 A），许多超参数也是硬编码的。模型架构：torchvision 为视觉任务提供了各种模型系列，从 AlexNet 到 R-CNN ，其中许多模型都是在大型基准数据集上进行预训练的。
复制链接

扫一扫