Krishna Kumar, The University of Texas at Austin. Joseph Vantassel, Texas Advanced Computing Center, UT Austin. Yongjin Choi, The University of Texas at Austin.
Graph Network-based Simulator (GNS) is a generalizable, efficient, and accurate machine learning (ML)-based surrogate simulator for particulate and fluid systems using Graph Neural Networks (GNNs). GNS code is a viable surrogate for numerical methods such as Material Point Method, Smooth Particle Hydrodynamics and Computational Fluid dynamics. GNS exploits distributed data parallelism to achieve fast multi-GPU training. The GNS code can handle complex boundary conditions and multi-material interactions.
MeshNet is a scalable surrogate simulator for any mesh-based models like Finite Element Analysis (FEA), Computational Fluid Dynammics (CFD), and Finite Difference Methods (FDM).
Training GNS/MeshNet on simulation data
# For particulate domain,
python3 -m gns.train --data_path="<input-training-data-path>" --model_path="<path-to-load-save-model-file>" --ntraining_steps=100
# For mesh-based domain,
python3 -m meshnet.train --data_path="<input-training-data-path>" --model_path="<path-to-load-save-model-file>" --ntraining_steps=100Resume training
To resume training specify model_file and train_state_file:
# For particulate domain,
python3 -m gns.train --data_path="<input-training-data-path>" --model_path="<path-to-load-save-model-file>" --model_file="model.pt" --train_state_file="train_state.pt" --ntraining_steps=100
# For mesh-based domain,
python3 -m meshnet.train --data_path="<input-training-data-path>" --model_path="<path-to-load-save-model-file>" --model_file="model.pt" --train_state_file="train_state.pt" --ntraining_steps=100Rollout prediction
# For particulate domain,
python3 -m gns.train --mode="rollout" --data_path="<input-data-path>" --model_path="<path-to-load-save-model-file>" --output_path="<path-to-save-output>" --model_file="model.pt" --train_state_file="train_state.pt"
# For mesh-based domain,
python3 -m meshnet.train --mode="rollout" --data_path="<input-data-path>" --model_path="<path-to-load-save-model-file>" --output_path="<path-to-save-output>" --model_file="model.pt" --train_state_file="train_state.pt"Render
# For particulate domain,
python3 -m gns.render_rollout --output_mode="gif" --rollout_dir="<path-containing-rollout-file>" --rollout_name="<name-of-rollout-file>"
# For mesh-based domain,
python3 -m gns.render --rollout_dir="<path-containing-rollout-file>" --rollout_name="<name-of-rollout-file>"In particulate domain, the renderer also writes .vtu files to visualize in ParaView.
GNS prediction of Sand rollout after training for 2 million steps.
In mesh-based domain, the renderer writes .gif animation.
Meshnet GNS prediction of cylinder flow after training for 1 million steps.
`train.py` in GNS (particulate domain)
mode (Enum)
This flag is used to set the operation mode for the script. It can take one of three values; 'train', 'valid', or 'rollout'.
batch_size (Integer)
Batch size for training.
noise_std (Float)
Standard deviation of the noise when training.
data_path (String)
Specifies the directory path where the dataset is located.
The dataset is expected to be in a specific format (e.g., .npz files).
It should contain metadata.json.
If --mode is training, the directory should contain train.npz.
If --mode is testing (rollout), the directory should contain test.npz.
If --mode is valid, the directory should contain valid.npz.
model_path (String)
The directory path where the trained model checkpoints are saved during training or loaded from during validation/rollout.
output_path (String)
Defines the directory where the outputs (e.g., rollouts) are saved,
when the --mode is set to rollout.
This is particularly relevant in the rollout mode where the predictions of the model are stored.
output_filename (String)
Base filename to use when saving outputs during rollout.
Default is "rollout", and the output will be saved as rollout.pkl in output_path.
It is not intended to include the file extension.
model_file (String)
The filename of the model checkpoint to load for validation or rollout (e.g., model-10000.pt). It supports a special value "latest" to automatically select the newest checkpoint file. This flexibility facilitates the evaluation of models at different stages of training.
train_state_file (String)
Similar to model_file, but for loading the training state (e.g., optimizer state). It supports a special value "latest" to automatically select the newest checkpoint file. (e.g., training_state-10000.pt)
ntraining_steps (Integer)
The total number of training steps to execute before stopping.
nsave_steps (Integer)
Interval at which the model and training state are saved.
lr_init (Float)
Initial learning rate.
lr_decay (Float)
How much the learning rate should decay over time.
lr_decay_steps (Integer)
Steps at which learning rate should decay.
cuda_device_number (Integer)
Base CUDA device (zero indexed). Default is None so default CUDA device will be used.
n_gpus (Integer)
Number of GPUs to use for training.
`train.py` in MeshNet (mesh-based domain)
mode (String)
This flag is used to set the operation mode for the script. It can take one of three values; 'train', 'valid', or 'rollout'.
batch_size (Integer)
Batch size for training.
data_path (String)
Specifies the directory path where the dataset is located.
The dataset is expected to be in a specific format (e.g., .npz files).
If --mode is training, the directory should contain train.npz.
If --mode is testing (rollout), the directory should contain test.npz.
If --mode is valid, the directory should contain valid.npz.
model_path (String)
The directory path where the trained model checkpoints are saved during training or loaded from during validation/rollout.
output_path (String)
Defines the directory where the outputs (e.g., rollouts) are saved,
when the --mode is set to rollout.
This is particularly relevant in the rollout mode where the predictions of the model are stored.
model_file (String)
The filename of the model checkpoint to load for validation or rollout (e.g., model-10000.pt). It supports a special value "latest" to automatically select the newest checkpoint file. This flexibility facilitates the evaluation of models at different stages of training.
train_state_file (String)
Similar to model_file, but for loading the training state (e.g., optimizer state). It supports a special value "latest" to automatically select the newest checkpoint file. (e.g., training_state-10000.pt)
cuda_device_number (Integer)
Allows specifying a particular CUDA device for training or evaluation, enabling the use of specific GPUs in multi-GPU setups.
rollout_filename (String)
Base name for saving rollout files. The actual filenames will append an index to this base name.
ntraining_steps (Integer)
The total number of training steps to execute before stopping.
nsave_steps (Integer)
Interval at which the model and training state are saved.
We use the numpy .npz format for storing positional data for GNS training. The .npz format includes a list of tuples of arbitrary length where each tuple corresponds to a differenet training trajectory and is of the form (position, particle_type). The data loader provides INPUT_SEQUENCE_LENGTH positions, set equal to six by default, to provide the GNS with the last INPUT_SEQUENCE_LENGTH minus one positions as input to predict the position at the next time step. The position is a 3-D tensor of shape (n_time_steps, n_particles, n_dimensions) and particle_type is a 1-D tensor of shape (n_particles).
The dataset contains:
- Metadata file with dataset information
(sequence length, dimensionality, box bounds, default connectivity radius, statistics for normalization, ...):
{
"bounds": [[0.1, 0.9], [0.1, 0.9]],
"sequence_length": 320,
"default_connectivity_radius": 0.015,
"dim": 2,
"dt": 0.0025,
"vel_mean": [5.123277536458455e-06, -0.0009965205918140803],
"vel_std": [0.0021978993231675805, 0.0026653552458701774],
"acc_mean": [5.237611158734309e-07, 2.3633027988858656e-07],
"acc_std": [0.0002582944917306106, 0.00029554531667679154]
}
- npz containing data for all trajectories
(particle types, positions, global context, ...):
Training datasets for Sand, SandRamps, and WaterDropSample are available on DesignSafe Data Depot [@vantassel2022gnsdata].
We provide the following datasets:
WaterDropSample(smallest dataset)SandSandRamps
Download the dataset DesignSafe DataDepot. If you are using this dataset please cite Vantassel and Kumar., 2022
We also use the numpy .npz format for storing data for training meshnet GNS.
The dataset contains:
- npz containing python dictionary describing mesh data and relevant dynamics at mesh nodes for all trajectories. The dictionary includes
{pos: (ntimestep, nnodes, ndims), node_type: (ntimestep, nnodes, ntypes), velocity: (ntimestep, nnodes, ndims), pressure: (ntimestep, nnodes, 1), cells: (ntimestep, ncells, 3)}
The dataset is shared on DesignSafe DataDepot. If you are using this dataset please cite Kumar and Choi., 2023
GNS uses pytorch geometric and CUDA. These packages have specific requirements, please see [PyG installation]((https://pytorch-geometric.readthedocs.io/en/latest/notes/installation.html) for details.
CPU-only installation on Linux
conda install -y pytorch torchvision torchaudio cpuonly -c pytorch
conda install -y pyg -c pyg
conda install -y pytorch-cluster -c pyg
conda install -y absl-py -c anaconda
conda install -y numpy dm-tree matplotlib-base pyevtk -c conda-forge You can use the WaterDropletSample dataset to check if your gns code is working correctly.
To test the code you can run:
pytest test/
To test on the small waterdroplet sample:
git clone https://github.com/geoelements/gns-sample
TMP_DIR="./gns-sample"
DATASET_NAME="WaterDropSample"
mkdir -p ${TMP_DIR}/${DATASET_NAME}/models/
mkdir -p ${TMP_DIR}/${DATASET_NAME}/rollout/
DATA_PATH="${TMP_DIR}/${DATASET_NAME}/dataset/"
MODEL_PATH="${TMP_DIR}/${DATASET_NAME}/models/"
ROLLOUT_PATH="${TMP_DIR}/${DATASET_NAME}/rollout/"
python -m gns.train --data_path=${DATA_PATH} --model_path=${MODEL_PATH} --ntraining_steps=10
- to setup a virtualenv
sh ./build_venv.sh- check tests run sucessfully.
- start your environment
source start_venv.sh pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
pip3 install torch_geometric
pip3 install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.3.0+cpu.html
pip3 install -r requirements.txtGNS can be trained in parallel on multiple nodes with multiple GPUs.
GNS scaling results on TACC Frontera GPU nodes with RTX-5000 GPUs.
GNS scaling result on TACC lonestar6 GPU nodes with A100 GPUs.
python -m torch.distributed.launch --nnodes=1 --nproc_per_node=[GPU_PER_NODE] --node_rank=[LOCAL_RANK] --master_addr=[MAIN_RANK] gns/train_multinode.py [ARGS] On each node, run
python -m torch.distributed.launch --nnodes=[NNODES] --nproc_per_node=[GPU_PER_NODE] --node_rank=[LOCAL_RANK] --master_addr=[MAIN_RANK ]gns/train_multinode.py [ARGS] PyTorch version of Graph Network Simulator and Mesh Graph Network Simulator are based on:
- https://arxiv.org/abs/2002.09405 and https://github.com/deepmind/deepmind-research/tree/master/learning_to_simulate
- https://arxiv.org/abs/2010.03409 and https://github.com/deepmind/deepmind-research/tree/master/meshgraphnets
- https://github.com/echowve/meshGraphNets_pytorch
This code is based upon work supported by the National Science Foundation under Grant OAC-2103937.
Kumar, K., & Vantassel, J. (2023). GNS: A generalizable Graph Neural Network-based simulator for particulate and fluid modeling. Journal of Open Source Software, 8(88), 5025. https://doi.org/10.21105/joss.05025
- Vantassel, Joseph; Kumar, Krishna (2022) “Graph Network Simulator Datasets.” DesignSafe-CI. https://doi.org/10.17603/ds2-0phb-dg64 v1
- Kumar, K., Y. Choi. (2023) "Cylinder flow with graph neural network-based simulator." DesignSafe-CI. https://doi.org/10.17603/ds2-fzg7-1719
