MLaaS4HEP:面向高能物理的机器学习即服务
摘要
机器学习(ML)将在欧洲核子研究中心(CERN)即将开展的高亮度大型强子对撞机(HL‐LHC)项目中发挥重要作用。未来十年,LHC实验将收集前所未有的艾级规模数据,这一工作需要全新的方法来训练和使用机器学习模型。本文讨论了一种面向高能物理领域的机器学习即服务流水线(MLaaS4HEP),该流水线提供三个独立的层:用于以原生ROOT数据格式读取高能物理(HEP)数据的数据流层;利用分布式ROOT文件训练机器学习模型的数据训练层;以及通过HTTP协议为预训练机器学习模型提供预测服务的数据推理层。这种模块化设计使得能够从远程存储设施(例如世界范围LHC计算网格(WLCG)基础设施)读取ROOT文件,在大规模上进行数据训练,并将数据输入用户偏好的机器学习框架。作为服务的TensorFlow(TFaaS)实现的推理层,可为现有基础设施以及高能物理领域内外的应用程序提供对预训练机器学习模型的便捷访问。特别是,我们展示了MLaaS4HEP架构在一个物理案例中的应用,即最初使用定制N元组在CMS中进行的t̄t希格斯分析。我们详细介绍了使用分布式ROOT文件训练机器学习模型的过程,讨论了MLaaS和TFaaS方法在选定物理分析中的性能,并将结果与传统方法进行了比较。
关键词 大数据 · LHC · 数据管理 · 机器学习
引言
随着欧洲核子研究中心LHC计划的推进,高能物理领域的数据增长开始加速。到第二轮运行结束时,欧洲核子研究中心的实验已经达到了拍字节(PB)级别,每年产生O(100) PB的数据。新的HL‐LHC计划将进一步扩展至艾字节级,而机器学习在高能物理中的应用将至关重要[1]。机器学习技术已成功应用于在线和离线重建程序,在探测器模拟、对象重建、识别、蒙特卡洛(MC)生成等方面应用也带来了巨大收益[2]。正如《高能物理中的机器学习社区白皮书》中所指出的,计算机科学专家未能充分参与解决高能物理机器学习挑战的部分原因在于,高能物理数据存储在ROOT数据格式中,这种格式在高能物理领域之外鲜为人知。ROOT数据格式[3]用于在基于树的数据结构中存储高能物理事件,其中单个事件的大小无法预先确定,例如每个物理事件中的电子数量各不相同。相反,现有的机器学习框架依赖于固定大小的单个事件数据表示,通常存储在CSV[4],NumPy[5], HDF5[6]等数据格式中。因此,基于事件的数据结构无法直接输入到现有的机器学习框架中,需要在框架层面或本文讨论的数据输入层面采取特殊处理。这一点以及其他原因导致了机器学习与高能物理两个社区之间的人为隔阂。例如,在最近的Kaggle挑战赛中[7——9]高能物理数据被转换为CSV数据格式,以便非高能物理领域的机器学习从业者能够参与竞赛。
本文讨论了高能物理的机器学习即服务(MLaaS)架构,文中称之为MLaaS4HEP,该架构包含两个独立部分。第一部分,即MLaaS4HEP框架[10],提供了一种方式将基于ROOT的高能物理数据原生读入用户选择的Python机器学习框架中。第二部分,即TensorFlow as a Service(TFaaS)框架[11],可用于托管预训练机器学习模型,并通过HTTP协议获取预测结果。这种方法可供高能物理领域内外的物理学家或其他专家使用,因为它仅依赖于Python库。它能够访问本地或远程数据存储,且无需对实验专用框架进行任何修改或集成。这种模块化设计使得可以通过WLCG站点远程访问PB级数据集来训练机器学习模型,而无需进行数据转换,即无需将ROOT数据格式转换为扁平数据格式,也无需后续存储这些数据集以供底层机器学习框架使用。因此,通过所讨论的机器学习即服务架构,可以有效缩小高能物理与机器学习社区之间现有的差距。
本文的组织结构如下。第2节概述了相关工作以及所提出的解决方案的关键方面。第3节介绍了MLaaS4HEP架构及其工作流程的详细信息。第4节展示了MLaaS4HEP在物理应用案例中的性能结果和验证。第5节概述了可能的未来发展方向,第6节给出了总结。
相关工作与解决方案
机器学习即服务是工业界一个众所周知的概念,主要的IT公司都为其客户提供此类解决方案。例如,Amazon ML、Microsoft Azure ML Studio、Google预测API和机器学习引擎,以及IBM Watson都是这一概念的典型实现(参见[12])。通常,机器学习即服务涵盖各种机器学习任务,例如通过REST API进行的数据预处理、模型训练与评估以及推理。尽管提供商提供了大量接口和API,但这些服务大多数情况下旨在覆盖标准使用案例,例如自然语言处理、图像分类、计算机视觉和语音识别。虽然可以向这些平台提供自定义机器学习代码库,但由于多种原因,其在高能物理领域的应用仍然非常有限。例如,高能物理ROOT数据格式无法直接用于任何服务提供商的API中。因此,对于大规模数据集而言,运营成本(例如将ROOT文件转换为机器学习即服务提供商API所使用的数据格式、数据管理以及数据预处理)可能非常高昂。从动态大小的基于事件的树形格式到固定大小的数据表示的数据展平功能也不存在。因此,我们发现现成的商业解决方案在高能物理使用案例中大多不可行或效率低下(在成本和功能方面)。这种情况在未来可能会改变,随着各类倡议的发展,例如欧洲核子研究中心OpenLab[13],继续与几乎所有上述服务提供商紧密合作。
同时,高能物理领域的各种研发活动正在进行中。例如,hls4ml项目[14]致力于在FPGA上实现机器学习推理,而SonicCMS项目[15]旨在提供协处理器上的最优网络推理服务。这两个项目都专注于优化推理阶段,而非完整的机器学习流水线,即从读取数据、训练模型到提供预测的整个流程。
另一种解决方案使用Spark平台进行数据处理和机器学习训练[16]。尽管该方案看起来非常有前景,但它需要将数据摄入到CERN EOS文件系统或HDFS/Spark基础设施中。因此,无法便捷地访问位于WLCG站点或此类专用基础设施之外的数据。此外,可能需要在Keras API之上使用基于Spark的库(Analytics Zoo、BigDL),而用户在机器学习框架选择上的灵活性受到限制。
如前一节所述,通常需要在数据或软件框架层面进行额外的转换。例如,在CMS中,用于喷注标识算法的深度神经网络(DNN)依赖于TensorFlow(TF)队列系统,并使用自定义操作内核读取ROOT树,然后将其输入到TensorFlow等机器学习模型中[17]。尽管这在特定用例中是一种有趣的方案,但它并非一种“即服务”解决方案。
最后,我们发现目前尚无现成的最终产品能够无需额外工作即可实现针对分布式高能物理数据的机器学习即服务,并且能透明地集成现有的基于Python的机器学习框架以对高能物理数据进行机器学习训练,而本工作旨在填补这一空白。
所提出方案的新颖性
所提出的解决方案的原创性贡献如下。
-
我们提供对存储在基于事件树的ROOT数据格式中的高能物理数据集的透明访问,以供用户选择的现有基于Python的机器学习框架使用。这些框架通常设计用于处理基于行的数据结构,例如NumPy数组、CSV文件等。第3.1节和3.2节讨论的所提出的解决方案依赖于uproot库[18]和XrootD协议[19],用于从本地文件系统或远程站点读取小型或大型基于树的ROOT文件。它将ROOT数据的锯齿状数组1表示进行转换,并通过将输入/输出流进行基于向量或矩阵的转换,将其输入到机器学习框架中。这为使用非高能物理领域的常用机器学习框架(如PyTorch[20], Keras[21], TensorFlow[22], fast.ai[23],等)提供了可能,可以利用分布式数据集训练机器学习模型,从而吸引非高能物理领域的机器学习从业者参与高能物理的机器学习活动。
-
在第4.2节中,我们展示了所提出的解决方案可以作为基于自定义扁平元组(源自生产级ROOT文件)的传统高能物理分析的有效替代方案。接着,在第4.3节中,我们针对特定分析提供了性能研究,并在第4.4节中讨论了超大规模数据集的性能预测2。
-
我们开发了一个独立的张量即服务框架[11],作为此项工作的一部分,该框架通过HTTP协议提供对各种基于张量的机器学习模型的访问。尽管类似功能在各种工业解决方案中已经存在,但它们大多集成在其服务栈中,对于科研社区而言可能难以负担或无法获取,而科研社区更希望拥有一种高效、可扩展的开源替代方案。
-
最后,所提出的模块化架构可以轻松地在任何高能物理实验中进行调整,既可以作为完整的流水线使用,也可部分使用,且无需对现有框架或基础设施进行修改。
MLaaS4HEP 架构
典型的机器学习工作流包括三个步骤:获取数据训练所需的数据,使用机器学习框架训练模型,以及利用训练好的模型进行预测。在我们的机器学习即服务解决方案MLaaS4HEP[10],中,该工作流可分别抽象为数据流、数据训练和推理阶段。这些组件中的每一个都可以紧密集成到应用程序设计中,也可以单独组合和使用。具体选择主要由特定的使用案例决定。我们可以定义这些层次如下(见图1)。
- 数据流层 :负责读取本地和/或远程ROOT文件,并将数据批次以数据流形式传输到数据训练层。该层的实现需要ROOT输入输出层支持远程输入输出文件访问;
- 数据训练层 :它是对TensorFlow等标准机器学习库的轻量封装,PyTorch等。它从数据流层以数据块的形式接收数据,将其从基于ROOT TTree的表示转换为适用于底层机器学习框架的格式,并用于训练用途;
- 数据推理层 :它指的是服务架构中的推理部分。我们将其实现为TFaaS[11],允许上传预训练模型,并通过HTTP协议从客户端应用查询这些模型。
尽管这些层的实现可能因实验而异(或其他科学领域),但它们可以轻松地被泛化,并成为通用机器学习即服务平台的基础。MLaaS4HEP框架[10]实现了数据流和数据训练层,我们分别在第3.1节和3.2节中提供其详细信息。在第3.3节中,我们提供了在MLaaS4HEP框架中实现并用于第4节所展示研究的机器学习训练工作流的技术细节。数据推理层作为独立的TFaaS[11]框架实现,因为它可以在高能物理之外使用,其详细信息在第3.3.4节中讨论。
数据流层
数据流层负责从本地或远程数据存储中流式传输数据。最初,ROOT文件的读取主要依赖于C++或PyRoot框架,但近年来ROOT I/O的发展显著简化并加速了从Python访问ROOT数据的过程。主要开发工作由DIANA‐HEP计划支持的uproot[18]框架完成[25]。uproot库使用NumPy[5]调用,将ROOT文件中的数据块快速转换为NumPy数组。它支持多种功能,包括对ROOT TBranch的局部读取、非扁平TTrees、非TTree直方图等。该库依靠数据缓存和并行处理来实现高吞吐量。数据可以从本地ROOT文件读取,也可以通过XrootD协议远程读取[19]。
在我们实现的机器学习即服务中(见第3.5节),该层由一个Python生成器构成[26],能够从本地或远程文件中读取数据块。该Python生成器的输出是一个包含扁平和锯齿状数组属性的NumPy数组。这种实现在访问大型数据集时具有高效性,因为它不需要将整个数据集加载到训练节点的内存中。此外,它还可用于并行化机器学习工作流管道中的数据流。块大小的选择应根据所处理事件的复杂性、可用网络带宽和硬件资源来决定,详见第4.3节的讨论。
数据训练层
该层将数据流层提供的高能物理ROOT数据(以锯齿数组形式呈现)转换为应用程序所使用的扁平数据格式[1, 17]。锯齿数组(见图2)是高能物理实验中产生的变大小事件数据的紧凑表示形式。
高能物理基于树的数据表示针对数据存储进行了优化,但并不直接适用于机器学习框架。因此,需要进行一定的数据转换,将基于树的数据结构转化为扁平数据结构,以供机器学习框架使用。我们探索了两种可能的转换方式:带有填充值的向量表示(见图3)以及用户选择的相空间内数据的矩阵表示(见图5)。
高能物理事件在不同事件属性上具有不同的维度。例如,单个事件可能包含不同数量的粒子。因此,在将ROOT事件转换为锯齿数组表示时,应适当地进行展平和填充。为此,我们采用两遍处理流程。在第一遍遍历所有事件时,我们确定每个属性的维度及其最小/最大值。尽管该流程在超大规模数据集(即太字节或拍字节规模)上可能不可行,但它可以被可轻松被采用近似最小/最大值和裁剪过程的替代方法所替换。在第二次处理过程中,我们将锯齿数组属性映射为具有适当大小和填充的单个向量表示(见图3)。此外,我们在该阶段对每个属性进行适当的归一化。此层可轻松抽象为一个Python装饰器,以支持用户直接提供的多种归一化过程实现。
我们还保留一个独立的掩码向量(见图4),以区分已分配的填充值(例如,NaN 或零)与属性的真实值。这在某些类型的神经网络中可能很重要,例如自编码器(AE)[27],,其中输入向量中填充值的位置可在解码阶段使用。
或者,可以从锯齿数组获得矩阵表示(见图5)。例如,空间坐标通常是高能物理数据集的一部分,因此,可用于事件的矩阵表示。这种方法可以解决向量表示中必须选择表示大小的问题,但它自身也存在空间矩阵粒度选择的问题。例如,在最简单的情况下,可以使用二维矩阵表示3(见图5)。然而,该图像的单元格大小是无法预先确定的。选择单元格大小可能会在事件内引入数据冲突问题,例如,不同的粒子可能具有位于同一单元格内的(X,Y)值对。此类歧义可通过增加矩阵粒度或使用附加属性来解决,例如通过单元格空间的更高维度实现。但此类更改将增加矩阵表示的稀疏性和矩阵大小,因此在训练时需要更多的计算资源。
机器学习训练工作流实现
我们使用Python编程语言实现了数据流和数据训练层,并且我们已在MLaaS4HEP仓库中提供[10]并采用MIT许可证。数据训练层被抽象化以支持任何基于Python的机器学习框架:TensorFlow、PyTorch及其他4。我们使用了两个参数来控制框架内的数据流。参数N_chunk用于控制数据流层从本地或远程存储读取的数据块大小。而参数N_batch定义了批次大小,即底层机器学习框架在每个训练周期中使用的事件数量。我们将数据流层读取的一组事件称为数据块,将机器学习训练循环使用的一组事件称为批次。
为了训练由用户代码(外部提供)定义的机器学习模型,MLaaS4HEP框架使用包含ROOT文件中相应比例事件的数据块。图6展示了数据流层和数据训练层中所用数据流的示意图。
第一遍(用1表示)表示MLaaS4HEP管道中用于创建规格文件的读取部分。该部分通过以数据块形式(其大小由用户预先固定)读取所有ROOT文件来实现,从而逐块更新规格文件中的信息。规格文件包含有关ROOT文件的所有信息:参差的分支的维度、每个分支的最小值和最大值,以及每个ROOT文件的事件数量5。流程图的第二部分如2所示,代表机器学习训练阶段。在循环的第一个周期中,当事件尚未读取时,我们从第i个文件N_chunk中读取f_i的c_i个事件,并将其存储到第i个数据块N_chunk⋅n_i/N_tot中。然后从中取出n_i个事件,其中f_i是来自文件N_tot的事件数量,而n_i是在步骤1中计算的信息)。对所有文件f_i进行事件读取和预处理。在从不同文件中正确混合创建了包含N_chunk个事件的数据块后,使用这些事件来训练机器学习模型。
训练阶段使用从该数据块中提取的数据批次进行,并运行一定数量的训练轮次。批次大小N_batch和训练轮次数N_epochs由用户预先确定。如果N_chunk不是N_batch的整数倍,则用于训练模型的最后一个批次包含少于N_batch个事件。然后我们回到循环的起始位置,如果存储在数据块c_i中的所有事件已经被读取,则从文件f_i中读取N_chunk个事件;否则,从数据块c_i中读取适当数量的事件(N_chunk⋅n_i/N_tot)。训练过程继续进行,如果文件尚未完全读取,则整个流水线从第2点的开头重新开始,直到所有事件被读取,在每个周期创建一个新的事件数据块,用于对该机器学习模型进行N_epochs轮次的训练。在周期结束时,即当我们读取了所有文件中的全部事件并完成了各个训练轮次的机器学习训练后,所有文件中包含的事件均已被读取,模型的训练过程完成,生成一个可用于物理分析的模型。
所讨论的训练过程可以应用于多种使用案例。尽管我们在第4节中使用了该方法,但不应将其视为使用MLaaS4HEP框架训练数据集的唯一方式。具体的机器学习策略最终由最终用户决定,应采取适当的步骤来检查模型的收敛性、选择合适的指标以监控训练周期等。例如,当数据集无法完全加载到训练节点的内存时,可采用其他解决方案,如使用随机梯度下降[28]模型。在这种情况下,机器学习训练工作流应调整为在每个训练周期中使用整个数据集。在此特定情况下,批量和数据块的概念将一致。
数据推理层
数据推理层可以通过多种方式实现。它既可以与应用框架紧密集成(例如,CMS和ATLAS实验分别在其CMSSW-DNN[29]和LTNN[30]解决方案中采用了这种方法),也可以作为服务(即服务)方案进行开发。前者的优势在于降低每个处理事件的推理步骤延迟,而后者则更容易推广,并且独立于内部基础设施。例如,它可以轻松集成到云平台中,可用作预训练模型的仓库,并能跨实验边界提供模型服务。然而,数据推理层的速度,即提供预测的吞吐量,可能因所选技术的不同而有所差异。选择HTTP协议可确保易于适配,而gRPC协议虽然能提供最佳性能,但需要专用的客户端。我们决定基于HTTP协议,将数据推理层实现为TensorFlow即服务架构[11]。
我们评估了多种机器学习框架,最终决定在推理阶段使用TensorFlow图[22]。TF模型以静态形式表示计算图,即数学计算、图边和数据流在运行时已明确定义。由于TF库提供的API支持,可以通过多种编程语言读取TF模型。此外,TF图针对GPU和TPU进行了高度优化。我们选择Go编程语言[31]来实现MLaaS4HEP框架的推理部分,主要基于以下因素:Go语言原生支持通过协程和通道实现并发;它是谷歌开发并使用的语言,与TF库集成良好;它提供最终的静态可执行文件,显著简化了在本地和各种(云)服务提供商上的部署。我们还选择了REST接口。客户端可以通过该接口将他们的TF模型上传到服务器,并用于其推理需求。基于REST API(端点)开发了Python和C++客户端,得益于HTTP协议,其他客户端也可以轻松开发。TFaaS框架可用于高能物理之外的场景,通过HTTP协议将任何基于TF的模型上传至TFaaS服务进行托管6。
MLaaS4HEP:概念验证原型
当MLaaS4HEP框架的所有层级开发完成后,我们使用可通过XrootD服务器访问的ROOT文件成功测试了该系统的可运行原型。数据以每块1千个事件的数据块形式读取,单个数据块大小约为4 MB。我们在本地机器上对该原型进行了测试,并成功将其部署到GPU节点上。为了进一步验证MLaaS4HEP框架,我们决定将其应用于真实物理分析,见第4节,其中我们研究了本地和远程数据访问、不同数据块的使用、文件的随机访问等。
真实案例场景
为了验证MLaaS4HEP方法,我们决定在真实的物理用例上测试该基础设施。这使我们能够测试MLaaS4HEP框架的性能,并从物理角度验证其结果。我们决定采用t̄t希格斯分析(t̄tH(b̄b))在增强型全强子末态[32, 33],因其与分析小组具有较高的关联性。在接下来的小节中,我们将讨论:
- t̄tH(b̄b)全强子分析策略(第4.1节);
- MLaaS4HEP验证(第4.2节);
- 使用物理用例的MLaaS4HEP性能结果(第4.3节);
- MLaaS4HEP预期性能(第4.4节);
- TFaaS性能结果(第4.5节)。
t̄tH(b̄b)全强子分析策略
在本小节中,我们提供了用于测试机器学习即服务性能并验证其在真实物理用例中功能的t̄t Higgs分析的详细信息。希格斯玻色子被认为是近几十年来高能物理领域最重要的发现。在被预言近50年后,它于2012年由ATLAS和CMS合作组织在欧洲核子研究中心的大型强子对撞机(LHC)上发现[34, 35]。自那以后,已开展了大量分析工作以更高精度测量其性质。
在标准模型框架中,希格斯玻色子通过类似汤川的相互作用与费米子耦合,该机制使费米子的质量正比于其耦合强度。最重的顶夸克主导了与希格斯玻色子的耦合。对顶夸克-希格斯耦合的直接测量利用的是树图过程。t̄tH产生在研究顶夸克-希格斯汤川耦合中起着重要作用,因为其他产生机制(如胶子-胶子融合)涉及圈图,其中超出标准模型(BSM)物理的贡献可能在圈图中未被察觉地进入。最高的分支比(≈25%)由全强子衰变道体现,即H(b̄b)以及全强子t̄t。由t̄t对产生的W玻色子衰变为一对轻夸克,而希格斯玻色子衰变为一对b̄b(见图7)。在末态中,至少存在八个部分子(初始态和末态辐射还可能产生更多),其中四个为底(b)夸克。尽管分支比最高,全喷注末态仍极具挑战性。它主要受到巨大的量子色动力学(QCD)LHC上的多喷注产生,由于存在大量喷注,该通道具有较大的不确定性。t̄tH的所有衰变产物均可观测,因此同时代表了完全重建的唯一可能性。
在13 TeV质心能量下,可通过p_T产生具有很高t̄tH的顶夸克。如果其洛伦兹提升足够高,则其衰变产物会高度准直地进入一个单一的宽喷注,称为boosted喷注。特别是,我们关注t̄tH(b̄b)的全喷注末态分析,其中末态中至少有一个喷注是boosted喷注,且希格斯玻色子衰变为一对由底夸克强子化形成的、可良好分辨的喷注。
为了识别t̄tH(b̄b)事件中包含的解析态希格斯粒子衰变,CMS在分析中使用了基于提升决策树(BDT)的机器学习模型[32, 33],并且在TMVA框架内完成了训练[36]。蒙特卡洛模拟提供了用于训练的事件,这些事件选自t̄tH样本以及两个主要的背景样本,即量子色动力学(QCD)和t̄t。匹配到两个b标记喷注系统的t̄tH事件中具有分辨出的希格斯玻色子的事件被视为信号事件。相反,未匹配的t̄tH事件,以及所有QCD和t̄t事件均被视为背景事件。信号事件和背景事件都必须满足某些选择标准,例如至少包含一个boosted喷注、不含轻子、通过信号触发等。该选择旨在挑选出boosted的全喷注类事件。
衰变)
MLaaS4HEP 验证
为了验证MLaaS4HEP功能相对于标准BDT方法的有效性,我们决定使用上一节讨论的分辨希格斯分析中的一组ROOT文件。本次实验的目标是证明MLaaS4HEP框架能够提供一种有价值的替代方案,并在预定义的一组指标基础上,提供与传统分析相当的结果。为此,我们选择使用一个通用的机器学习模型,并比较在MLaaS4HEP内部和外部获得的结果。具体而言,我们探索了以下方法:
- 使用MLaaS4HEP读取和归一化事件,并训练机器学习模型;
- 使用MLaaS4HEP读取和归一化事件,并使用Jupyter笔记本在MLaaS4HEP外部执行机器学习模型的训练;
- 使用Jupyter笔记本完成整个流水线,而不使用MLaaS4HEP。
最初,我们使用通过上一节中讨论的选择标准的ROOT文件进行分析。最终数据集包含八个包含背景事件的ROOT文件和一个包含信号事件的文件。每个文件有27个分支,总共包含35万个事件,该数据集的总大小为28 MB。信号事件与背景事件数量的比率约为10.8%。该数据集被分为三部分,分别用于训练(64%)、验证(16%)和测试(20%)。我们使用了Keras[21]具有两个隐藏层的序列神经网络,隐藏层分别由128和64个神经元组成,并采用0.5的dropout正则化在层之间。最后,我们以每批次100个事件的批次大小对模型进行了5个训练轮次的训练。
该练习的结果如图8所示,表明不同方法的性能相似。我们并未试图复现和/或匹配标准物理分析中获得的精确AUC数值,我们发现我们的结果(以AUC分数衡量)与物理分析中使用的BDT模型相当。当使用第3.3节中所述的训练工作流策略对上述机器学习模型进行基于数据块的机器学习训练时,该模型的收敛性仍然有效。然而,正如我们在第3.3节中指出的那样,一旦用户选择了特定的物理应用场景和机器学习模型,应验证该模型的收敛性,并在必要时调整适当的机器学习训练工作流。
MLaaS4HEP 性能
在本节中,我们提供了MLaaS4HEP性能测试的详细信息:该框架的可扩展性及其在不同存储层上的基准测试。为此,我们使用了所有可用的ROOT文件,且未施加任何物理筛选。这为我们提供了一个包含2850万个事件的数据集,共有74个分支(22个平坦分支和52个锯齿状分支),总大小约为10.1 GB。
我们在macOS(笔记本电脑)和CentOS 7 Linux系统上运行MLaaS4HEP框架进行了所有测试,其中macOS配置为2.2 GHz Intel Core i7双核、8 GB内存,CentOS 7 Linux配置为4个虚拟CPU的Intel Core处理器Haswell 2.4 GHz、7.3 GB内存的CERN虚拟机。ROOT文件从三个数据中心读取:博洛尼亚(BO)、比萨(PI)和巴里(BA)。平均可用带宽分别约为129 Mbit/s(标准误差(SDOM)参数为4 Mbit/s)和639 Mbit/s(SDOM = 39),前者使用macOS,后者使用CERN虚拟机(两种情况下的数值均基于10次试验得出)。
表1总结了我们在MLaaS4HEP管道的第一步(图6中的步骤1)中使用不同配置和10万事件的块大小所获得的I/O数据。该表提供了读取文件所用时间、计算规格值所用时间、完成步骤1的总时间,以及读取和规格计算步骤的事件吞吐量。
| 读取时间 (s) | 特征计算时间 (s) | 完成时间 步骤1 (秒) | 事件吞吐量 for 读取 + specs comp. (事件/秒) |
|---|---|---|---|
| macOS与本地文件 1633(9) | 958(2) | 2599(11) | 11055(49) |
| 带有远程文件的macOS (BO) 2365(49) | 974(10) | 3353(57) | 8585(149) |
| 带有本地文件的虚拟机 1131(3) | 963(2) | 2102(5) | 13690(34) |
| 带有远程文件的虚拟机(BO) 2455(68) | 959(2) | 3427(67) | 8396(158) |
| 带有远程文件的虚拟机(BA) 2304(88) | 961(2) | 3279(89) | 8801(241) |
| 带有远程文件的虚拟机(PI) 2129(41) | 1044(78) | 3186(83) | 9047(228) |
图9展示了在不同试验中,以块大小为函数的读取数据的事件吞吐量。在所有情况下,我们均未发现明显的峰值。较大的块大小可能会导致某些问题,例如在CERN虚拟机的情况下,可能会达到底层硬件的限制,如较大的内存占用。使用本地文件时,我们发现读取时间更短、事件吞吐量更高;而对于远程文件,结果主要受可用带宽的影响(即处理节点与存储数据的站点之间的链路连接性)。
在MLaaS4HEP管道第二步的性能研究中(2,见图6),我们关注的是数据读取部分以及数据预处理步骤(即包括数据转换)以及在MLaaS4HEP训练步骤中花费的时间。如第3.3节所述,存在一个文件循环,可用于按适当的事件比例构建用于训练机器学习模型的数据块。如果包含第i个ROOT文件事件的数据块为空或已完全处理,则会从第i个文件读取一个新的事件数据块,并将读取时间添加到创建该数据块所花费的总时间内(见图6)。换句话说,创建一个数据块所花费的时间由n次读取操作的时间以及预处理事件的时间之和构成。创建单个数据块的事件吞吐量和预处理单个数据块的事件吞吐量如表2所示。在图10中,我们展示了不同试验下创建数据块的事件吞吐量随块大小的变化情况。
我们发现,使用macOS或CERN虚拟机创建数据块的时间几乎相同,使用本地文件或远程文件时也大致相似。显然,对于远程文件,读取时间相应增加,创建数据块的时间也随之增加,但这一差异相当小。可以忽略不计。例如,我们花费约90秒创建一个包含10万事件的数据块,对应的事件吞吐量约为1.1k事件/秒,如表2所示。
| 事件吞吐量 for 创建块(事件/秒) | 事件吞吐量用于 预处理块(事件/秒) |
|---|---|
| 带有本地文件的macOS 1102(11) | 1157(7) |
| 带有远程文件的macOS (BO) 1057(17) | 1138(4) |
| 带有本地文件的虚拟机 1209(11) | 1247(2) |
| 带有远程文件的虚拟机(BO) 1110(32) | 1243(5) |
| 带有远程文件的虚拟机(BA) 1071(19) | 1153(4) |
| 带有远程文件的虚拟机(PI) 1152(18) | 1234(5) |
块大小的选择由用户决定,没有预定义的“最佳”值。我们建议用户从较低的块大小值开始,例如1千,,并根据资源可用性逐步增加。例如,在我们最初的原型实现中(见第3.5节),我们使用1千事件作为数据块大小,而在本节讨论的性能研究中,我们将块大小扩展到10万事件。
实际的机器学习训练时间与MLaaS4HEP框架无关,因为它由底层机器学习框架(例如Keras或PyTorch)、所用机器学习模型的复杂性以及可用硬件资源决定。特别是,使用第4.2节中介绍的简单机器学习模型和10万事件的块大小,我们发现每个数据块用于将数据正确划分为训练、验证和测试部分的时间约为1秒(在macOS和CERN虚拟机上几乎相等),而5个训练轮次的训练时间在macOS上约为11秒,在CERN虚拟机上约为13秒。
在实现MLaaS4HEP框架的过程中,我们解决了[24]中结果存在的若干瓶颈。例如,我们将读取时间提升了10倍,这得益于通过展平事件数组并计算每个分支的最小/最大值来更好地处理锯齿状数组。此外,在数据预处理步骤中,我们使用列表推导式替代事件内的循环,使性能提升了2.8倍。在macOS上,MLaaS4HEP框架预处理10万事件耗时约86秒,其中36%、26%、27%和6%的耗时分别用于将每个事件提取并转换为NumPy数组列表,归一化步骤、固定维度和创建掩码向量。
总之,我们证明了MLaaS4HEP方法可适用于所讨论的高能物理分析。使用10吉字节数据(约2850万个事件),我们获得了以下结果:
- MLaaS4HEP框架能够处理本地和远程文件;
- 读取本地ROOT文件(含规格计算)的吞吐量约为13.7k事件/秒,远程文件的吞吐量约为9k事件/秒;
- 预处理步骤的吞吐量峰值达到1.2k事件/秒。
MLaaS4HEP 性能预测
根据我们在前一节中介绍的研究,我们发现,对于处理2850万个事件(或10吉字节数据),MLaaS4HEP在流水线的第一步(1见图6)上大约需要35分钟,即获取所有已处理事件中各属性的最小/最大边界。MLaaS4HEP流水线的第二步(2见图6)大约需要7小时。这段时间包括从ROOT文件读取所有数据块、预处理事件(将锯齿状数组转换为固定维度的扁平NumPy数组、数据归一化等数据转换操作),以及将数据输入机器学习框架。实际的机器学习训练时间取决于用户提供的模型,并不代表MLaaS4HEP的性能。在第4.2节报告的研究中,这额外增加了一小时。因此,我们估计,在使用相同硬件资源的情况下,步骤1对于TB和PB规模的数据集将分别耗时O(100)小时和O(100k)小时,而步骤2的时间分别为O(1k)小时和O(1M)小时,此外还需加上训练机器学习模型所需的时间。5节将进一步讨论这些估算结果表明,为了处理TB或PB规模的数据集,需要对MLaaS4HEP流水线进行进一步优化,可能涉及I/O并行化、分布式机器学习训练及其他优化技术。在此阶段,我们的主要目标是证明MLaaS4HEP管道的可行性,并验证其在真实物理用例中的使用,而非在TB/拍字节规模上进行真正的机器学习训练。在第5节中,我们将讨论可以进一步改进的地方。
TFaaS 性能
TFaaS服务的性能测试使用了多种机器学习模型,从简单的图像分类到第4.2节中开发和讨论的机器学习模型。具体来说,我们在CentOS 7 Linux、16核、30 GB内存上运行的TFaaS服务器进行了多次基准测试。测试分为两种模式:使用100个并发客户端进行1000次调用,以及使用200个并发客户端进行5000次调用。在向TFaaS服务器发送和获取数据时,我们测试了JSON和ProtoBuffer[37]两种数据格式。在这两种情况下,我们都达到了∼500请求/秒的吞吐量。这些数值是通过提供一个用于第4.1节中讨论的物理分析的、具有27个特征和1024x1024隐藏层的中等规模预训练模型获得的。对于图像分类数据集(MNIST)也发现了类似的性能表现。TFaaS的实际性能将取决于所部署的机器学习模型的复杂性和可用的硬件资源。尽管单个TFaaS服务器可能不如集成解决方案高效,但它可以轻松地进行横向扩展,例如使用Kubernetes或其他集群编排解决方案,从而为并发客户端提供所需的吞吐量。此外,它还将应用层/框架与推理阶段解耦,可通过HTTP协议轻松集成到任何现有基础设施中。
未来方向
在前一节中,我们讨论了MLaaS4HEP在真实高能物理分析范围内的使用情况。我们发现以下内容:
- MLaaS4HEP的使用对所选高能物理数据集是透明的,即数据可从本地或远程存储中读取;
- 所讨论的架构与高能物理实验无关,可与任何现有的基于Python的机器学习方法结合使用框架以及轻松集成到现有基础设施中;
- 数据可以从远程存储中分块读取,从而实现对大型数据集的持续机器学习训练,并进一步实现并行化。
这些观察结果为使用现有的基于Python的开源机器学习框架在大型数据集(可能达到PB级别)上训练机器学习模型开辟了可能性。因此,我们预计机器学习即服务的方法在高能物理领域具有广泛的应用前景。例如,这项工作的未来方向可能包括利用该架构简化对云和高性能计算资源的访问,以支持训练与推理任务。这可以成为一种有吸引力的选择,用于开放高性能计算资源以支持高能物理中的大规模机器学习训练,同时结合所需的安全措施、资源供给以及对WLCG站点的远程数据访问。要实现这一方向,还需要进一步的工作。下面,我们将讨论一组可能的改进方案。
数据流层
为了改进数据流层,可以实现一个多线程I/O层。这可以通过将数据读取器代码库封装成一个服务来实现,该服务在接收到上游层的请求后并行提供数据块。此外,可以从XrootD服务器预取数据块到本地缓存中,以提高I/O吞吐量。特别是,目前正在进多项研发工作,旨在展示在计算中心(如高性能计算、国家Tier中心等)为动态按需分析服务(DODAS)提供智能缓存[38]的能力。此类DODAS设施可通过将ROOT文件预取到本地缓存中并用于机器学习训练,从而减少在数据流层上花费的时间。
数据训练层
如果能够实现数据I/O并行性,则可以通过实施分布式训练进一步改进[39]。在这方面有一些研发进展,包括采用Dask Python框架[40],或使用分布式Keras[41],或使用基于MPI的Python框架进行分布式训练[42],或在HDFS[43] Spark基础设施上使用MLflow框架+,这些方法探索了任务和数据并行性两种途径。机器学习框架的当前格局正在快速变化,我们应该将MLaaS4HEP调整以适应现有和未来的机器学习框架及创新。例如,开放网络交换格式[44]为模型从一个框架迁移到另一个框架打开了大门。这可能为使用MLaaS4HEP提供可能性下一代开源机器学习框架,并确保最终用户不会被锁定在特定框架上。例如,我们正在研究将PyTorch[20]和fast.ai[23]模型自动转换为TensorFlow,之后可通过TFaaS服务上传和使用[11]。
如第3.2节所述,将锯齿状数组输入机器学习框架有多种方法,该方向的研发工作正在进行中。例如,对于自编码器模型,带有填充值的向量表示应始终保留一个类型转换向量,以便后续将数据的向量表示解码回锯齿状数组或ROOT TTree数据结构。我们还希望探索锯齿状数组数据的矩阵表示,并研究其是否可应用于某些使用案例,例如在量热术或轨迹重建中,其中对象的图像表示可以被使用。
数据推理层
在推理方面,可以采用多种方法。如上所述,通过从HTTP切换到基于gRPC的解决方案(例如SONIC[11]),可以进一步提高TFaaS的吞吐量,[45]该方案能够基于FPGA和GPU基础设施提供快速的推理层。当前TFaaS的实现可作为预训练模型的仓库,通过HTTP协议提供机器学习模型服务,使得这些模型能够轻松地在实验边界或不同领域之间共享。例如,当前TFaaS的实现支持对上传模型的可视化检查、版本控制、标签功能等。我们预计下一步将发展为具备灵活搜索功能、扩展的模型标签功能和版本控制的预训练模型仓库。这可以通过为机器学习模型提供带有适当元数据描述的专用服务来实现。例如,此类元数据可以包含模型参数、所用软件的详细信息、版本发布、数据输入以及性能输出。配备适当的搜索引擎后,用户便可搜索与其使用场景相关的可用机器学习模型。
MLaaS4HEP服务
所提出的架构使我们能够将训练与推理层作为独立的服务进行开发和部署。可根据需要使用不同的资源提供者并动态扩展,例如,可通过商业云按需提供GPU/TPU用于特定模型的训练用途,而推理TFaaS服务则可部署在其他位置,例如某个计算中心的专用Kubernetes集群上。例如,当实验产生的数据放置在WLCG站点上时,即可实现复杂深度学习模型的持续训练。训练服务将接收到关于新产生的数据的一系列通知可用的数据,并重新训练特定模型。当新的机器学习模型准备就绪后,可以轻松推送到TFaaS,并立即供最终用户使用,而无需对现有基础设施进行任何干预,作为持续开发与持续集成(Continuous Development and Continuous Integration)工作流程的一部分。TFaaS还可进一步适配以使用现场可编程门阵列来加速推理阶段。我们预计,这种方法在高亮度LHC时代对于高能物理实验可能更加灵活且具有成本效益。因此,我们计划在此方向上开展更多的研发研究,并利用现有资源评估进一步的MLaaS4HEP服务。
摘要
本文中,我们提出了一种使用原生ROOT数据格式(来自本地或远程存储)训练高能物理机器学习模型的现代方法。MLaaS4HEP包含三个层:作为MLaaS4HEP框架组成部分的数据流层和数据训练层[10],以及基于TensorFlow库在TFaaS框架中实现的数据推理层。所有三层均被实现为独立组件。数据流层依赖uproot库从ROOT文件(本地或远程)读取数据并生成NumPy(锯齿)数组。数据训练层将输入的锯齿数组转换为向量表示,并将其传递给用户提供的机器学习框架。最后,数据推理层被实现为一个独立的HTTP服务。我们预计它可在多种使用案例中发挥作用,例如在物理分析中快速评估机器学习模型,或在可定期构建新模型的在线应用中。TFaaS的实现使其本身可作为高能物理预训练模型的仓库,并可成为任何团队敏捷机器学习开发周期中的宝贵组件,从小型物理分析小组到跨实验合作均可受益。
我们实现的灵活架构支持在大量分布式高能物理ROOT数据上进行机器学习训练,而无需将数据实际下载到本地存储中。我们证明了该架构能够读取通过WLCG基础设施上的XrootD协议提供的本地和分布式数据集。我们使用CMS官方t̄t希格斯分析(t̄tH(bb))在增强型全强子末态中对MLaaS4HEP架构进行了验证,并获得了与传统物理分析相当的机器学习模型性能,传统分析依赖于从ROOT文件中提取数据至自定义N元组。
1499

被折叠的 条评论
为什么被折叠?



