[发明专利]基于知识蒸馏的模型训练方法、装置、设备及存储介质

说明书

本发明适用于人工智能技术领域，尤其涉及一种基于知识蒸馏的模型训练方法、装置、设备及存储介质，该方法通过获取满足目标条件的第一模型和不满足目标条件的第二模型，根据第一模型的输出构建的优化损失函数，得到更新后的第二模型，计算并确定目标相似度，构建相似度损失函数，并将相似度损失函数与优化损失函数的和作为目标损失函数，对第二模型进行训练，得到满足目标条件的第二模型。将第一模型与第二模型中网络层之间的表征相似度作为第二模型中损失函数的一部分，充分利用了网络中间层的信息，增加了第二模型学习第一模型的能力和范围，从而提高第二模型训练时的稳定性和收敛性。

技术领域

本发明涉及人工智能领域，尤其涉及一种基于知识蒸馏的模型训练方法、装置、设备及存储介质。

背景技术

目前，深度学习神经网络已成功应用于各种计算机视觉应用，如图像分类、对象检测和语义分割，大型的深度学习模型训练必须从非常大的、高度冗余的数据集中训练得到，但是数据集的数据量较大时模型训练需要占据大量的时间和存储空间，因此，为了缩短训练时间和减少资源占据，使用知识蒸馏方法对大型深度学习网络进行压缩得到了广泛运用，知识蒸馏方法对教师网络与学生网络的匹配度要求较高，而当前的知识蒸馏方法只会对训练样本集的标签进行优化，无法应对教师网络与学生网络的匹配程度不高的情况，导致的训练过程不稳定、不收敛的问题。因此，如何改进知识蒸馏的训练过程，以提高学生网络训练过程的稳定性、收敛性成为亟待解决的问题。

发明内容

基于此，有必要针对上述技术问题，提供一种基于知识蒸馏的模型训练方法、装置、设备及存储介质，以解决训练过程中不稳定、不收敛的问题。

本申请实施例的第一方面提供了一种基于知识蒸馏的模型训练方法，所述方法包括：

获取满足目标条件的第一模型和不满足目标条件的第二模型，第一模型包括M个网络层，第二模型包括N个网络层，N、M均为大于零的整数；

根据第一模型的输出构建的优化损失函数，更新第二模型的初始损失函数，得到更新后的第二模型；

计算第一模型中M个网络层分别与更新后的第二模型中N个网络层的表征相似度，通过预设选取条件，确定目标相似度；

根据目标相似度，构建相似度损失函数，并将相似度损失函数与优化损失函数的和作为目标损失函数；

使用训练集对第二模型进行训练，直至目标损失函数收敛，得到满足目标条件的第二模型。

本申请实施例的第二方面提供了一种基于知识蒸馏的模型训练装置，所述装置包括：

获取模型模块，用于获取满足目标条件的第一模型和不满足所述目标条件的第二模型，所述第一模型包括M个网络层，所述第二模型包括N个网络层，N、M均为大于零的整数；

更新模块，用于根据第一模型的输出构建的优化损失函数，更新第二模型的初始损失函数，得到更新后的第二模型；

目标相似度确定模块，用于计算第一模型中M个网络层分别与更新后的第二模型中N个网络层的表征相似度，通过预设选取条件，确定目标相似度；

目标损失函数确定模块，用于根据目标相似度，构建相似度损失函数，并将相似度损失函数与优化损失函数的和作为目标损失函数；

训练模块，用于使用训练集对第二模型进行训练，直至目标损失函数收敛，得到满足目标条件的第二模型。

第三方面，本发明实施例提供一种计算机设备，所述计算机设备包括处理器、存储器以及存储在所述存储器中并可在所述处理器上运行的计算机程序，所述处理器执行所述计算机程序时实现如第一方面所述的基于知识蒸馏的模型训练方法。

第四方面，本发明实施例提供一种计算机可读存储介质，所述计算机可读存储介质存储有计算机程序，所述计算机程序被处理器执行时实现如第一方面所述的基于知识蒸馏的模型训练方法。