发明内容
为了解决联邦学习在不均衡流量分类任务中通信代价过大和精度降低的问题,本发明提出如下技术方案:
一种基于不平衡数据分布的联邦学习网络流量分类模型训练方法,用于中心节点,所述训练方法包括:
中心节点在一个训练周期接收本地节点发送的本地模型权重和平衡度,其中,本地节点的流量数据在本地模型中训练得到本地模型权重,平衡度由所述流量数据所在的流量数据集计算所得;
将本地模型权重和平衡度加权聚合得到聚合后的全局模型;
中心节点将所述聚合后的全局模型向本地节点发送;
确定中心节点发送的聚合后的全局模型的准确率,根据所述准确率确定本地节点训练的模型;
至中心节点所加权聚合得到的全局模型收敛时所处于的训练周期结束训练,将结束时所处训练周期的所述全局模型的权重作为网络流量分类模型权重。
在一种实施例中,确定中心节点发送的聚合后的全局模型的准确率,根据所述准确率确定本地节点训练的模型的具体方法是
聚合后的全局模型的准确率不低于本地节点的本地模型的准确率,则中心节点发送的聚合后的全局模型由本地节接收,本地节点训练由中心节点发送的所述聚合后的全局模型,在下一个训练周期将训练后的所述全局模型作为新的本地模型,向中心节点发送新的本地模型权重和平衡度;聚合后的全局模型的准确率低于本地节点的本地模型的准确率,本地节点训练在所述一个训练周期发送的本地模型权重所对应的本地模型,在下一个训练周期将训练后的所述本地模型作为新的本地模型,向中心节点发送新的本地模型权重和平衡度。
在一种实施例中,所述将本地模型权重和平衡度加权聚合得到聚合后的全局模型的步骤包括
计算取得本地节点的流量数据集的平衡度;
根据所述平衡度,获取反映本地节点k数据量占所有节点数据量的比例,本地节点k平衡度占所有节点平衡度的比例的单向权重,
根据单项权重和本地节点的本地模型权重,计算聚合后的权重。
在一种实施例中,计算流量数据集的平衡度由公式(1)~(3)表示:
其中,p代表真实的流量数据分布,q代表在基准数据集中各类样本按照1:1:1的数据分布,DKL(p||q)表示真实流量数据与基准数据之间的KL距离,p(x)为真实条件下某类样本占所有样本的概率,q(x)为理想条件下某类样本占所有样本的概率,C为流量数据的类别数,n表示本地节点的数量,计算得到的平衡度A满足0≤A≤1的性质,A越趋近于1表示数据集越平衡,A越趋近于0表示数据集越不平衡。
在一种实施例中,中心节点计算单向权重,并根据单项权重和本地节点的权重计算聚合后的权重:
计算单项权重由公式(4)表示:
其中,n为所有本地节点总的数据量,n
k为单个本地节点的数据量,D
k为本地节点k数据量占所有节点数据量的比例,S
k为本地节点k平衡度占所有节点平衡度的比例,
为所有本地节点总的平衡度,A
k为本地节点k的平衡度;
计算聚合后的权重由公式(5)表示:
其中,
为加权聚合后的全局模型,
为本地节点的权重,α=0.7,layer l≠bn表示局部网络中的bn层不参与聚合。
在一种实施例中,所述的网络流量分类模型包含依次连接的四个阶段,分别为Conv1、Conv2_x、Conv3_x、Conv4_x、Conv5_x,其中Conv1部分输入为H×W,包含1个卷积核为7×7,步长为2的卷积层,输出尺寸为H×W;Conv2_x部分输入为H×W,包含2个残差块与一个池化层,其中每个残差块里有2个卷积层,三个卷积核大小依次为3×3、3×3,每个卷积层后为批归一化以及Relu激活函数,池化层为2×2大小的最大池化操作,输出尺寸为H/2×W/2;Conv3_x部分输入为H/2×W/2,包含依次连接的2个残差块与1个池化层,输出尺寸为H/4×W/4;Conv4_x部分输入为H/4×W/4,包含依次连接的2个残差块与1个池化层,输出尺寸为H/8×W/8;Conv5_x部分输入为H/8×W/8,包含依次连接的2个残差块与1个池化层,输出尺寸为H/16×W/16。
在一种实施例中,根据所述的网络流量分类模型分类结果和本地模型权重wi的方法是:对Conv5_x部分输出的特征图进行全局平均池化操作以及全连接操作,将输出特征图转化为一维向量,再对所述一维向量进行Softmax操作,得到分类结果与本地模型权重wi。
在一种实施例中,本地节点收集流量数据,并对所述收集到的流量数据进行预处理操作,预处理操作后再计算流量数据的平衡度及训练本地模型,所述预处理操作包括如下步骤:
S1.对20类网络流量数据拆分,将原始网络流量数据拆分为离散的网络流;
S2.对所述离散的网络流数据的包头信息进行剔除;
S3.数据清洗:对所述剔除包头的数据包里的重复内容或空白内容进行删除;
S4.对所述清洗后数据进行784字节的固定长度的截取,如果流量内容大于784字节,则将超过部分进行剔除,如果小于,则在末尾补充字节,其中,每个字节有8位构成,而标注工作是在每条网络流量数据的最后一个字节之后添加网络流量的类别标签;
S5.将所述截取后784维的pcap格式的流量数据转化为28×28维的灰度图像,最终映射为本地节点网络模型的输入。
在一个方面上,本发明还涉及一种由任一项所述的训练方法训练所得网络流量分类模型。
在一个方面上,本发明还涉及一种基于不平衡数据分布的联邦学习网络流量分类方法,由所述的网络流量分类模型用于对网络流量数据分类。
本发明与现有技术相比,具有如下优点和有益效果:本发明通过搭建基于不平衡数据分布的联邦学习网络流量分类算法,一方面可以协助多方共同学习一个准确且通用的神经网络模型,而无需公开和共享他们的本地用户数据集;另一方面通过在本地节点的模型接收策略上的改进,可以有效地减小通信代价;除此之外,中心节点上新的加权聚合方法可以有效地解决流量数据不平衡分布的问题。
具体实施方式
下面将结合附图和技术方案,对本发明的实施过程进行详细描述。
实施例1:本发明提供一种基于不平衡数据分布的联邦学习网络流量分类模型的训练方法,该方法可以一定程度上解决联邦学习在不均衡流量分类任务中通信代价过大和精度降低等问题,从而提升通信速度。该方法包括如下步骤:
S1、本地节点收集流量数据,并对所述收集到的流量数据进行预处理操作。
S2、计算所述预处理后的流量数据的平衡度。
S3、将述预处理后的流量数据放入分类模型中进行训练,得到本地模型权重。
S4、本地节点将所述本地模型权重与平衡度上传至中心节点进行加权聚合,得到聚合后的全局模型。
S5、本地节点下载所述聚合后的全局模型,并将所述聚合后的全局模型与聚合前的本地模型精度进行比较,判断是否接收全局模型。
S6、多次迭代步骤S3至S5,经过联邦学习中心节点端多次聚合更新至全局模型收敛后不再迭代,收敛得到的权重作为最终权重,所述最终权重用于进行网络流量分类。
在一种实施例中,所述S1具体包括:
本地节点收集流量数据,并对所述收集到的流量数据进行预处理操作。
所述本地节点从各类公开网站,平台收集网络流量,并将所述收集到的网络流量分为良性网络流量与恶性网络流量。其中所述良性网络流量包括FTP、BitTorrent、Outlook、Skype、SMB、Weibo、World of Warcraft、Facetime、Gmail、MySQL。所述恶性流量包括Cridex、Neris、Nsis-ay、Htbot、Geodo、Zeus、Miuref、Virut、Tinba,数据格式为pcap。对所述收集到的流量数据进行预处理操作生成idx格式的图像数据。
采用上述进一步方案的有益效果是:通过将pcap格式的流量变为图像数据,可以将网络流量分类问题转化为图片的分类问题,方便对其进行表征方面的分析,且给模型的进一步测试带来便利。
在一种实施例中,所述S2具体包括:
计算所述预处理后的流量数据的平衡度。
所述平衡度的计算方法为改进的KL距离,得到的平衡度与本地训练模型一同传入中心节点进行加权聚合。
采用上述进一步方案的有益效果是:量化每个本地节点数据集的平衡程度,并据此调节平衡策略,从而能够有效地提取本地模型的信息,提高流量分类精度。
在一种实施例中,所述S3具体包括:
将所述预处理后的流量数据放入分类模型中进行训练,得到本地模型权重。
所述分类模型为残差网络(ResNet18),所述残差网络包括1个卷积层、8个残差块、以及全连接层。所述残差块均为2层卷积层。
采用上述进一步方案的有益效果是:将ResNet18作为整个流量分类模型的主干网络,其作用为提取网络流量信息中的特征,供后续分类使用。相较于其他流量分类网络,Resnet18的模型参数更小,从而可以让本地节点与全局节点更快的进行通信,减少通信时间。
在一种实施例中,所述S4具体包括:
本地节点将所述本地模型权重与平衡度上传至中心节点进行加权聚合,得到聚合后的全局模型。
所述加权聚合是根据各客户端的流量数据的数据量与平衡度进行加权平均,能反映出本地节点占全部所有上传节点的比例,同时所述本地模型中所有BN层不参与到聚合中。能够提高异质性。
采用上述进一步方案的有益效果是:采用所述加权聚合的方法可以有效的解决流量数据的不平衡问题,从而提高流量分类精度,降低通信成本。
在一种实施例中,所述S5具体包括:
本地节点下载所述聚合后的全局模型,并将所述聚合后的全局模型与聚合前的本地模型精度进行比较,判断是否接收全局模型。
所述接收全局模型流程为,利用本地节点的数据集对聚合后的全局模型进行验证,当验证精度高于聚合前的精度,则接收全局模型。反之,则保留聚合前的本地模型。
采用上述进一步方案的有益效果是:有选择性的接收全局模型可以有效的提升模型收敛速度。
在一种实施例中,所述S6具体包括:
多次迭代S3至S5,经过联邦学习中心节点端多次聚合更新至全局模型收敛后不再迭代,最后中心节点聚合得到的权重作为最终权重,所述最终权重用于进行网络流量分类。
在一种更详细的实施例中,如图1所示,为本发明提供的一种基于联邦学习的网络流量分类方法,包括:
步骤1:本地节点收集流量数据,并对所述收集到的流量数据进行预处理操作。
所述本地节点从各类公开网站,平台收集网络流量,并将所述收集到的网络流量分为良性网络流量与恶性网络流量。其中所述良性网络流量包括FTP、BitTorrent、Outlook、Skype、SMB、Weibo、World of Warcraft、Facetime、Gmail、MySQL。所述恶性流量包括Cridex、Neris、Nsis-ay、Htbot、Geodo、Zeus、Miuref、Virut、Tinba,数据格式为pcap。每种类型网络流量50000条。
所述预处理方法具体分为5个步骤,如图2所示。第一,对20类网络流量数据拆分,将原始网络流量数据拆分为离散的网络流。第二,对所述离散的网络流数据的包头信息进行剔除。第三,数据清洗:对所述剔除包头的数据包里的重复内容或空白内容进行删除。第四,对所述清洗后数据进行784字节的固定长度的截取,如果流量内容大于784字节,则将超过部分进行剔除,如果小于,则在末尾补充字节。其中,每个字节有8位构成,而标注工作是在每条网络流量数据的最后一个字节之后添加网络流量的类别标签。第五,将所述截取后784维的pcap格式的流量数据转化为28×28维的灰度图像(如图3所示),最终映射为本地节点网络模型的输入。
基于不平衡数据分布的联邦学习网络流量分类算法具体如下:
在步骤2前需要统计参与联邦学习的本地节点的集合Ct,集合Ct是由max(Kρ,1)得到的。其中,K表示为本地节点的数量,ρ表示一个0到1之间的小数,。函数max(Kρ,1)表示取Kρ与1之间的最大值,表明服务器每次通信时至少随机选取1个客户端。
步骤2:计算所述预处理后的流量数据的平衡度。
所述流量数据的平衡度A采用改进的KL距离得到的,具体公式如下:
其中,p代表真实的流量数据分布,q代表在基准数据集中,各类样本按照1:1:1的数据分布。p(x)为真实条件下某类样本占所有样本的概率,q(x)为理想条件下某类样本占所有样本的概率,C为流量数据的类别数。计算得到的平衡度A满足0≤A≤1的性质。A越趋近于1表示数据集越平衡,A越趋近于0表示数据集越不平衡。
本发明通过改进的KL距离,将现有的平衡度A的区间由(-∝,2)改进为0≤A≤1,负无穷在权重计算中难以表达,因此导致平衡度在(-∝,2)区间表达时,本领域难以将平衡度用于权重计算,而由本发明实验数据可知,将平衡度使用在数据集中加权聚合,可以提高流量分离模型的分类准确率,因此改进的KL距离的计算实现了平衡度A的区间为0≤A≤1评价平衡度,使得平衡度A能够参与到模型的加权计算中,从而提高模型的分类准确性。
步骤3:将所述预处理后的流量数据放入分类模型中进行训练,得到局部模型权重。
所述分类模型为如图4所示的残差网络(ResNet18),并使用了在COCO数据集上的预训练模型。本发明采用的ResNet101,包含依次连接的四个阶段,分别为(Conv1、Conv2_x、Conv3_x、Conv4_x、Conv5_x),其中Conv1部分输入为H×W,包含1个卷积核为7×7,步长为2的卷积层,输出尺寸为H×W。Conv2_x部分输入为H×W,包含2个残差块与一个池化层,其中每个残差块里有2个卷积层,三个卷积核大小依次为3×3、3×3,每个卷积层后为批归一化以及Relu激活函数,池化层为2×2大小的最大池化操作,输出尺寸为H/2×W/2。Conv3_x部分输入为H/2×W/2,包含依次连接的2个残差块与1个池化层,输出尺寸为H/4×W/4。Conv4_x部分输入为H/4×W/4,包含依次连接的2个残差块与1个池化层,输出尺寸为H/8×W/8。Conv5_x部分输入为H/8×W/8,包含依次连接的2个残差块与1个池化层,输出尺寸为H/16×W/16。对Conv5_x部分输出的特征图进行全局平均池化操作,以及全连接操作,将输出特征图转化为一维向量,最后对所述一维向量进行Softmax操作,得到分类结果与本地模型权重wi。
步骤4:本地节点将所述本地模型权重与平衡度上传至中心节点进行加权聚合,得到聚合后的全局模型。
具体地,本地节点在初次训练需上传步骤2得到的平衡度,然后中心节点根据公式4计算得到单项权重
其中,n为所有本地节点总的数据量,nk为单个本地节点的数据量,Dk为本地节点k数据量占所有节点数据量的比例,Sk为本地节点k平衡度占所有节点平衡度的比例。
接着,联邦学习的中心节点采用加权聚合方法计算出聚合后的权重
加权聚合方法计算公式为:
其中,
为加权聚合后的全局模型,
为本地节点的权重,α=0.7,(layer l≠bn)为局部网络中的bn层不参与聚合。
步骤5:本地节点下载所述聚合后的全局模型,并将所述聚合后的全局模型与聚合前的本地模型精度进行比较,判断是否接收全局模型。
所述判断是否接收全局模型的方法为利用本地节点的数据集对聚合后的全局模型进行验证,当验证准确率高于聚合前的准确率,则接收全局模型,本地节点训练在聚合后的全局模型下进行下一轮训练。反之,则保留聚合前的本地模型,本地节点训练在聚合前的本地模型下进行下一轮训练。
步骤6:多次迭代步骤S3至S5,经过联邦学习中心节点端多次聚合更新至全局模型收敛后不再迭代,收敛得到的权重作为最终权重,所述最终权重用于进行网络流量分类。
如图5所示为基于加权聚合联邦学习的网络流量分类架构图,整个训练过程本地节点间的数据是不共享的,且无需将流量数据上传至中心节点。整个训练过程使用的计算资源为三张英伟达GTX 3090Ti显卡,深度学习训练框架为Pytorch,服务器模拟10个本地节点与1个中心节点。图6为本地节点与中心节点的任务流程图。
图7a、7b、7c分别为本发明算法在本地迭代轮数为10,全局迭代轮数为50下,数据平衡度A=1、0.75、0.5,对比FedAvg、FedBn、FedProx三种联邦学习策略的实验对比图,可以发现本发明在数据不平衡下,训练精度与收敛速度均优于另外三种策略。
上面对本发明的实施方法做出了详细说明,但本发明不限于上述实施方法。本发明通过对联邦学习聚合策略进行改进,提升了网络流量分类精度,减小通信代价。除此之外,中心节点上新的加权聚合方法有效地解决流量数据的不平衡分布。
本发明实施例还提供一种电子设备,该电子设备包括:存储器、处理器及存储在存储器上并可在处理器上运行的计算机程序,处理器执行计算机程序时实现上述实施例提供的方法中的步骤。本发明实施例提供的电子设备能够实现上述方法实施例中的各个实施方式,以及相应有益效果。
本发明实施例还提供一种计算机可读存储介质,计算机可读存储介质上存储有计算机程序,该计算机程序被处理器执行时实现本发明实施例提供的方法,且能达到相同的技术效果。
本领域普通技术人员可以理解实现上述实施例方法中的全部或部分流程,是可以通过计算机程序来指令相关的硬件来完成的,该程序可存储于一计算机可读取存储介质中,该程序在执行时,可包括如上述各方法的实施例的流程。其中,上述存储介质可为磁碟、光盘、只读存储记忆体(Read-Only Memory,ROM)或随机存取存储器(Random AccessMemory,简称RAM)等。
以上所揭露的仅为本发明较佳实施例而已,当然不能以此来限定本发明之权利范围,因此依本发明权利要求所作的等同变化,仍属本发明所涵盖的范围。