基于拉格朗日方法的流体模拟与连续卷积

本代码库包含了我们ICLR 2020论文的代码。我们展示了如何使用连续卷积将基于粒子的流体模拟网络训练为CNN。该代码允许你生成数据、训练自己的模型或仅运行预训练模型。

如果您觉得这份代码有用，请引用我们的论文(pdf)：

@inproceedings{Ummenhofer2020Lagrangian,
        title     = {Lagrangian Fluid Simulation with Continuous Convolutions},
        author    = {Benjamin Ummenhofer and Lukas Prantl and Nils Thuerey and Vladlen Koltun},
        booktitle = {International Conference on Learning Representations},
        year      = {2020},
}

我们建议您关注此代码库以获取最新更新。

依赖项

Tensorflow 2.3或PyTorch 1.6
Open3D 0.11或更高版本，包含ML模块 (https://github.com/intel-isl/Open3D/)
SPlisHSPlasH 2.4.0 (用于生成训练数据和流体粒子采样，https://github.com/InteractiveComputerGraphics/SPlisHSPlasH)
Tensorpack DataFlow (用于读取数据，pip install --upgrade git+https://github.com/tensorpack/dataflow.git)
python-prctl (Tensorpack DataFlow所需；依赖libcap-dev，通过apt install libcap-dev安装)
msgpack (pip install msgpack)
msgpack-numpy (pip install msgpack-numpy)
python-zstandard (pip install zstandard https://github.com/indygreg/python-zstandard)
partio (https://github.com/wdas/partio)
SciPy
OpenVDB及其Python绑定 (可选，用于创建表面网格，https://github.com/AcademySoftwareFoundation/openvdb)
plyfile (可选，用于创建表面网格，pip install plyfile)

这些版本与我们在Ubuntu 18.04系统上测试的配置相匹配。生成训练数据需要SPlisHSPlasH 2.4.0（确保以Release模式编译）。我们建议对所有其他包使用最新版本。

通过pip安装Open3D 0.11及更高版本

ML模块已包含在Open3D 0.11及更高版本中，可以简单地通过以下命令安装：

pip install open3d

确保您的ML框架版本与Open3D中ML操作构建时使用的版本匹配。对于Open3D 0.11，这是CUDA 10.1、TensorFlow 2.3和PyTorch 1.6。如果无法匹配此配置，建议从源代码构建Open3D。

从源代码构建包含ML模块的Open3D

目前，为了使本代码库中的代码正常工作，需要从源代码构建Open3D。要构建包含TensorFlow和PyTorch ML操作的Open3D，请执行以下步骤：

git clone --recursive https://github.com/intel-isl/Open3D.git
# 查看Open3D/util/scripts/install-deps-ubuntu.sh文件
# 了解依赖项并安装它们。更多说明请参阅Open3D文档

mkdir Open3D/build
cd Open3D/build

# 这将为TensorFlow和PyTorch构建ML操作。
# 如果您不需要两个框架，可以将不需要的框架设置为OFF。
cmake .. -DCMAKE_BUILD_TYPE=Release -DBUILD_TENSORFLOW_OPS=ON -DBUILD_PYTORCH_OPS=ON -DBUILD_CUDA_MODULE=ON -DGLIBCXX_USE_CXX11_ABI=OFF
make install-pip-package

运行预训练模型

预训练的网络权重文件位于TensorFlow的scripts/pretrained_model_weights.h5和PyTorch的scripts/pretrained_model_weights.pt中。

以下代码在示例场景上运行网络：

cd scripts
# 使用TensorFlow
./run_network.py --weights pretrained_model_weights.h5 \
                 --scene example_scene.json \
                 --output example_out \
                 --write-ply \
                 train_network_tf.py
# 或者使用PyTorch
./run_network.py --weights pretrained_model_weights.pt \
                 --scene example_scene.json \
                 --output example_out \
                 --write-ply \
                 train_network_torch.py

该脚本会将粒子位置以.ply文件格式输出点云，可以使用Open3D进行可视化。请注意，从.obj文件采样初始流体体积需要使用SPlisHSPlasH。

训练网络

数据生成

数据生成脚本位于datasets子文件夹中。要生成训练和验证数据，请按以下步骤操作：

在datasets/splishsplash_config.py脚本中设置SPlisHSPlasH的DynamicBoundarySimulator路径。
在datasets文件夹中运行脚本：
```
cd datasets
./create_data.sh
```

数据下载

如果您想跳过数据生成步骤，可以从以下链接下载训练和验证数据。

默认数据	34GB	链接
DPI大坝崩溃	24GB	链接
6k盒子数据	23GB	链接

对于默认数据，训练集是使用本仓库中的脚本生成的，验证数据与论文中使用的数据相对应。

DPI大坝崩溃数据是使用DPI-Nets仓库的代码生成的。请注意，数据已经按比例缩放以匹配我们方法使用的粒子半径。有关缩放因子的更多信息，请参阅scripts/dambreak.yaml配置文件。

6k盒子数据是默认数据的简化版本，具有恒定数量的粒子，并始终使用简单的盒子作为环境。

训练脚本

要使用生成的数据训练模型，只需在scripts文件夹中运行train_network_x.py脚本之一即可。

cd scripts
# TensorFlow版本
./train_network_tf.py default.yaml
# PyTorch版本
./train_network_torch.py default.yaml

这些脚本将创建一个train_network_tf_default或train_network_torch_default文件夹，其中包含快照和日志文件。可以使用Tensorboard查看日志文件。

评估网络

要评估网络，请按如下方式运行scripts/evaluate_network.py脚本：

./evaluate_network.py --trainscript train_network_tf.py --cfg default.yaml
# 或
./evaluate_network.py --trainscript train_network_torch.py --cfg default.yaml

这将创建train_network_{tf,torch}_default_eval_50000.json文件，其中包含帧对之间的各个误差。

脚本还将打印总体误差。如果您使用生成的数据，输出应该类似于： {'err_n1': 0.000859004137852537, 'err_n2': 0.0024183266885233934, 'whole_seq_err': 0.030323669719872864}

请注意，由于数据生成中的变化，这些数字与论文中的数字不同：