From e00711d7e50f7296e17dfd61f0ac299f49a432d6 Mon Sep 17 00:00:00 2001 From: Artur Toshev Date: Sat, 24 Feb 2024 03:31:12 +0000 Subject: [PATCH] update readme and requirements --- README.md | 27 +++++++++++++++------------ docs/requirements.txt | 2 ++ requirements_cuda.txt | 2 ++ 3 files changed, 19 insertions(+), 12 deletions(-) diff --git a/README.md b/README.md index befa950..6463b79 100644 --- a/README.md +++ b/README.md @@ -74,7 +74,7 @@ pip install --upgrade jax[cuda12_pip]==0.4.20 -f https://storage.googleapis.com/ ### MacOS Currently, only the CPU installation works. You will need to change a few small things to get it going: - Clone installation: in `pyproject.toml` change the torch version from `2.1.0+cpu` to `2.1.0`. Then, remove the `poetry.lock` file and run `poetry install --only main`. -- Configs: You will need to set `f32: True` and `num_workers: 0` in the `configs/` files. +- Configs: You will need to set `dtype=float32` and `train.num_workers=0`. Although the current [`jax-metal==0.0.5` library](https://pypi.org/project/jax-metal/) supports jax in general, there seems to be a missing feature used by `jax-md` related to padding -> see [this issue](https://github.com/google/jax/issues/16366#issuecomment-1591085071). @@ -85,37 +85,37 @@ A general tutorial is provided in the example notebook "Training GNS on the 2D T ### Running in a local clone (`main.py`) Alternatively, experiments can also be set up with `main.py`, based on extensive YAML config files and cli arguments (check [`configs/`](configs/)). By default, the arguments have priority as 1) passed cli arguments, 2) YAML config and 3) [`defaults.py`](lagrangebench/defaults.py) (`lagrangebench` defaults). -When loading a saved model with `--load_ckp` the config from the checkpoint is automatically loaded and training is restarted. For more details check the [`experiments/`](experiments/) directory and the [`run.py`](experiments/run.py) file. +When loading a saved model with `load_ckp` the config from the checkpoint is automatically loaded and training is restarted. For more details check the [`runner.py`](lagrangebench/runner.py) file. **Train** For example, to start a _GNS_ run from scratch on the RPF 2D dataset use ``` -python main.py --config configs/rpf_2d/gns.yaml +python main.py config=configs/rpf_2d/gns.yaml ``` Some model presets can be found in `./configs/`. -If `--mode=all`, then training (`--mode=train`) and subsequent inference (`--mode=infer`) on the test split will be run in one go. +If `mode=all` is provided, then training (`mode=train`) and subsequent inference (`mode=infer`) on the test split will be run in one go. **Restart training** -To restart training from the last checkpoint in `--load_ckp` use +To restart training from the last checkpoint in `load_ckp` use ``` -python main.py --load_ckp ckp/gns_rpf2d_yyyymmdd-hhmmss +python main.py load_ckp=ckp/gns_rpf2d_yyyymmdd-hhmmss ``` **Inference** -To evaluate a trained model from `--load_ckp` on the test split (`--test`) use +To evaluate a trained model from `load_ckp` on the test split (`test=True`) use ``` -python main.py --load_ckp ckp/gns_rpf2d_yyyymmdd-hhmmss/best --rollout_dir rollout/gns_rpf2d_yyyymmdd-hhmmss/best --mode infer --test +python main.py load_ckp=ckp/gns_rpf2d_yyyymmdd-hhmmss/best rollout_dir=rollout/gns_rpf2d_yyyymmdd-hhmmss/best mode=infer test=True ``` -If the default `--out_type_infer=pkl` is active, then the generated trajectories and a `metricsYYYY_MM_DD_HH_MM_SS.pkl` file will be written to the `--rollout_dir`. The metrics file contains all `--metrics_infer` properties for each generated rollout. +If the default `eval.infer.out_type=pkl` is active, then the generated trajectories and a `metricsYYYY_MM_DD_HH_MM_SS.pkl` file will be written to `eval.rollout_dir`. The metrics file contains all `eval.infer.metrics` properties for each generated rollout. ## Datasets -The datasets are hosted on Zenodo under the DOI: [10.5281/zenodo.10021925](https://zenodo.org/doi/10.5281/zenodo.10021925). When creating a new dataset instance, the data is automatically downloaded. Alternatively, to manually download them use the `download_data.sh` shell script, either with a specific dataset name or "all". Namely +The datasets are hosted on Zenodo under the DOI: [10.5281/zenodo.10021925](https://zenodo.org/doi/10.5281/zenodo.10021925). If a dataset is not found in `dataset_path`, the data is automatically downloaded. Alternatively, to manually download the datasets use the `download_data.sh` shell script, either with a specific dataset name or "all". Namely - __Taylor Green Vortex 2D__: `bash download_data.sh tgv_2d datasets/` - __Reverse Poiseuille Flow 2D__: `bash download_data.sh rpf_2d datasets/` - __Lid Driven Cavity 2D__: `bash download_data.sh ldc_2d datasets/` @@ -144,7 +144,8 @@ We provide three notebooks that show LagrangeBench functionalities, namely: ┃ ┗ 📜utils.py ┣ 📂evaluate # Evaluation and rollout generation tools ┃ ┣ 📜metrics.py - ┃ ┗ 📜rollout.py + ┃ ┣ 📜rollout.py + ┃ ┗ 📜utils.py ┣ 📂models # Baseline models ┃ ┣ 📜base.py # BaseModel class ┃ ┣ 📜egnn.py @@ -157,6 +158,7 @@ We provide three notebooks that show LagrangeBench functionalities, namely: ┃ ┣ 📜strats.py # Training tricks ┃ ┗ 📜trainer.py # Trainer method ┣ 📜defaults.py # Default values + ┣ 📜runner.py # Runner wrapping training and inference ┗ 📜utils.py ``` @@ -195,7 +197,7 @@ pytest ### Clone vs Library LagrangeBench can be installed by cloning the repository or as a standalone library. This offers more flexibility, but it also comes with its disadvantages: the necessity to implement some things twice. If you change any of the following things, make sure to update its counterpart as well: -- General setup in `experiments/` and `notebooks/tutorial.ipynb` +- General setup in `lagrangebench/runner.py` and `notebooks/tutorial.ipynb` - Configs in `configs/` and `lagrangebench/defaults.py` - Zenodo URLs in `download_data.sh` and `lagrangebench/data/data.py` - Dependencies in `pyproject.toml`, `requirements_cuda.txt`, and `docs/requirements.txt` @@ -232,3 +234,4 @@ The associated datasets can be cited as: The following further publications are based on the LagrangeBench codebase: 1. [Learning Lagrangian Fluid Mechanics with E(3)-Equivariant Graph Neural Networks (GSI 2023)](https://arxiv.org/abs/2305.15603), A. P. Toshev, G. Galletti, J. Brandstetter, S. Adami, N. A. Adams +2. [Neural SPH: Improved Neural Modeling of Lagrangian Fluid Dynamics](https://arxiv.org/abs/2402.06275), A. P. Toshev, J. A. Erbesdobler, N. A. Adams, J. Brandstetter diff --git a/docs/requirements.txt b/docs/requirements.txt index a19b4eb..acdb3c6 100644 --- a/docs/requirements.txt +++ b/docs/requirements.txt @@ -9,12 +9,14 @@ jax_md>=0.2.8 jmp>=0.0.4 jraph>=0.0.6.dev0 matscipy>=0.8.0 +omegaconf>=2.3.0 optax>=0.1.7 ott-jax>=0.4.2 pyvista PyYAML sphinx==7.2.6 sphinx-rtd-theme==1.3.0 +toml>=0.10.2 torch==2.1.0+cpu wandb wget diff --git a/requirements_cuda.txt b/requirements_cuda.txt index 3c969e3..abb46bb 100644 --- a/requirements_cuda.txt +++ b/requirements_cuda.txt @@ -11,10 +11,12 @@ jax_md>=0.2.8 jmp>=0.0.4 jraph>=0.0.6.dev0 matscipy>=0.8.0 +omegaconf>=2.3.0 optax>=0.1.7 ott-jax>=0.4.2 pyvista PyYAML +toml>=0.10.2 torch==2.1.0+cpu wandb wget