:no_entry: [已弃用] 知识蒸馏工具包

演示

用法

知识蒸馏训练

论文"缩小大脚怪:减少wav2vec 2.0的占用空间"的代码

这个工具包允许你使用知识蒸馏来压缩机器学习模型。要使用这个工具包,你需要提供一个教师模型、一个学生模型、用于训练和验证的数据加载器,以及一个推理流程。这个工具包基于PyTorch和PyTorch Lightning,因此教师和学生模型需要是PyTorch神经网络模块,数据加载器需要是PyTorch数据加载器。

<a name="demo"></a> 演示

我们提供了两个使用此工具包并压缩机器学习模型的演示。在这些演示中,我们展示了如何创建学生和教师模型、推理流程、训练和验证数据加载器,并将它们传递到知识蒸馏工具包中。

压缩resnet:

压缩wav2vec 2.0: 此笔记本

<a name="usage"></a>用法

定义推理流程

class InferencePipeline:
    def __init__(self):
        # 如有必要进行设置
        
    def run_inference_pipeline(self, model, data_loader):
        # 获取模型在验证数据集上的准确率
        return {"inference_result": accuracy}

具体示例请参见此处。

定义学生模型

import torch.nn as nn

class StudentModel(nn.Module):
    def forward(self, ):
        # 学生模型的前向传播

具体示例请参见此处。

定义教师模型

import torch.nn as nn

class TeacherModel(nn.Module):
    def forward(self, ):
        # 教师模型的前向传播

具体示例请参见此处。

将所有内容整合在一起并开始知识蒸馏

import torch
inference_pipeline = InferencePipeline()
student_model = StudentModel()
teacher_model = TeacherModel()
train_data_loader = torch.utils.data.DataLoader(train_dataset)
val_data_loaders = {"一个验证数据集": torch.utils.data.DataLoader(val_dataset)}

KD_example = KnowledgeDistillationTraining(train_data_loader = train_data_loader,
                                           val_data_loaders = val_data_loaders,
                                           inference_pipeline = inference_pipeline,
                                           student_model = student_model,
                                           teacher_model = teacher_model)
KD_example.start_kd_training()                                

<a name="kd_train"></a> 开始知识蒸馏训练!

要开始知识蒸馏训练,你需要先实例化KnowledgeDistillationTraining类,然后调用start_kd_training方法。

在下表中,我们展示了KnowledgeDistillationTraining类的构造函数接受的参数。

推理流水线如何工作？

该工具包使用推理流水线来测试学生模型。inference_pipeline类需要实现run_inference_pipeline方法。此方法的目的是获取学生模型在验证数据集上的性能。

我们通过以下代码示例来说明如何创建推理流水线。我们将model和data_loader传递给run_inference_pipeline。model是student_model，data_loader是验证数据加载器。在使用此工具包时，你应该已经准备好这两个参数，因为你需要它们来实例化KnowledgeDistillationTraining类。在run_inference_pipeline中，我们从data_loader中获取每个数据样本，然后将其传递给model。对于每个数据样本，我们根据学生模型的预测和真实标签计算准确率。最后，我们计算总体accuracy并将其作为字典返回。在返回的字典中，inference_result应该对应总体准确率。

class inference_pipeline:

    def __init__(self):
        # 构造方法是可选的

    def run_inference_pipeline(self, model, data_loader):
        accuracy = 0
        model.eval()
        with torch.no_grad():
            for i, data in enumerate(data_loader):
                X, y = data[0].to(self.device), data[1].to(self.device)
                outputs = model(X)
                predicted = torch.max(outputs["prob"], 1)[1]
                accuracy += predicted.eq(y.view_as(predicted)).sum().item()
        accuracy = accuracy / len(data_loader.dataset)
        return {"inference_result": accuracy}

上面的代码只是一个示例，你可以按照自己的方式创建推理流水线。只需记住两个规则：

1. inference_pipeline类只需要实现run_inference_pipeline。run_inference_pipeline在验证数据集上测试学生模型。

2. run_inference_pipeline应返回一个字典，例如{"inference_result": 衡量学生模型在验证数据集上性能的数值}。

损失函数如何工作？

在进行知识蒸馏训练时，我们需要计算损失函数。我们总是有一个知识蒸馏损失，它是教师模型和学生模型概率分布之间的KL散度。通常还会在损失函数中添加一个有监督训练损失。例如，如果学生模型是图像分类网络，有监督训练损失可以是交叉熵损失。最终的损失函数会结合不同的损失，这样我们就可以使用最终损失计算梯度并更新学生模型（显然，我们不会手动执行这个过程，PyTorch会自动为我们完成）。 这个工具包如何结合不同的损失（例如知识蒸馏损失、监督训练损失等）并形成最终损失？我们从final_loss_coeff_dict中获取系数，并将它们与final_loss_components中的损失相乘。final_loss_coeff_dict和final_loss_components都是字典，这两个字典中的键必须匹配，这样我们才知道如何将系数与损失对应起来。例如，final_loss_components总是包含"kd_loss"（知识蒸馏损失的键），所以final_loss_coeff_dict中必须有一个"kd_loss"键。使用这个工具包时，用户需要提供final_loss_coeff_dict，但不需要提供final_loss_components，因为我们内部会形成这个字典。如果你还想将监督训练损失（学生损失）添加到最终损失中，可以在学生模型的前向传播结束时返回它。具体示例可以参考这里。