本技术属于人工智能及高性能计算技术领域,特别涉及一种提高深度算子网络模型推理求解和训练速度的方法。本发明针对PI‑DeepONet模型进行多层次的优化,提高其推理求解速度和训练速度。本发明定制SGEMM算子,在PI‑DeepONet常用的模型尺寸范围内,SGEMM算子相较于CUDA计算平台的基础线性代数库中最好的单精度矩阵乘法可取得最高1.5倍的加速比。针对模型中使用的modified MLP的特殊结构设计融合算子,模型的端到端推理速度获得2倍以上的加速比,图形处理器计算时间最高可被缩短5.7倍。本发明使用数据并行对模型进行分布式训练,并使用通信掩盖的策略,可扩展性保持在83%‑85%。
背景技术
偏微分方程的求解在工程技术、生物数学和图像处理等许多科研领域的研究中都起着至关重要的作用。例如,在石油工程领域中,通过求解伯格斯方程、达西渗流方程等流体力学方程,可以得到油藏中流体的速度和压力分布情况,构建出真实油田的数值模型,进而对储量评估、产能预测和油田开发方案预测等方面的决策起到指导性作用,帮助提高石油勘探开发的效率和经济效益。
长期以来,偏微分方程的求解一直是一项非常重要且具有挑战性的任务。一方面,许多物理系统的数学建模是通过偏微分方程构建的,求解偏微分方程是这些领域中开展科学研究的重要手段,也是许多工业软件和物理计算的基础。经典的求解方法包括有限差分法和有限元法等,它们通常在给定的定义域上进行复杂的网格划分,将连续的问题域划分为离散的网格,并在划分的结果上使用简单的函数进行近似,最终得到方程的近似解。另一方面,使用传统数值方法求解偏微分方程存在着若干难点:首先,这些方法都依赖于网格划分,为某一特定的求解问题设计合适的网格这一过程本身就具有挑战性;然后,这些方法在计算过程每一步都伴随着大量的迭代计算,这往往意味着很高的计算代价;此外,使用数值方法求解偏微分方程,其流程可能会因为问题的不同而差异巨大。
最近,基于深度学习的方法成为了求解偏微分方程的另一类有效手段,并广泛应用于流体力学、热传导方程、材料学等许多领域。这类方法已经可以取得较好的精度;同时,相较于传统数值方法,训练好的深度学习模型在求解速度方面有巨大的优势。
自上世纪90年代起,已有若干研究致力于使用神经网络来求解偏微分方程。其中,Dissanayake和Phan Thien将偏微分方程的求解问题转化为无约束最小化问题,提出了使用神经网络求解偏微分方程的算法,并使用其算法求解了泊松方程和非线性热传导方程,这项工作是使用神经网络求解偏微分方程这一研究方向的开端。但受限于当时计算机的算力和所使用的神经网络模型的局限性,这些方法往往只能用于简单偏微分方程的求解,因此它们在当时并未能得到广泛的关注和足够的重视。
2019年,Raissi等人提出了一项具有开创性的工作PINN,他们在深度神经网络中使用偏微分方程的控制方程、初始值和边界值条件来构造损失函数约束网络模型的收敛方向,相较于单纯使用标签数据来进行训练,这种思路下训练得到的模型更能符合方程本身的条件。在此之后,许多PINN的变体被提出。UQPINN结合了PINN和对抗学习的策略来处理各种偏微分方程。fPINN结合了PINN和时间离散化技术用以求解时间分数阶偏微分方程。PPINN通过将一个长时问题分解为多个独立的短时问题来减小算法的训练代价并更高效地处理长时问题。但不论进行了何种程度的改进,其核心思想都是将偏微分方程本身的条件加入到深度神经网络网络模型中来约束模型的训练,因此它们都是基于物理信息的深度神经网络。
另一类方法是使用具有独特结构的算子神经网络来求解偏微分方程,这类方法的网络模型逼近的对象不再是偏微分方程的解函数,而是将基函数映射为解函数的算子,这类方法往往具有更好的泛化性。其中最具代表性的两个工作分别是Lu Lu等提出的深度算子网络(DeepONet)和ZongYi Li等提出的傅里叶神经算子(FNO)。之后,基于这两种算子神经网络的改进工作和应用工作也被相继提出。S-DeepONet通过在DeepONet的分支网络中引入时序网络结构来用于输入数据具有时序性的问题。Goswami S等提出一种提高深度算子网络模型推理求解和训练速度的方法基于物理原理的V-DeepONet用于脆性断裂分析。GegeWen等提出U-FNO用于多相流问题的求解。Meer Mehran Rashid等使用FNO来学习二维复合材料的力学响应。
虽然基于算子神经网络求解偏微分方程的方法具有泛用性强的优点,但因为其纯数据驱动的方式而无法应对标签数据较少甚至没有标签数据的情况。而在现实问题中,一些领域的真实数据可能难以大量获取或者获取成本非常高昂,只有非常少量的标签数据,因此也有一些工作尝试将算子神经网络与基于物理信息的神经网络的思想结合,在算子神经网络的模型中加入偏微分方程的条件作为约束。PI-DeepONet是一种提高深度算子网络模型推理求解和训练速度的方法先进的用于求解偏微分方程的深度学习模型,PI-DeepONet在DeepONet网络模型的基础上,用偏微分方程的控制函数和初边值条件构建损失函数,实现了在完全没有标签数据的情况下依然能够训练网络模型至收敛,并取得了较好的准确性。
但是,在使用PI-DeepONet求解偏微分方程时,模型的训练和推理速度会因为模型的一些结构特点受到限制:较窄的全连接层宽度和较大的批次大小会导致低效的矩阵乘法计算,许多逐元素算子的频繁启动也会带来额外的开销,这将限制模型的推理速度;另一方面,由于使用了PINN结构,大量求导计算被引入训练过程中,模型训练的时间开销也变得难以忽视。
实现思路