Performance and caching

A session pays two fixed costs before its first prediction, the download of the weights from Hugging Face and the JAX compilation of the network. The compilation is paid again for each new batch size. Prediction itself is far cheaper than either.

Setup cost

preload_models loads the three estiMINT gradient-boosted models and both stateMINT checkpoints, downloading them if they are not cached.

import time
from estimint import preload_models

t0 = time.perf_counter()
eir_models, emulator_models = preload_models()
print(f"preload_models: {time.perf_counter() - t0:.2f} s")
print("estiMINT :", list(eir_models))
print("stateMINT:", list(emulator_models))
preload_models: 7.37 s
estiMINT : ['prevalence', 'hbr', 'eir_to_hbr']
stateMINT: ['prevalence', 'cases']

Before the weights are cached, most of that time is the download, whose size is given in Installation. Once they are cached, it is deserialisation from disk. preload_models tracks main rather than a pinned tag, so load the checkpoints yourself with from_pretrained(..., revision=...) when the exact weights matter. Call preload_models once at the top of a script.

The cache

Weights land in the standard Hugging Face cache.

from huggingface_hub import constants

print(constants.HF_HUB_CACHE)
/home/cosmo/.cache/huggingface/hub

Two environment variables control the cache. Both must be set before the first import.

Variable Effect
HF_HOME Relocates the whole cache. Point it at a shared or scratch filesystem on a cluster, where a home directory quota will not hold 80 MB of weights.
HF_HUB_OFFLINE=1 Forbids network access. Loads succeed from the cache and fail loudly otherwise. Set it on a compute node with no outbound route, and pre-populate the cache from the login node. 
export HF_HOME=/scratch/$USER/hf
export HF_HUB_OFFLINE=1

Batching

.predict is a single compiled forward pass over the whole batch. JAX compiles the network on first use and recompiles it whenever the batch size changes.

import numpy as np
from stateMINT.model import Mamba2Regressor

prevalence = Mamba2Regressor.from_pretrained(
    "dide-ic/stateMINT", predictor="prevalence",
)

district = {
    "eir": 31.96, "dn0_use": 0.311, "dn0_future": 0.485,
    "Q0": 0.85, "phi_bednets": 0.80, "seasonal": 0,
    "routine": 0.0, "itn_use": 0.70, "irs_use": 0.0,
    "itn_future": 0.70, "irs_future": 0.0, "lsm": 0.0,
}

batch = [dict(district, eir=float(e)) for e in np.linspace(5, 60, 50)]
len(batch)
50

The first call pays the compilation cost. The second does not.

t0 = time.perf_counter()
prevalence.predict(batch)
first = time.perf_counter() - t0

t0 = time.perf_counter()
prevalence.predict(batch)
second = time.perf_counter() - t0

print(f"first call  (compile + run): {first * 1000:7.0f} ms")
print(f"second call (run only)     : {second * 1000:7.0f} ms")
first call  (compile + run):     857 ms
second call (run only)     :     398 ms

The same fifty scenarios then run as fifty separate calls, each compiled for a batch of one.

prevalence.predict([district])          # pay the batch-of-one compilation

t0 = time.perf_counter()
for scenario in batch:
    prevalence.predict([scenario])
looped = time.perf_counter() - t0

print(f"one call, 50 scenarios : {second * 1000:7.0f} ms   ({second / 50 * 1000:.1f} ms each)")
print(f"50 calls, 1 scenario   : {looped * 1000:7.0f} ms   ({looped / 50 * 1000:.1f} ms each)")
print(f"speed-up               : {looped / second:.1f}x")
one call, 50 scenarios :     398 ms   (8.0 ms each)
50 calls, 1 scenario   :     975 ms   (19.5 ms each)
speed-up               : 2.4x

A call enters the compiled forward pass once, whatever the batch size. The loop pays it fifty times over.

A list of a few hundred rows is still a single call, and still under a second. Build the whole list first and then call .predict once. The timings came off whatever machine rendered the cells, so trust the ratio rather than the milliseconds.

CPU and GPU

The timings above are CPU timings, and a CPU is adequate at the scale of a list of scenarios, since the emulator is small. It is two Mamba2 layers with d_model 256. A GPU runs the same forward pass faster still, with the measured figures, the comparison against a full malariasimulation run, and what the speed costs in accuracy set out in Accuracy and speed.

For a sweep of thousands of scenarios, a CUDA 12 build of JAX is available as an extra.

pip install "mintstate[gpu]"

Nothing in the API changes when a GPU is present. JAX places the computation on the accelerator if it finds one, and without a CUDA-enabled JAX it falls back to CPU and warns, which is the only sign the GPU is not in use.

See also

A list of scenarios runs several campaigns against one baseline and settings in one call. stateMINT API documents the inference surface argument by argument.