Platform for running and training AI models for Seismology

The EPOS AI Platform is a versatile solution that simplifies the technological barriers for utilizing and creating Machine Learning applications in the field of Seismology. It streamlines the process of applying artificial intelligence algorithms in research projects by enabling the execution of pre-trained models. Additionally, it provides the capability to retrain these models with custom, including private, datasets. Researchers can conduct training experiments and easily compare them with prior results, facilitating the optimization of methods to identify the most effective model. Furthermore, the platform supports the development of new models, and extensions to existing ones.

The EPOS AI Platform facilitates the sharing of results with close collaborators. The libraries utilized to construct the tools and create example solutions are de facto standards within the research community. This makes it possible not only to share the outcomes but also the models themselves.

Running AI models

On Episodes Platform

TODO write a simple guide

On user local machines or HPC clusters

The AI tools are available as an ordinary Python packages. We recommend for installation to use Anaconda or Mambaforge, but it should be possible to use any other Python installations or distributions.

Mambaforge installation

Mambaforge is a Python distribution based on Conda. This is a preferred way to run scripts and notebooks distributed on our Platform. It is possible to install and run on other python distributions, but we provide support only for Mambaforge/Anaconda.

The installation starts with downloading the Mambaforge binary from the official project site for your platform. Then follow the instructions in the official guide.

We prepared a Conda environment with all the AI tools installed. To create the environment please:

Download the epos-ai-tools.yml from our repository.
If necessary activate the Mambaforge environment.
Run the installation
```
mamba env create -f epos-ai-tools.yml
```
Then to activate the environment it is necessary to run for each new shell session:
```
conda activate epos-ai-tools
```
To check if you have enabled the environment correctly you should see the name of the environment in the shell prompt.

Available Tools

The installation comes with the following applications:

gpd_tool - an application for automatic detection of the first arrival time of the P and S waves based on Generalized Seismic Phase Detection with Deep Learning by Ross, Z. E., Meier, M.-A., Hauksson, E., and T. H. Heaton (2018)
phasenet_tool - an application for automatic detection of the first arrival time of the P and S waves based on PhaseNet: a deep-neural-network-based seismic arrival-time picking method by Zhu, W., & Beroza, G. C. (2019)

and official packages such as:

Seisbench - an open-source python toolbox for machine learning in seismology,
ObsPy - an open-source project dedicated to provide a Python framework for processing seismological data,
Jupyter Notebook and JupyterLab - a web services for interactive computing,
IPython - a powerful interactive shell,
PyTorch - an optimized tensor library for deep learning using GPUs and CPUs

Tools for picking P and S waves

TODO

Training AI models

Our Platform supports the user with training their own weights for predefined models and even to create their own models. In the first case this could help to improve the quality of the model by training in on data coming from a particular seismic networks (see e.g. Johnson et al. 2021).

To start with training we recommend to use the our HPC clusters which provide enough computing power and access to GPU. The prerequisites are following:

Active account on PLGrid Portal. To create an account see the official guide available in English and Polish.
Computing grant on PLGrid Infrastructure. The resources needed for training are highly depended on the size of a training dataset. For a dataset of about 10⁴ samples a single training session lasts about 10 minutes on a single GPU for models like GPD or PhaseNet. Please keep in mind that fine tuning of the model hyper-parameters usually requires tu run training sessions multiple times. Therefore applying for a grant, please use this numbers only as a starting point.
We recommend to apply for resources on a cluster Athena or Ares, but it is possible to run the training on other clusters as well. Please apply for the CPU and GPU computing time, as the GPU significantly accelerates the training process.
Mambaforge installation. Please follow the above instruction.
Environment for training AI models. Please install the epos-ai-train environment, by downloading epos-ai-train.yml from our repository and running:
```
mamba env create -f epos-ai-train.yml
```
Activate the environment it. is necessary to run for each new shell session:
```
conda activate epos-ai-train
```
To check if you have enabled the environment correctly you should see the name of the environment in the shell prompt.

Running jobs on an HPC cluster with SLURM

To submit a job on an HPC cluster use the following job script scheme:

#!/bin/bash

## Set you JOB_NAME to to make it easier to see in the job queue
#SBATCH --job-name=JOB_NAME

## Max task execution time (format is HH:MM:SS)
#SBATCH --time=00:15:00

## Name of grant to which resource usage will be charged
#SBATCH --account=GRANT_ID

## Name of partition
#SBATCH --partition=plgrid

## Number of allocated nodes
#SBATCH --nodes=1

## Number of tasks per node (by default this corresponds to the number of cores allocated per node)
#SBATCH --cpus-per-task=1

## Change to sbatch working directory
cd $SLURM_SUBMIT_DIR

## Activate Python environment
source PATH_TO_MAMBAFORGE_INSTALLATION/mambaforge/bin/activate
conda activate epos-ai-train

## Your usual invoke method e.g.
python train.py

To enable GPU access for the job please add to the top:

#SBATCH --gres=gpu:1

where the number :1 states the number of GPU devices, can be changed to higher numbers. To find which partitions provide access to GPU, please run:

sinfo -o '%P || %N || %G' | column -t

if in the last column is not null then, it indicates that the partition have access to GPU devices. Please setup the correct partition in the job script.

To find more about submitting jobs please see in corresponding documentation for Ares or Athena.

Available Tools

Seisbench - an open-source python toolbox for machine learning in seismology,
Weight & Biases SDK - a library to track experiments, version and iterate on datasets, evaluate model performance, reproduce models, and manage your ML workflows end-to-end,
ObsPy - an open-source project dedicated to provide a Python framework for processing seismological data,
Jupyter Notebook and JupyterLab - a web services for interactive computing,
IPython - a powerful interactive shell,
PyTorch - an optimized tensor library for deep learning using GPUs and CPUs,
CUDA - provides a development environment for creating high performance GPU-accelerated applications, a backend library for PyTorch

Weights & Biases

Weights & Biases is the machine learning platform for developers to build better models faster. Use W&B's lightweight, interoperable tools to quickly track experiments, version and iterate on datasets, evaluate model performance, reproduce models, visualize results and spot regressions, and share findings with colleagues. The main features are:

Track, compare, and visualize your ML models,
Scalable, customizable hyperparameter search and optimization,
Visualize and compare every experiment,
Quickly find and re-run previous model checkpoints,
Model registry and lifecycle management,
Lightweight dataset and model versioning,
Collaborative Reporting.

The Weights & Biases consists of two components: web-based platform on https://epos-ai.grid.cyfronet.pl and a public Python package wandb. The wandb Python library API calls serve to integrate an ML application, facilitating model performance tracking and sending metrics to a central platform. With web interface of this platform, one can then visualize and compare the model results.

Acknowledgements

The development of the EPOS AI platform and its integration with the EPISODES Platform was partially funded by the EPOS-PL+ project (No POIR.04.02.00-00-C005/19-00), co-financed by the European Union from the funds of the European Regional Development Fund (ERDF).

We want to express our appreciation for the excellent efforts of the Seisbench developers, and we strongly advise citing their publication:

Woollam, J., Münchmeyer, J., Tilmann, F., Rietbrock, A., Lange, D., Bornstein, T., Diehl, T., Giuchi, C., Haslinger, F., Jozinović, D., Michelini, A., Saul, J., & Soto, H. (2022). SeisBench - A Toolbox for Machine Learning in Seismology. in Seismological Research Letters https://doi.org/10.1785/0220210324

We express our gratitude to Weights & Biases for generously granting us an academic license of their web platform service.