---
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.