DancingNow/swag-train-bundle

SWaG Empty-8 Training Bundle

Train/ is a portable download, training, and sampling bundle. It does not use absolute local paths, ~/.seisbench, or the SeisBench Python API. The default workflow downloads already processed data from the Hugging Face Dataset DancingNow/swag-processed-data; it does not reprocess anything. Checkpoints, metrics, and generated samples stay below Train/.

Model interface

The model is the no-high-resolution-skip SWaG backbone with eight reserved continuous condition slots:

prediction = model(noisy_waveforms, diffusion_timesteps, conditions)

conditions has shape [batch, 8]. During base training all eight values are zero. The condition projection is zero-initialized, preserving a clean base initialization while reserving the interface for later transfer learning.

Training data interface

The processed STEAD files use this waveform interface and component order:

data:   float32 [N, 6000, 3]  # ENZ, 100 Hz, 60 seconds
labels: float32 [N, 2]        # [P sample, S sample]

Each 6000-sample window is standardized to zero mean and unit standard deviation. No bandpass filter is applied. Only rows marked earthquake_local with valid P/S arrivals are written to the processed STEAD files; noise is never used for training.

The Iquique files are downloaded from the iquique/ directory of the Hugging Face dataset DancingNow/swag-processed-data. They were processed previously from the original SeisBench Iquique release and already use the same [N,6000,3] ENZ waveform interface and P/S labels. They are downloaded as-is and are not processed again by this bundle's default workflow.

Install

Run from the directory containing Train/:

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install -r Train/requirements.txt

The system must also provide the hf command and a Hugging Face login. Use a new virtual environment as shown; reusing an environment with incompatible preinstalled PyTorch packages can cause import errors unrelated to this bundle.

For a private dataset, log in before downloading:

hf auth login

For CUDA training, install the PyTorch build matching the machine's CUDA driver if the default pip build is not suitable.

Download Processed Data

Download the already processed STEAD and Iquique files:

bash Train/run_download_prepare.sh

Outputs are written to:

Train/data/processed/from_raw/stead/train/stead_100hz_60s_train.h5
Train/data/processed/from_raw/stead/test/stead_100hz_60s_test.h5
Train/data/processed/from_raw/iquique/

The command downloads from DancingNow/swag-processed-data into the exact directory expected by the training configuration. It performs no cropping, standardization, label filtering, or other conversion.

The repository is currently Private, so the target computer must be logged in to the DancingNow account (or otherwise have read access). The download is resumable through the Hugging Face cache; rerunning the command is safe.

To use another revision or destination:

HF_REVISION=main \
HF_LOCAL_DIR=/path/to/project/Train/data/processed/from_raw \
bash Train/run_download_prepare.sh

Training

The default runnable configuration uses 8 GPUs, a global batch size of 256 (32 samples per GPU), no gradient accumulation, 15 epochs, and saves checkpoints every 5 epochs:

bash Train/run_train_local.sh --num-gpus 8

The GPU count is a launcher parameter. For example, use --num-gpus 1, --num-gpus 4, or --num-gpus 8. It defaults to 8 and may also be set with NUM_GPUS. The global batch size must be divisible by the GPU count.

For a new machine, the complete download-to-training sequence can be started with:

bash Train/run_all_local.sh --num-gpus 8

The training entrypoint itself is independent of a scheduler:

python Train/training/train.py --config Train/configs/stead_empty8_local.yaml

The transfer-learning configuration in Train/Transfer_and_Test/configs/config.yaml also uses a global training batch size of 256. Generation uses a global batch size of 256 in both local and transfer-learning launchers. Training and generation distribute their work across the requested GPUs and only the main process writes the final checkpoint or merged HDF5 output.

If another machine cannot fit 256 waveforms in GPU memory, reduce the configured batch size before running. If distributed training is required, launch this same Python entrypoint with that environment's own distributed launcher or job scheduler. No scheduler-specific submission script is included.

Generate 100 samples

After epoch 15 finishes:

bash Train/run_generate_100.sh

This also defaults to 8 GPUs. To select a different count:

bash Train/run_generate_100.sh --num-gpus 4

This performs ancestral DDPM sampling with 25 steps and seed 2026. The output is Train/results/stead_empty8_local/generated_100/generated_ddpm25_seed2026.h5.