一种预训练模型抽选框架的构建方法及装置

1.本发明涉及对比自监督学习领域，尤其涉及一种预训练模型抽选框架的构建方法及装置。


背景技术：

2.模型自监督预训练是一个重要且具有挑战性的计算机视觉任务，可以通过在相应的提取框架中提取自监督预训练模型实现，在医学图像诊断，图像分割等有标注数据的领域具有广泛的应用。
3.在各个应用场景下，基于现有的提取框架进行模型自监督预训练，达到类似coco数据集的标注数量级，成本太高，难以实现，一般只能获取少量低成本的标注数据，数据受限导致模型训练十分困难，一般都需要选取预训练好的模型在下游数据集微调权重。与此同时，在各种不同的应用场景下，可使用的硬件资源差异十分明显，能够部署的模型也不相同，针对不同模型需要进行单独的预训练，模型的复用性很差，十分浪费硬件资源。


技术实现要素：

4.本发明提供一种种预训练模型抽选框架的构建方法及装置，用以解决现有技术中对现有对比学习自监督方法模型收敛慢，以及下游模型选择耗费过多资源的缺陷，可以提升模型的收敛速度、减少资源浪费和提升模型的复用性。
5.第一方面，本发明提供了一种预训练模型抽选框架的构建方法，包括：选取图像数据集和自监督对比学习框架；根据所述图像数据集和所述自监督对比学习框架对构建的超网预训练模型进行训练，得到训练好的超网预训练模型；选取下游迁移任务和下游迁移数据集；在基于所述自监督对比学习框架获取的采样空间中筛选符合预设条件的第一模型，基于所述下游迁移任务和所述下游迁移数据集计算所述第一模型与所述训练好的超网预训练模型的相似度；基于所述相似度的计算结果，确定与所述训练好的超网预训练模型共享权重的目标预训练模型，得到预训练模型抽选框架。
6.进一步地，根据所述图像数据集和所述自监督对比学习框架对构建的超网预训练模型进行训练，得到训练好的超网预训练模型，包括：将所述图像数据集输入所述自监督对比学习框架进行计算，得到构建的超网预训练模型的损失函数；基于所述损失函数对所述构建的超网预训练模型进行训练，得到训练好的超网预训练模型。
7.进一步地，所述将所述图像数据集输入所述自监督对比学习框架进行计算，得到构建的超网预训练模型的损失函数，包括：在每一训练轮次中，将所述图片数据集中的图片数据按预设数目的批次进行划分，并将每一批次的每个图片数据分别进行两次数据增强，得到所述每一批次的每张图片数据对应的两组数据增强的图片数据；基于所述自监督对比学习框架的特征提取骨干网络，设定所述采样空间，在所述采样空间中随机选取一个模型结构，基于选取的模型结构对所述监督对比学习框架中的梯度更新分支网络进行修改；将所述每一批次的每张图片数据对应的两组数据增强的图片数据，分别输入所述监督对比学
习框架中的修改的梯度更新分支网络和非梯度更新分支网络进行计算，得到所述损失函数。
8.进一步地，所述将每一批次的每个图片数据分别进行两次数据增强，得到所述每一批次的每张图片数据对应的两组数据增强的图片数据，包括：将所述每一批次的每个图片数据分别进行两次缩放、翻转、颜色转化和裁剪，得到所述每一批次的每张图片数据对应的两组数据增强的图片数据。
9.进一步地，所述基于所述下游迁移任务和所述下游迁移数据集计算所述第一模型与所述训练好的超网预训练模型的相似度，包括:基于所述下游迁移任务，筛选所述第一模型与所述训练好的超网预训练模型在所述下游迁移数据集中进行推理得到的第一模型中间特征相似图和训练好的超网预训练模型中间特征相似图；对所述第一模型中间特征相似图和所述训练好的超网预训练模型中间特征相似图进行相似度计算，得到所述第一模型与所述训练好的超网预训练模型的相似度。
10.进一步地，所述基于所述相似度的计算结果，确定与所述训练好的超网预训练模型共享权重的目标预训练模型，包括:将与所述训练好的超网预训练模型的相似度最大的第一模型作为所述目标预训练模型。
11.第二方面，本发明还提供例如一种预训练模型抽选框架的构建装置，包括：第一选取模块，用于选取图像数据集和自监督对比学习框架；训练模块，用于根据所述图像数据集和所述自监督对比学习框架对构建的超网预训练模型进行训练，得到训练好的超网预训练模型；第二选取模块，用于选取下游迁移任务和下游迁移数据集；计算模块，用于在基于所述自监督对比学习框架获取的采样空间中筛选符合预设条件的第一模型，基于所述下游迁移任务和所述下游迁移数据集计算所述第一模型与所述训练好的超网预训练模型的相似度；确定模块，用于基于所述相似度的计算结果，确定与所述训练好的超网预训练模型共享权重的目标预训练模型，得到预训练模型抽选框架。
12.第三方面，本发明实施例还提供了一种电子设备，包括存储器、处理器及存储在存储器上并可在处理器上运行的计算机程序，所述处理器执行所述程序时实现如第一方面所述的预训练模型抽选框架的构建方法的步骤。
13.第四方面，本发明实施例还提供了一种非暂态计算机可读存储介质，其上存储有计算机程序，该计算机程序被处理器执行时实现如第一方面所述的预训练模型抽选框架的构建方法的步骤。
14.第五方面，本发明实施例还提供了一种计算机程序产品，其上存储有可执行指令，该指令被处理器执行时使处理器实现如第一方面所述的预训练模型抽选框架的构建方法的步骤。
15.本发明提供的预训练模型抽选框架的构建方法及装置，通过选取图像数据集和自监督对比学习框架；根据图像数据集和自监督对比学习框架对构建的超网预训练模型进行训练，得到训练好的超网预训练模型；选取下游迁移任务和下游迁移数据集；在基于自监督对比学习框架获取的采样空间中筛选符合预设条件的第一模型，基于下游迁移任务和下游迁移数据集计算第一模型与所述训练好的超网预训练模型的相似度；基于相似度的计算结果，确定与训练好的超网预训练模型共享权重的目标预训练模型，得到预训练模型抽选框架。通过确定目标预训练模型得到的预训练模型抽选框架，可以实现高效的下游任务定制
化抽取，抽取出的模型具有极佳的泛化能力，能够更好的适应下游任务，无需任何下游有监督训练即可进行模型抽取。在实际应用场景中，下游任务有标签数据较多时，相较于常规方法进行模型选取时需要耗费大量gpu资源在下游任务上进行梯度更新训练，耗费过多资源，本方法仅需要模型在下游数据集上进行推理，无需任何gpu的硬件需求即可选出最优模型。
附图说明
16.为了更清楚地说明本发明或现有技术中的技术方案，下面将对实施例或现有技术描述中所需要使用的附图作简单地介绍，显而易见地，下面描述中的附图是本发明的一些实施例，对于本领域普通技术人员来讲，在不付出创造性劳动的前提下，还可以根据这些附图获得其他的附图。
17.图1为本发明提供的预训练模型抽选框架的构建方法实施例的流程示意图；
18.图2为本发明提供的获取训练好的超网预训练模型的方法实施例的流程示意图；
19.图3为本发明提供的获取损失函数的方法实施例的流程示意图；
20.图4为本发明提供的采样空间示意图；
21.图5为本发明提供的相似度计算方法实施例的流程示意图；
22.图6为本发明提供的构建预训练模型抽选框架的线管图；
23.图7为本发明提供的构建预训练模型抽选框架的流程图；
24.图8为本发明提供的预训练模型抽选框架的构建装置实施例的结构组成示意图；
25.图9示例了一种电子设备的实体结构示意图。
具体实施方式
26.为使本发明的目的、技术方案和优点更加清楚，下面将结合本发明中的附图，对本发明中的技术方案进行清楚、完整地描述，显然，所描述的实施例是本发明一部分实施例，而不是全部的实施例。基于本发明中的实施例，本领域普通技术人员在没有作出创造性劳动前提下所获得的所有其他实施例，都属于本发明保护的范围。
27.图1为本发明提供的预训练模型抽选框架的构建方法实施例的流程示意图。如图1所示，该预训练模型抽选框架的构建方法，可以包括以下步骤：
28.s101，选取图像数据集和自监督对比学习框架。
29.在步骤s101中，图像数据集可以为imagenet，或者也可以为coco，本发明实施例对此不作限定。自监督对比学习框架可以为moco框架、simclr框架、byol框架或者dino框架，本发明实施例对此不作限定。
30.s102，根据图像数据集和自监督对比学习框架对构建的超网预训练模型进行训练，得到训练好的超网预训练模型。
31.s103，选取下游迁移任务和下游迁移数据集。
32.在步骤s103中，迁移下游任务可以为图片分类，或者可以为图片检测，或者还可以为图片分割，本发明实施例对此不作限定。下游迁移数据集可以为pascal voc、coco、或pets等模型最终任务部署所用数据集，本发明实施例对此不作限定。
33.s104，在基于自监督对比学习框架获取的采样空间中筛选符合预设条件的第一模型，基于下游迁移任务和下游迁移数据集计算第一模型与训练好的超网预训练模型的相似
度。
34.在步骤s104中，预设条件可以为预先指定flops、参数量或者最大内存限制等，本发明实施例对此不作限定。
35.s105，基于相似度的计算结果，确定与训练好的超网预训练模型共享权重的目标预训练模型，得到预训练模型抽选框架。
36.在步骤s105中，可以将与训练好的超网预训练模型的相似度最大的第一模型作为目标预训练模型，得到目标预训练模型的过程即为构建预训练模型首选框架的过程。
37.本发明实施例提供的预训练模型抽选框架的构建方法，通过选取图像数据集和自监督对比学习框架；根据图像数据集和自监督对比学习框架对构建的超网预训练模型进行训练，得到训练好的超网预训练模型；选取下游迁移任务和下游迁移数据集；在基于自监督对比学习框架获取的采样空间中筛选符合预设条件的第一模型，基于下游迁移任务和下游迁移数据集计算第一模型与训练好的超网预训练模型的相似度；基于相似度的计算结果，确定与训练好的超网预训练模型共享权重的目标预训练模型，得到预训练模型抽选框架。通过确定目标预训练模型得到的预训练模型抽选框架，可以实现高效的下游任务定制化抽取，抽取出的模型具有极佳的泛化能力，能够更好的适应下游任务，无需任何下游有监督训练即可进行模型抽取。在实际应用场景中，下游任务有标签数据较多时，相较于常规方法进行模型选取时需要耗费大量gpu资源在下游任务上进行梯度更新训练，耗费过多资源，本方法仅需要模型在下游数据集上进行推理，无需任何gpu的硬件需求即可选出最优模型。
38.图2为本发明提供的获取训练好的超网预训练模型的方法实施例的流程示意图。如图2所示，该获取训练好的超网预训练模型的方法，可以包括以下步骤：
39.s201，将图像数据集输入自监督对比学习框架进行计算，得到构建的超网预训练模型的损失函数。
40.s202，基于损失函数对构建的超网预训练模型进行训练，得到训练好的超网预训练模型。
41.在步骤s201和步骤s202中，可以根据反向传播算法和sgd优化器来减小损失函数，以训练构建的超网预训练模型，经过多次迭代训练得到训练好的超网预训练模型。
42.本发明提供的获取训练好的超网预训练模型的方法，通过损失函数的获取和动态网络结构的设计，可在一次训练中完成一个可裁剪的动态超网，能够提取出满足各种运行需求的预训练子网。
43.图3为本发明提供的获取损失函数的方法实施例的流程示意图。
44.如图3所示，该获取损失函数的方法，可以包括以下步骤：
45.s301，在每一训练轮次中，将图片数据集中的图片数据按预设数目的批次进行划分，并将每一批次的每个图片数据分别进行两次数据增强，得到每一批次的每张图片数据对应的两组数据增强的图片数据。
46.在步骤s301中，预设数目可以为4次，或者可以为16次，或者还可以为48次，本发明实施例对此不作限定。假设图片数据集中共有256张图片数据，预设数目为4次，那么每个批次有64张图片数据，将每一个批次的64张图片数据的每一张图片数据分别进行两次数据增强，可以得到每个批次64张图片数据对应的两组图片数据，即每个批次具有两组图片数据，每组图片数据包括64张图片数据。
47.s302，基于自监督对比学习框架的特征提取骨干网络，设定采样空间，在采样空间中随机选取一个模型结构，基于选取的模型结构对监督对比学习框架中的梯度更新分支网络进行修改。
48.在步骤s302中，残差神经网络作为自监督对比学习框架的特征提取骨干网络。采用残差神经网络作为特征提取骨干网络时的采样空间可以如图4所示，采样空间规定了模型的深度和宽度的变化范围，从而确定整个可能的模型空间。
49.自监督对比学习框架包括梯度更新分支网络和非梯度更新分支网络，根据随机选取的模型对梯度更新分支网络进行修改，使修改的梯度更新分支网络与选取的模型具有相同的结构。
50.s303，将每一批次的每张图片数据对应的两组数据增强的图片数据，分别输入监督对比学习框架中的修改的梯度更新分支网络和非梯度更新分支网络进行计算，得到损失函数。
51.在步骤s303中，将每一批次的每张图片数据对应的两组数据增强的图片数据，分别输入监督对比学习框架中的修改的梯度更新分支网络和非梯度更新分支网络，可以计算得到每个分支的损失函数，并将这两个损失函数统计记为l。
52.本发明提供实施例提供的获取损失函数的方法，在训练中的每个迭代轮次，梯度更新分支的网络结构会在模型空间中进行随机采样，非梯度更新分支一直保持为梯度更新分支模型空间中最大网络结构，可以获得有利于使构建的超网预训练模型收敛的损失函数。
53.在一些可选的实施例中，将每一批次的每个图片数据分别进行两次数据增强，得到每一批次的每张图片数据对应的两组数据增强的图片数据，可以包括：将每一批次的每个图片数据分别进行两次缩放、翻转、颜色转化和裁剪，得到每一批次的每张图片数据对应的两组数据增强的图片数据。其中，对每个图片数据的缩放是随机多尺度的缩放。
54.图5为本发明提供的相似度计算方法实施例的流程示意图。如图5所示，该相似度计算方法，可以包括以下步骤：
55.s501，基于下游迁移任务，筛选第一模型与训练好的超网预训练模型在下游迁移数据集中进行推理得到的第一模型中间特征相似图和训练好的超网预训练模型中间特征相似图。
56.在步骤s501中，中间特征相似图可以为分类任务选取模型最终输出的特征图，检测选取backbone网络每个stage输出的特征图，分割选取backbone最后一个stage输出的特征图。
57.第一模型与训练好的超网预训练模型在下游迁移数据集中进行推理，得到相应的中间特征图集合，基于下游迁移任务，在得到的中间特征图集合中筛选第一模型中间特征相似图和训练好的超网预训练模型中间特征相似图。
58.s502，对第一模型中间特征相似图和训练好的超网预训练模型中间特征相似图进行相似度计算，得到第一模型与训练好的超网预训练模型的相似度。
59.在步骤s502中，通过计算第一模型中间特征相似图和训练好的超网预训练模型中间特征相似图进行相似度，确定第一模型与训练好的超网预训练模型的相似度。
60.图6为本发明提供的构建预训练模型抽选框架的线管图。如图6所示，构建预训练
模型抽选框架的步骤可以包括：对输入的图片进行两次数据增强，得到两组数据增强的图片，其中一组图片输入动态超网，另一组图片输入固定网络结构，得到两个损失函数，将两个损失函数统一记为损失函数，基于损失函数对构建的超网预训练模型进行训练，得到训练好的超网预训练模型；基于下游任务，如图片分类和图片分割，得到从采样空间筛选的模型的中间特征图片与训练好的超网预训练模型的中间特征图片，对二者进行相似度计算，可以得到目标预训练模型，完成预训练模型抽选框架的构建。
61.图7为本发明提供的构建预训练模型抽选框架的流程图。如图7所示，构建预训练模型抽选框架的步骤可以包括：选取图像数据集和自监督算法框架；确定构建的超网预训练模型的采样空间；对图片数据集进行预处理；在采样空间中随机采样一个模型，对自监督算法框架的梯度分支网络进行修改；将经过预处理的图片数据集输入自监督算法框架的修改的梯度分支网络和非梯度分支网络进行前向传播，计算得到损失函数；基于损失函数对构建的超网预训练模型进行训练；在采样空间内筛选符合限定条件的模型，确定下游迁移任务和下游迁移数据集；将训练好的超网预训练模型和筛选的模型在下游迁移数据集进行推理；根据下游迁移任务选取推理得到的训练好的超网预训练模型的中间特征图和筛选的模型中间特征图；基于二者进行相似度计算，筛选出最优模型，即目标预训练模型。
62.图8为本发明提供的预训练模型抽选框架的构建装置实施例的结构组成示意图。如图8所示，该预训练模型抽选框架的构建装置，包括：
63.第一选取模块801，用于选取图像数据集和自监督对比学习框架；
64.训练模块802，用于根据图像数据集和自监督对比学习框架对构建的超网预训练模型进行训练，得到训练好的超网预训练模型；
65.第二选取模块803，用于选取下游迁移任务和下游迁移数据集；
66.计算模块804，用于在基于自监督对比学习框架获取的采样空间中筛选符合预设条件的第一模型，基于下游迁移任务和下游迁移数据集计算第一模型与训练好的超网预训练模型的相似度；
67.确定模块805，用于基于相似度的计算结果，确定与训练好的超网预训练模型共享权重的目标预训练模型，得到预训练模型抽选框架。
68.可选地，训练模块802，包括：
69.第一计算单元，用于将图像数据集输入自监督对比学习框架进行计算，得到构建的超网预训练模型的损失函数；
70.训练单元，用于基于损失函数对构建的超网预训练模型进行训练，得到训练好的超网预训练模型。
71.可选地，第一计算单元，包括：
72.增强子单元，用于在每一训练轮次中，将图片数据集中的图片数据按预设数目的批次进行划分，并将每一批次的每个图片数据分别进行两次数据增强，得到每一批次的每张图片数据对应的两组数据增强的图片数据；
73.修改子单元，用于基于自监督对比学习框架的特征提取骨干网络，设定采样空间，在采样空间中随机选取一个模型结构，基于选取的模型结构对监督对比学习框架中的梯度更新分支网络进行修改；
74.计算子单元，用于将每一批次的每张图片数据对应的两组数据增强的图片数据，
分别输入监督对比学习框架中的修改的梯度更新分支网络和非梯度更新分支网络进行计算，得到损失函数。
75.可选地，增强子单元，还用于将每一批次的每个图片数据分别进行两次缩放、翻转、颜色转化和裁剪，得到每一批次的每张图片数据对应的两组数据增强的图片数据。
76.可选地，计算模块804，包括:
77.筛选单元，用于基于下游迁移任务，筛选第一模型与训练好的超网预训练模型在下游迁移数据集中进行推理得到的第一模型中间特征相似图和训练好的超网预训练模型中间特征相似图；
78.第二计算单元，用于对第一模型中间特征相似图和训练好的超网预训练模型中间特征相似图进行相似度计算，得到第一模型与训练好的超网预训练模型的相似度。
79.可选地，确定模块805，包括:
80.确定单元，用于将与训练好的超网预训练模型的相似度最大的第一模型作为目标预训练模型。
81.图9示例了一种电子设备的实体结构示意图，如图9示，该电子设备可以包括：处理器(processor)901、通信接口(communicationsinterface)902、存储器(memory)903和通信总线904，其中，处理器901，通信接口902，存储器903通过通信总线904完成相互间的通信。处理器901可以调用存储器903中的逻辑指令，以执行如下预训练模型抽选框架的构建方法：
82.选取图像数据集和自监督对比学习框架；根据图像数据集和自监督对比学习框架对构建的超网预训练模型进行训练，得到训练好的超网预训练模型；选取下游迁移任务和下游迁移数据集；在基于自监督对比学习框架获取的采样空间中筛选符合预设条件的第一模型，基于下游迁移任务和下游迁移数据集计算第一模型与训练好的超网预训练模型的相似度；基于相似度的计算结果，确定与训练好的超网预训练模型共享权重的目标预训练模型，得到预训练模型抽选框架。
83.此外，上述的存储器903中的逻辑指令可以通过软件功能单元的形式实现并作为独立的产品销售或使用时，可以存储在一个计算机可读取存储介质中。基于这样的理解，本发明的技术方案本质上或者说对现有技术做出贡献的部分或者该技术方案的部分可以以软件产品的形式体现出来，该计算机软件产品存储在一个存储介质中，包括若干指令用以使得一台计算机设备(可以是个人计算机，服务器，或者网络设备等)执行本发明各个实施例方法的全部或部分步骤。而前述的存储介质包括：u盘、移动硬盘、只读存储器(rom，read-onlymemory)、随机存取存储器(ram，randomaccessmemory)、磁碟或者光盘等各种可以存储程序代码的介质。
84.另一方面，本发明实施例还提供一种计算机程序产品，所述计算机程序产品包括存储在非暂态计算机可读存储介质上的计算机程序，所述计算机程序包括程序指令，当所述程序指令被计算机执行时，计算机能够执行上述各实施例提供的预训练模型抽选框架的构建方法：
85.选取图像数据集和自监督对比学习框架；根据图像数据集和自监督对比学习框架对构建的超网预训练模型进行训练，得到训练好的超网预训练模型；选取下游迁移任务和下游迁移数据集；在基于自监督对比学习框架获取的采样空间中筛选符合预设条件的第一
模型，基于下游迁移任务和下游迁移数据集计算第一模型与训练好的超网预训练模型的相似度；基于相似度的计算结果，确定与训练好的超网预训练模型共享权重的目标预训练模型，得到预训练模型抽选框架。
86.又一方面，本发明还提供一种非暂态计算机可读存储介质，其上存储有计算机程序，该计算机程序被处理器执行时实现以执行上述各实施例提供的预训练模型抽选框架的构建方法：
87.选取图像数据集和自监督对比学习框架；根据图像数据集和自监督对比学习框架对构建的超网预训练模型进行训练，得到训练好的超网预训练模型；选取下游迁移任务和下游迁移数据集；在基于自监督对比学习框架获取的采样空间中筛选符合预设条件的第一模型，基于下游迁移任务和下游迁移数据集计算第一模型与训练好的超网预训练模型的相似度；基于相似度的计算结果，确定与训练好的超网预训练模型共享权重的目标预训练模型，得到预训练模型抽选框架。
88.以上所描述的装置实施例仅仅是示意性的，其中作为分离部件说明的模块可以是或者也可以不是分开的，作为模块显示的部件可以是或者也可以不是第二模块，即可以位于一个地方，或者也可以分布到多个网络模块上。可以根据实际的需要选择其中的部分或者全部模块来实现本实施例方案的目的。本领域普通技术人员在不付出创造性的劳动的情况下，即可以理解并实施。
89.通过以上的实施方式的描述，本领域的技术人员可以清楚地了解到各实施方式可借助软件加必需的通用硬件平台的方式来实现，当然也可以通过硬件。基于这样的理解，上述技术方案本质上或者说对现有技术做出贡献的部分可以以软件产品的形式体现出来，该计算机软件产品可以存储在计算机可读存储介质中，如rom/ram、磁碟、光盘等，包括若干指令用以使得一台计算机设备(可以是个人计算机，服务器，或者网络设备等)执行各个实施例或者实施例的某些部分的方法。
90.最后应说明的是：以上实施例仅用以说明本发明的技术方案，而非对其限制；尽管参照前述实施例对本发明进行了详细的说明，本领域的普通技术人员应当理解：其依然可以对前述各实施例所记载的技术方案进行修改，或者对其中部分技术特征进行等同替换；而这些修改或者替换，并不使相应技术方案的本质脱离本发明各实施例技术方案的精神和范围。