Stanford CS224W: Machine Learning with Graphs

By Haochen Shi, Peng Chen, Shiyu Li as part of the Stanford CS 224W Final Project


代码下载：https://t.zsxq.com/hznTQ

请索引第39个项目



01 概述


许多重要任务需要对复杂的物理过程进行高分辨率模拟。例如，从事基于模型规划的科学家希望了解具有复杂物理特性的物体（例如可变形物体）在环境变化后（例如，如本文[ 3]所示，被机械臂的夹爪夹住）将如何变形。此外，电影行业也渴望制作逼真的特效，例如大规模爆炸，以打造更引人注目的场景。


一种很有前景的复杂物理模拟方法是训练一个能够学习如何模拟的图神经网络（GNN）。在本教程中，我们将介绍为什么这种方法适合处理这项任务，并提供构建能够模拟复杂物理的 GNN 的分步指南。本教程基于 DeepMind 的论文《Simulate Complex Physics with Graph Networks》 [1]。


这是DeepMind公司原始论文的预告视频。来源：https://arxiv.org/abs/2002.09405


02 为什么不使用传统的物理模拟器？


在 Taichi 物理引擎中，使用 V100 GPU 耗时 8 小时模拟了 1 亿个粒子。来源：https ://github.com/taichi-dev/taichi_elements


虽然传统的物理模拟器功能强大，但它们也存在一些重要的缺点：1）对于传统的物理模拟器来说，要获得高质量的大规模模拟结果既昂贵又耗时；2）要设置物理模拟器，我们需要完全了解物体和环境的物理参数，而这些参数在某些情况下是极其难以知道的。


为了解决这些问题，许多计算机科学家致力于利用机器学习方法来模拟复杂的物理现象。这些机器学习模型无需人类提供额外信息，即可直接从视觉输入中学习物体的复杂物理特性，并对物体的未来状态做出准确高效的预测。


03 为什么 GNN 是完成这项任务的理想选择？

3.1 图神经网络的输入


非结构化数据和结构化数据。


图神经网络 (GNN) 是神经网络的一个特殊子集，它接受结构化程度较低的数据（例如图）作为输入，而其他神经网络（例如卷积神经网络 (CNN) 和 Transformer）只能接受结构化程度较高的数据（例如网格和序列）。所谓“结构化程度较低”，是指输入数据可以具有任意形状和大小，并且可以具有复杂的拓扑关系。


在基于粒子的物理模拟中，我们以所有粒子的非结构化位置信息作为输入，这启发了我们使用图神经网络的想法。


3.2 排列等变性


图神经网络 (GNN) 的一个关键特征，也是它与其他神经网络的区别所在，是其排列等价性。也就是说，图中的节点没有固定的顺序，因此我们如何“排列”图中的节点并不影响 GNN 的计算结果。


由于在基于粒子的模拟中，物体的粒子是“相同的”，因此当对它们应用物理定律时，它们是置换等变的。所以，像图神经网络（GNN）这样的置换等变模型适合模拟粒子间的相互作用。


3.3 图神经网络的结构


图神经网络（GNN）层的结构也影响我们选择使用GNN来模拟物理过程的决定。GNN每一层的关键步骤如下：



直观地说，基于粒子的物理模型与上述 GNN 结构完美契合：1）节点代表粒子的物理状态；2）通过边传递消息类似于粒子之间的成对相互作用；3）邻居聚合模拟所有邻居对粒子的相互作用的总影响；4）更新步骤在计算所有相互作用后更新每个粒子的位置。


04 数据集

数据集中的水滴轨迹。


数据集可在此处找到(https://github.com/deepmind/deepmind-research/tree/master/learning_to_simulate)。它包含详尽的场景，包括具有不同物理性质（沙子、粘液和水）的粒子在不同场景（摇晃和掉落）下的运动轨迹。我们的任务是根据粒子的初始条件和系统的全局属性预测粒子的运动轨迹。


在本教程中，我们选择 WaterDrop 数据集。


4.1 任务设置


每个粒子的输入向量包含粒子的当前位置、过去5步的速度、代表粒子物理参数的特征，以及（如果适用）系统的全局属性（例如，外力）。图神经网络（GNN）的目标是预测每个粒子的加速度。利用GNN预测的加速度，欧拉积分器更新粒子的速度和位置。


4.2 损失和指标


损失函数和最终评价指标均由均方误差（MSE）定义。图神经网络（GNN）输出粒子的加速度，因此损失函数是加速度的均方误差，即：



然而，由图神经网络和欧拉积分器组成的学习型模拟器输出的是粒子的位置而不是加速度。因此，学习型模拟器的度量标准定义为位置的均方误差（MSE），如下所示：



05 基于图网络的模拟器架构


本教程中使用的模型是DeepMind提出的基于图网络的模拟器（GNS）[1]。在GNS中，节点代表粒子，边代表粒子之间的交互。GNS由编码器、处理器和解码器三个部分组成。为了更清晰地展示GNS的结构，我们将DeepMind的开源TensorFlow实现重构为一个更简洁的PyG版本。


5.1 编码器


编码器对图神经网络 (GNN) 的数据进行预处理。它会在任意两个距离小于阈值 R 的粒子之间添加一条边。边特征表示两个粒子之间的相对位移。在将这些特征输入 GNN 之前，编码器会分别通过多层感知器 (MLP) 对节点特征和边特征进行编码。


节点编码器屏蔽粒子的绝对位置信息，并将粒子的速度和物理参数编码到节点特征中。边缘编码器将粒子与其邻居之间的相对位置位移编码到边缘特征中。


5.2 处理器


该处理器是一个图神经网络（GNN），由多个相同的交互网络（IN）[2]层堆叠而成。在每个IN层中，节点和边都具有隐藏表示。

层内消息传递、聚合和更新的示意图。


在 PyTorch Geometric (PyG) 中，可以轻松实现诸如 IN 层之类的 GNN 层。在 PyG 中，GNN 层通常实现为 MessagePassing 类的子类。我们遵循此约定，并将 InteractionNetwork 类定义如下。



1）为图中的每条边构建一条消息。该消息的生成方法是将边的两个节点的特征与边本身的特征连接起来，然后用多层感知器（MLP）对连接后的向量进行变换。



2) 将每个节点所有传入边的消息聚合（求和）。



