[发明专利]基于知识蒸馏的模型训练方法及装置

说明书

本发明公开了一种基于知识蒸馏的模型训练方法及装置，涉及知识蒸馏技术领域，主要目的在于提高student模型的预测精度。本发明主要的技术方案为：利用第一数据集训练第一模型；基于数据增强技术对所述第一数据集进行扩展，得到第二数据集；利用所述第一模型、第一数据集与第二数据集训练第二模型，确定所述第二模型的损失函数。本发明用于训练高精度的student模型。

技术领域

本发明涉及知识蒸馏技术领域，尤其涉及一种基于知识蒸馏的模型训练方法及装置。

背景技术

知识蒸馏的概念首次提出于文章《Distilling the Knowledge in a NeuralNetwork》中，通过引入教师网络用以诱导学生网络的训练，实现知识迁移。因此，知识蒸馏是将一个网络的知识转移到另一个网络，两个网络可以是同构或者异构。做法是先训练一个teacher模型，然后使用这个teacher模型的输出和数据的真实标签去训练student模型。知识蒸馏，可以用来将网络从大网络转化成一个小网络，并保留接近于大网络的性能，以此解决模型在边缘段的部署硬件不足的问题。

但是，在知识蒸馏的过程中，由于student模型的结构更为简单，使得该student模型的精度只能是尽可能接近teacher模型的精度，这就使得在实际应用过程中，基于知识蒸馏得到的student模型的精度无法满足应用需求。

发明内容

鉴于上述问题，本发明提出了一种基于知识蒸馏的模型训练方法及装置，主要目的在于提高student模型的预测精度。

为达到上述目的，本发明主要提供如下技术方案：

第一方面，本发明提供一种基于知识蒸馏的模型训练方法，包括：

利用第一数据集训练第一模型；

基于数据增强技术对所述第一数据集进行扩展，得到第二数据集；

利用所述第一模型、第一数据集与第二数据集训练第二模型，确定所述第二模型的损失函数。

优选的，利用所述第一模型、第一数据集与第二数据集训练第二模型，确定所述第二模型的损失函数，包括：

利用第一模型的逻辑输出训练所述第二模型，得到第一损失函数；

利用第一数据集与第二数据集训练所述第二模型，得到第二损失函数；

根据所述第一损失函数与第二损失函数确定所述第二模型的损失函数。

优选的，利用第一模型的逻辑输出训练所述第二模型，得到第一损失函数，包括：

将同一样本数据分别输入所述第一模型与第二模型，得到第一逻辑输出与第二逻辑输出；

基于所述第一逻辑输出与第二逻辑输出的均方误差确定为所述第一损失函数。

优选的，利用第一数据集与第二数据集训练所述第二模型，得到第二损失函数，包括：

利用第二模型对所述第一数据集与第二数据集中的数据进行预测，确定所述数据的预测标注信息；

根据所述数据的预测标注信息与所述数据携带的标注信息的交叉熵，确定所述第二损失函数。

优选的，根据所述第一损失函数与第二损失函数确定所述第二模型的损失函数，包括：

获取所述第一损失函数与第二损失函数的占比系数；

根据所述占比系数对所述第一损失函数与第二损失函数进行加权求和，确定所述第二模型的损失函数。

优选的，在利用第一数据集训练第一模型之前，所述方法还包括：

构建第一模型，所述第一模型的模型结构中含有第二模型的模型结构。