机器学习与流体动力学：谷歌AI利用「ML+TPU」实现流体模拟数量级加速

2021-02-08 02:47:56

机械之心报道

[原文来自：www.ii77.com]

作者：魔王

谷歌 AI 比来一项研究表明，行使机械进修和硬件加快器可以改善流体模拟，且不损害正确率或泛化机能。

[本文来自：www.ii77.com]

流体数值模拟对于建模多种物理现象而言非常主要，如天色、天气、空气动力学和等离子体物理学。流体能够用纳维 - 斯托克斯方程来描述，但大规模求解这类方程仍属难题，受限于解决最小时空特征的较量成本。这就带来了正确率和易处理性之间的衡量。

弗成压缩流体平日由如上纳维 - 斯托克斯方程来建模。

比来，来自谷歌 AI 的研究人员行使端到端深度进修改善较量流体动力学（CFD）中的近似，以建模二维涡流。对于湍流的直接数值模拟（direct numerical simulation, DNS）和大涡模拟（large eddy simulation, LES），该方式获得的正确率与基线求解器沟通，尔后者在每个空间维度的差别率是前者的 8-10 倍，因而该方式实现了 40-80 倍的较量加快。在较长模拟中，该方式仍能连结不乱，并泛化至练习所用流以外的力函数（forcing function）和雷诺数，这与黑箱机械进修方式正相反。此外，该方式还具备通用性，可用于随意非线性偏微分方程。

论文地址：https://arxiv.org/pdf/2102.01010.pdf

该研究作者之一、谷歌研究员 Stephan Hoyer 透露：这项研究表明，机械进修 + TPU 能够使流体模拟加快多达两个数量级，且不损害正确率或泛化机能。

至于结果若何呢？论文配合一作 Dmitrii Kochkov 展示了该研究提出的神经收集与 Ground truth、基线的结果对比：

首先是雷诺数 Re=1000 时，在 Kolmogorov 流上的结果对比：

其次是关于衰变湍流（decaying turbulence）的结果对比：

最后是雷诺数 Re=4000 时，在更复杂流上的结果对比：

方式简介

用非线性偏微分方程描述的复杂物理系统模拟对于工程与物理科学而言非常主要。然而，大规模求解这类方程并非易事。

谷歌 AI 这项研究提出一种方式来较量非线性偏微分方程解的正确时间演化，而且其使用的网格差别率比传统方式实现一致正确率要粗拙一个数量级。这种新型数值求解器不会对未解决的自由度取平均，而是使用离散方程，对未解决的网格给出逐点正确解。研究人员将受差别率损失影响最大的传统求解器组件替代为其学得的组件，行使机械进修发现了一些算法。

如下图 1a 所示，对于涡流的二维直接数值模拟，该研究提出的算法能够在每个维度的差别率粗拙 10 倍的情形下维持正确率不变，也就是说获得了 80 倍的较量时间改善。该模型进修若何对解的局部特征进行插值，从而可以正确泛化至分歧的流前提，如分歧受力前提，甚至分歧的雷诺数（图 1b）。

研究者还将该方式应用于涡流的高差别率 LES 模拟中，获得了雷同的机能提拔，在网格差别率粗拙 8 倍的情形下在 Re = 100, 000 LES 模拟中维持逐点正确率不变，实现约 40 倍的较量加快。

图 1：该研究提出方式与究竟概览。a）基线（direct simulation）与 ML 加快（learned interpolation）求解器的正确率与较量成本对比情形；b）练习与验证样本图示，展示出该模型壮大的泛化能力；c）该研究提出「learned interpolation」模型的单时间步构造，用卷积神经收集掌握尺度数值求解器对流较量中学得的近似。

