[发明专利]基于联邦学习的模型训练方法、电子设备及存储介质

说明书

本申请提供一种基于联邦学习的模型训练方法、电子设备及存储介质。其中，应用于本地客户端的方法包括：响应于一轮训练开始，获取初始本地模型以及初始全局模型，其中，初始本地模型为上一轮训练中更新完成后的本地模型，初始全局模型为上一轮训练中中心服务器返回的全局模型；利用本地数据集、初始本地模型以及初始全局模型，并基于对比学习更新初始本地模型；将更新完成后的本地模型的本地模型参数发送至中心服务器；接收中心服务器返回的全局模型，一轮训练结束。本申请提供的方法利用对比学习判断数据相似度的思想，解决本地客户端数据集的异构问题，提高模型在训练完成后的性能。

技术领域

本申请涉及人工智能技术领域，尤其涉及一种基于联邦学习的模型训练方法、电子设备及存储介质。

背景技术

随着人工智能技术在各领域的广泛应用，越来越多的任务都从人工处理变成计算机辅助，更加高效快捷。但要得到准确的人工智能模型，需要大量的样本提前进行训练，从而得到更加准确的结果。但由于不同样本数据集之间存在数据保密以及隐私的问题，将大量数据集传输至模型训练服务器时面临着隐私泄露的风险，同时由于样本量巨大，也给传输过程带来巨大的负载，数据之间存在着壁垒，模型无法得到深入训练。

联邦学习的出现打破了“数据孤岛”的僵局。联邦学习通过传输模型或模型参数而非原始数据来保护原始数据中包含的保密信息以及隐私信息，本地客户端基于各自的数据集训练出本地模型，并将其上传到中心服务器进行聚合，聚合后得到的全局模型再分发到本地客户端，进行迭代更新训练，因此全局模型能够在保护数据私密性的同时利用所有数据集进行训练。

但在实际应用中，不同的数据集往往呈现异构性，即各数据集类别分布不均，数量不均，特征也不完全一致，所以训练得到的全局模型往往偏向于贡献较大的一方，导致模型不能在每个客户端都表现出很好的性能，造成全局模型的性能衰减。

发明内容

有鉴于此，本申请的目的在于提出一种基于联邦学习的模型训练方法、电子设备及存储介质。

基于上述目的，本申请提供了一种基于联邦学习的模型训练方法，应用于本地客户端，包括：响应于一轮训练开始，获取初始本地模型以及初始全局模型，其中，所述初始本地模型为上一轮训练中更新完成后的本地模型，所述初始全局模型为上一轮训练中中心服务器返回的全局模型；利用本地数据集、所述初始本地模型以及所述初始全局模型，并基于对比学习更新所述初始本地模型，得到更新完成后的本地模型；将更新完成后的所述本地模型的本地模型参数发送至中心服务器；接收所述中心服务器根据所述本地模型返回的全局模型，所述一轮训练结束。

可选地，所述利用本地数据集、所述初始本地模型以及所述初始全局模型，并基于对比学习训练所述初始本地模型，得到更新完成后的本地模型，进一步包括：执行迭代过程，其中，所述迭代过程包括：将所述本地数据集中的数据输入至第一模型、第二模型以及第三模型，得到多个第一特征、多个第二特征以及多个第三特征，其中，所述第一模型为所述初始本地模型，所述第二模型为所述初始全局模型，所述第三模型仅在每一轮训练的第一轮更新时为所述初始本地模型；利用所述第一特征、所述第二特征以及所述第三特征计算对比损失；利用所述对比损失更新所述第三模型的模型参数；响应于所述第三模型的模型参数达到第一更新目标，停止所述迭代过程，并将所述第三模型作为更新完成后的所述本地模型。

可选地，所述第一模型、所述第二模型以及所述第三模型均包括主干网络、特征融合层以及检测头部网络；所述利用所述第一特征、所述第二特征以及所述第三特征计算对比损失，进一步包括：利用下式计算所述对比损失：

其中，为主干网络对比损失，为特征融合层对比损失，为检测头部网络对比损失，依次为数据输入所述第一模型的主干网络、特征融合层以及检测头部网络所得到的特征，依次为数据输入所述第二模型的主干网络、特征融合层以及检测头部网络所得到的特征，依次为数据输入所述第三模型的主干网络、特征融合层以及检测头部网络所得到的特征，sim(·)为余弦相似度函数，τ为温度参数。