swag-train-bundle / README.md
DancingNow's picture
Add files using upload-large-folder tool
923fdff verified
|
Raw History Blame Contribute Delete
5.06 kB
# 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:
```python
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:
```text
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/`:
```bash
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:
```bash
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
bash Train/run_download_prepare.sh
```
Outputs are written to:
```text
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:
```bash
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
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
bash Train/run_all_local.sh --num-gpus 8
```
The training entrypoint itself is independent of a scheduler:
```bash
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
bash Train/run_generate_100.sh
```
This also defaults to 8 GPUs. To select a different count:
```bash
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`.