研究者使用数据驱动离散化将微分算子插值到粗拙网格，且包管高正确率（图 1c）。具体而言，将求解底层偏微分方程的尺度数值方式内的求解器作为可微分编程进行练习，在 JAX 框架中写神经收集和数值方式（JAX 框架支撑反向模式主动微分）。这许可对整个算法执行端到端的梯度优化，与密度泛函理论、分子动力学和流体方面的之前研究雷同。研究者推导出的这些方式是特定于方程的，需要使用高差别率真值模拟练习粗拙差别率的求解器。因为偏微分方程的动态是局部的，是以高差别率模拟能够在小型域内实施。

该算法的工作流程如下：在每一个时间步中，神经收集在每个网格位置基于速度场生成隐向量，然后求解器的子组件使用该向量处理局部解构造。该神经收集为卷积收集，具备平移不变性，因而许可解构造在空间中是局部的。之后，使用尺度数值方式的组件执行纳维 - 斯托克斯方程对应的概括偏置，如图 1c 灰色框所示：对畅通量（convective flux）模型改善离散对流算子的近似；散度算子（divergence operator）基于有限体积法执行局部动量守恒；压力投影（pressure projection）实现弗成压缩性，显式时间步算子（explicit time step operator）使动态具备时间一连性，并许可额外时变力的插值。「在更粗拙网格上的 DNS」将传统 DNS 和 LES 建模的界线恍惚化，从而获得多种数据驱动方式。

该研究首要存眷两种 ML 组件：learned interpolation 和 learned correction。此处不再赘述，详情拜见原论文。

实验究竟

加快 DNS

一旦网格差别率无法捕获到解的最小细节，则 DNS 的正确率将快速下降。而该研究提出的 ML 方式极大地缓解了这一效应。下图 2 展示了雷诺数 Re = 1000 的情形下在 Kolmogorov 流上练习和评估模型的究竟。

而就较量效率而言，10 倍网格粗拙度的情形下，learned interpolation 求解器取得与 DNS 一致正确率的速度也要更快。研究者在单个谷歌云 TPU v4 内核上对该求解器进行了基准测试，谷歌云 TPU 是用于机械进修模型的硬件加快器，也适用于很多科学较量用例。在充沛大的网格巨细（256 × 256 甚至更大）上，该研究提出的神经收集可以很好地行使矩阵乘法单元，每秒浮点运算的吞吐量是基线 CFD 求解器的 12.5 倍。是以，尽管使用了 150 倍的算术运算，该 ML 求解器所用时间仍然仅有一致差别率下传统求解器的 1/12。三个维度（两个空间维度和一个时间维度）中有效差别率的 10 倍提拔，带来了 10^3/12 ≈ 80 倍的加快。

此外，研究者还考虑了三种分歧的泛化测试：大型域规模；非受迫衰减涡流；较大雷诺数的 Kolmogorov 流。

首先，研究者将同样的力泛化至较大的域规模。该 ML 模型获得了与在练习域中同样的机能，因为它们仅依靠流的局部特征（拜见下图 5）。

然后，研究者将在 Kolmogorov 流上练习的模型应用于衰减涡流。下图 3 表明，在 Kolmogorov 流 Re = 1000 上学得的离散模型的正确率能够成家以 7 倍差别率运行的 DNS。

最后，该模型能够泛化至更高的雷诺数吗？也就是更复杂的流。下图 4a 表明，该模型的正确率能够成家以 7 倍差别率运行的 DNS。鉴于该测试是在复杂度显著增加的流长进行的，是以这种泛化结果很不错。图 4b 对速度进行了可视化，表明该模型能够处理更高的复杂度，图 4c 的能谱进一步验证了这一点。

与其他 ML 模型进行对比

研究者将 learned interpolation 与其他 ML 方式的机能进行了对比，包罗 ResNet (RN) [50]、Encoder Processor-Decoder (EPD) [51, 52] 架构和之前介绍的 learned correction (LC) 模型。下图 5 展示了这些方式在所有考虑设置中的究竟。总体而言，learned interpolation (LI) 机能最佳，learned correction (LC) 紧随厥后。