知识蒸馏（尝试在ASR方向下WeNet中实现）

左手的ㄟ右手 2024-04-06 14:05 64阅读 0赞

#### 知识蒸馏（尝试在WeNet中实现知识蒸馏） ####

*  1、知识蒸馏简介
 *   *  1.1、论文
     *  1.2、目标蒸馏结构解释
 *  2、WeNet中关于知识蒸馏的思考
 *   *  2.1、WeNet结构
     *  2.2、CTC端的蒸馏
 *  3、WeNet知识蒸馏实验
 *  4、总结

## 1、知识蒸馏简介 ##

蒸馏的作用：模型压缩、加速训练  
蒸馏的知识：大模型logits层、激活、神经元、中间层特征、教师网络参数等  
蒸馏的结构：大的教师模型、小的学生模型（大小什么意思？可以理解为我们的层数少一些，网络结构少一些）

### 1.1、论文 ###

论文：《Distilling the Knowledge in a Neural Network》  
详细信息可自己查看该论文。论文中有提及ASR，因此本文主要是实现ASR中的知识蒸馏。

论文综述：《Knowledge Distillation A Survey》  
可查看知识蒸馏的基础知识与常用蒸馏方式。

### 1.2、目标蒸馏结构解释 ###

本文主要进行logits蒸馏（还可以有基于特征、基于关系的蒸馏方式）。其实主要的就是计算蒸馏中的两个loss，soft loss和hard loss。

*  Soft Loss  
    softoutput1（softtarget）：教师模型，经过教师模型后的输出除以温度T后再经过softmax层后得到softoutput1。  
    softoutput2：学生模型，经过学生模型后的输出除以温度T后再经过softmax层后得到softoutput2。  
    softoutput1与softoutput2求解Soft Loss（这里softloss我使用了KL相对熵loss）
 *  Hard Loss  
    hardoutput：学生模型，经过学生模型后的输出再经过softmax层后得到softoutput。  
    hardtarget：真实的实验数据  
    hardoutput与hardtarget求解Hard Loss（这里就是正常不进行蒸馏的训练）

![在这里插入图片描述][aa95232c03564d729f04c0e9c2ef27bf.png]

## 2、WeNet中关于知识蒸馏的思考 ##

### 2.1、WeNet结构 ###

如图可以看到，wenet包括一个贡共享的编码器和两个解码器结构，左解码器为CTC，右解码器为Attention。因此我们在进行蒸馏的时候应该考虑到他的两个输出，即CTC端的输出，以及Attention端的输出。  
![在这里插入图片描述][57c8774545ec4991ac6900288d535b9c.png]

### 2.2、CTC端的蒸馏 ###

在Wenet中进行蒸馏，这里考虑对Attention端直接进行蒸馏，即对输出直接进行loss求解。对CTC端分为不同级别的蒸馏，具体如下：

*  frame  
    正常训练中我们对数据的处理即处理到帧级别，因此wenet中进行frame级别的蒸馏也就是对encoder（教师模型和学生模型）的输出除以温度T，求解Soft Loss,这里的Soft Loss使用pytorch自带的nn.KLDivLoss.
 *  sequence\_max  
    上面提到了帧级别，这里打算使用序列级别的蒸馏。什么是序列级别的蒸馏？对什么进行序列级别的蒸馏？
    
    a、总所周知，wenet中有一端解码器为CTC解码，也就是在这一端我们使用pytorch自带的CTCloss计算损失，CTCloss的输入有四个（输入语音数据、输入语音数据长度、输入序列的标签、输入序列标签长度（speech, speech\_lengths, text, text\_lengths））。
    
    b、ctc有一类解码方式为greedy\_search解码，即找到topk，然后输出对应的标签表示（text’, text\_lengths’）。标签表示？那正好跟我们的CTCloss的输入一致，所以我们考虑让CTC端蒸馏的时候学习教师模型greedy\_search解码后的结果，即求解的学生loss为torch.nn.CTCLoss(speech, speech\_lenghts, text’, text\_lengths’)，这样就是学习的序列级的知识。

## 3、WeNet知识蒸馏实验 ##

实验：

教师模型：已经训练完成的大模型。

学生模型：教师模型的一半（什么意思呢？就是教师模型的编码器有12层，解码器有6层，做卷积的时候卷积核为31，那我学生模型的编码器就设置6层，解码器有3层，做卷积的时候卷积核为15）

温度：有试过1、1.2、2，发现温度为1时控制其他变量后效果更好，因此温度T=1，其他配置一致。

## 4、总结 ##

知识蒸馏在wenet中的优缺点（基于3中的实验而言）

**优点**：

*  模型大小：从513M压缩到85M（相较于量化而言压缩得更小）。
 *  语音wer：测试过包括aishell、magicdata、aidatatang、wenetspeech等测试集，wer基本一致，且在部分解码方式下知识蒸馏后的模型效果更好（相较于量化而言wer也更小）（att端直接目标蒸馏，CTC端frame蒸馏时，attention与attention\_rescoring解码效果优于不做蒸馏的小model，但ctc类的解码方法学生model性能劣于不做蒸馏的小model；att端直接目标蒸馏，CTC做序列蒸馏，即学生model学习教师模型greedy\_search解码后的结果，此时用到的是pytorch自带CTCloss，attention、attention\_rescoring、CTC类解码均优于不做蒸馏的小model）。
 *  训练速度：训练1wh数据，4张3090卡，batchsize为18，一个epoch能够从9小时变为8.5小时左右。

**缺点：**

*  实现过程：相比于wenet自带的量化功能而言，实现知识蒸馏比较麻烦。且对语音识别过程进行知识蒸馏而言，现在研究感觉还比较少，处于试错的阶段（比如对温度T的把握，就要多做尝试）
 *  蒸馏方法：直接目标蒸馏效果可能没有那么好，可以考虑优化蒸馏方法，比如优化CTC部分的蒸馏，或者考虑对不同层的特征知识进行蒸馏。

[aa95232c03564d729f04c0e9c2ef27bf.png]: https://image.dandelioncloud.cn/pgy_files/images/2024/04/06/0d90cb37c3e444f799007868d2ba23bf.png
[57c8774545ec4991ac6900288d535b9c.png]: https://image.dandelioncloud.cn/pgy_files/images/2024/04/06/87b1ad96eebd4f2fa035a604f7bbc53a.png