3）更新节点特征和边特征。每条边的新特征是其原特征加上该边上的消息之和。每个节点的新特征由其原特征加上消息的聚合值决定。



5.3 解码器


解码器使用多层感知器（MLP）解码粒子加速度及其处理器的节点嵌入。它也可以被视为图神经网络（GNN）的后处理器。解码后的加速度由欧拉积分器更新粒子的速度和位移。


5.4 整合所有元素


让我们把编码器、处理器和解码器放在一起！在 GNN 层之前，输入特征会通过 MLP 进行转换，从而在不增加 GNN 层数的情况下提高 GNN 的表达能力。在 GNN 层之后，最终输出（在本例中为粒子加速度）会从 GNN 层生成的特征中提取出来，以满足任务需求。



06 GNS 的运行模式


GNS 有两种工作模式：单步模式和滚动模式。在单步模式下，GNS 始终使用真实值进行预测。在滚动模式下，GNS 基于前一步的预测结果来预测下一步粒子的位置。因此，滚动模式下误差会随时间累积。

GNS的运行模式：单步模式和展开模式。每个箭头都从模型的输入指向模型的相应输出。


07 模型的主要特点

7.1 空间不变性的归纳偏置


粒子间的物理相互作用与其空间位置无关。为了保持这种不变性，编码器会屏蔽粒子的绝对位置信息，而是将粒子间的相对位移编码成特征。这使得模型能够对空间不变性具有很强的归纳偏置。


7.2 鲁棒性


在模拟复杂混沌系统物理特性时，误差随时间累积会对结果产生巨大的负面影响。为了训练模型，研究人员会在输入的真实数据中加入随机游走噪声，以模拟先前运行过程中产生的误差。这使得模型能够有效应对此类误差累积。


7.2 不同材质，同一模型


不同材料的物理特性各不相同。大多数模拟器仅适用于一种类型的材料[1]。GNS提出了一种模拟不同材料物理特性的通用方法。该方法能够对流体、刚性材料和可变形材料进行精确预测。


08 结果


这是使用 GeForce RTX 3080 Ti 显卡，在整个 WaterDrop 数据集上训练模型 5 个 epoch 后的结果，大约需要 14 个小时。

真实值（左）与我们的结果（右）的比较

训练损失和验证损失。

GNS 在单步模式下运行时的均方误差指标。

GNS 在部署模式下运行时的均方误差指标。


我们绘制了训练过程中的演化损失和指标曲线。损失曲线和单步均方误差（MSE）曲线表明模型在训练结束时收敛。但滚动均方误差（rollout MSE）指标曲线表明模型在长时间内稳定性较差。一种解释是，模拟系统本质上是混沌的，因此一个微小的数值误差就会导致系统行为发生完全不同的变化。


最后，我们在测试集上测量了该指标。我们的复现结果达到了单步均方误差 (MSE) 3.04e-9 和滚动均方误差 (MSE) 1.42e-2，考虑到我们计算能力有限，这些结果与原始论文相比是合理的。


09 结论


在本教程中，我们一步步地讲解了如何在实际应用中使用图神经网络（GNN），例如模拟复杂的物理现象！如果您对技术细节感兴趣，请阅读DeepMind的原始论文。


参考文献

  • A. Sanchez-Gonzalez, J. Godwin, T. Pfaff, R. Ying, J. Leskovec, and P. W. Battaglia. Learning to simulate complex physics with graph networks. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13–18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 8459–8468. PMLR, 2020. URL http://proceedings.mlr.press/v119/sanchez-gonzalez20a.html.

  • P. W. Battaglia, R. Pascanu, M. Lai, D. J. Rezende, and K. Kavukcuoglu. Interaction networks for learning about objects, relations and physics. In D. D. Lee, M. Sugiyama, U. von Luxburg, I. Guyon, and R. Garnett, editors, Advances in Neural Information Processing Systems 29: Annual Conference on Neural Information Processing Systems 2016, December 5–10, 2016, Barcelona, Spain, pages 4502–4510, 2016. URL https://proceedings.neurips.cc/paper/2016/hash/3147da8ab4a0437c15ef51a5cc7f2dc4-Abstract.html.

  • Y. Li, J. Wu, R. Tedrake, J. B. Tenenbaum, and A. Torralba. Learning particle dynamics for manipulating rigid bodies, deformable objects, and fluids. In ICLR, 2019.

微信群

内容中包含的图片若涉及版权问题，请及时与我们联系删除