File size: 3,980 Bytes
9e41588 222ef4f 9e41588 222ef4f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 | ---
license: mit
pretty_name: PoreML V1 — reference model checkpoints
library_name: pytorch
datasets:
- PoreML/PoreML_data
tags:
- physics
- computational-fluid-dynamics
- porous-media
- multiphase-flow
- scientific-machine-learning
- neural-operator
- 3d
---
# PoreML V1 — reference model checkpoints
The trained weights behind the PoreML benchmark: five reference models — `unet3d`, `fno3d`,
`p3d`, `transolver`, `abupt` — on the next-frame forecasting task of
[`PoreML/PoreML_data`](https://huggingface.co/datasets/PoreML/PoreML_data), in two phases.
| Phase | Folder | Runs | Report this checkpoint |
|---|---|---:|---|
| base training, 20 effective epochs | `train/` | 35 | `ckpts/best.pt` — best one-step validation `mae@phi` |
| push-forward fine-tune of each base run | `train_push/` | 35 | `ckpts/best_rollout.pt` — best 64-step rollout `mae@phi` |
35 = 5 models × 7 tasks: the `gen` (generated media) and `all` (generated ∪ micro-CT) splits of
drainage, GDL and trapping, plus underfill's single `all` split.
## Download
With the PoreML code repository checked out and installed (`uv sync`):
```bash
uv run poreml checkpoints --dry-run # what is there and how big it is (7 GB)
uv run poreml checkpoints --phase train_push --which best_rollout # the weights the paper reports
uv run poreml checkpoints --campaign drainage --model unet --kind gen
uv run poreml checkpoints # everything
```
Filters are `--phase`, `--campaign`, `--model`, `--kind` and `--which`; they intersect and each
is repeatable. Files land under `case/`, which is where the push configs' `train.init_from`, the
transfer studies (`case/shift`, `case/scale`) and `util/inference` look a finished run up —
nothing has to be moved. Rerunning resumes. Without the package:
```bash
hf download PoreML/PoreML_checkpoint --local-dir case --exclude README.md checkpoints.csv
```
## Layout
The repository root is the benchmark's `case/` folder:
```
<phase>/<campaign>/<model>_<kind>/ckpts/<run>/
config.yaml the full config the run was trained with; a checkpoint is always scored with it
run_meta.json status, resolved data and split (with sha256), model size, provenance, progress
metrics.csv one row per epoch: losses, learning rate, every validation metric
train_log.csv the training loss every 100 steps
ckpts/best.pt best one-step validation metric
ckpts/last.pt the final epoch
ckpts/best_rollout.pt best rollout metric (push-forward runs)
checkpoints.csv every weight file: phase, campaign, model, kind, path, epoch, bytes, sha256
```
A `.pt` is a `torch.save`d dict: `model` (the state dict), `config` (the model's name, params
and precision), `epoch`, `metrics`. `run_meta.json` of a push-forward run records the path and
sha256 of the base `best.pt` it started from; it matches `checkpoints.csv`.
Not included: optimiser state, the per-epoch rollout candidates, and stored rollout frames —
`poreml inference`, `poreml metric` and `poreml render` reproduce those from the weights.
## Use
```bash
RUN=case/train_push/drainage/unet_gen/ckpts/<run>
uv run poreml rollout -c $RUN/config.yaml --ckpt $RUN/ckpts/best_rollout.pt # one scored 64-step rollout per validation run
uv run python case/shift_push/submit.py --dry-run # the transfer studies see the runs as finished
```
The runs were trained on single NVIDIA H200 GPUs with PyTorch 2.11 (CUDA 12.8), `tf32`
everywhere except AB-UPT's `bf16`. Paths inside the text files are relative to the code
repository's root; the node's hostname has been removed from `run_meta.json`, and nothing else
was changed.
## Licence
MIT for the weights of `unet3d`, `fno3d`, `p3d` and `transolver`. The `abupt` weights are
derived from the AB-UPT architecture and follow its upstream Emmi AI Non-Production License:
research and evaluation use only.
|