Add --batch-size auto to predict, matching train

Estimates batch size from free GPU memory using the checkpoint's
hidden_dim/n_blocks, same as the train command.

Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This commit is contained in:
2026-06-22 08:23:58 +02:00
parent 46068f356f
commit 92d38cbed4
+39 -4
View File
@@ -255,8 +255,14 @@ def predict(
),
] = None,
batch_size: Annotated[
int, typer.Option("--batch-size", "-b", help="Inference batch size")
] = 4096,
str,
typer.Option(
"--batch-size",
"-b",
help="Inference batch size, or 'auto' to estimate from free GPU "
"memory (cuda devices only)",
),
] = "4096",
steps: Annotated[
int, typer.Option("--steps", "-s", help="Flow matching ODE steps")
] = 10,
@@ -266,6 +272,21 @@ def predict(
] = None,
) -> None:
"""Run trained model on a parquet file and save predictions."""
batch_size_auto = False
batch_size_value: Optional[int] = None
if batch_size.strip().lower() == "auto":
batch_size_auto = True
else:
try:
batch_size_value = int(batch_size)
except ValueError:
typer.echo(
f"error: --batch-size must be an integer or 'auto', "
f"got {batch_size!r}",
err=True,
)
raise typer.Exit(1)
_device = torch.device(device) if device else gconfig.auto_device()
typer.echo(f"device: {_device}")
@@ -279,6 +300,20 @@ def predict(
raise typer.Exit(1)
model_cfg = ckpt["model_config"]
if batch_size_auto:
try:
batch_size_value = gconfig.estimate_batch_size(
model_cfg["hidden_dim"], model_cfg["n_blocks"], _device
)
except ValueError as exc:
typer.echo(f"error: {exc}", err=True)
raise typer.Exit(1)
typer.echo(
f"batch_size: {batch_size_value} (auto-estimated from free GPU memory)"
)
assert batch_size_value is not None
bs = batch_size_value
pdg_map = {int(k): v for k, v in ckpt["pdg_map"].items()}
mat_map = {str(k): v for k, v in ckpt["mat_map"].items()}
cond_norm = Normalizer.from_dict(ckpt["normalizer"]["cond"])
@@ -323,8 +358,8 @@ def predict(
# Inference in batch_size slices
pred_parts = []
for start in range(0, N, batch_size):
end = min(start + batch_size, N)
for start in range(0, N, bs):
end = min(start + bs, N)
cc = torch.from_numpy(cond_cont[start:end]).float().to(_device)
ck = torch.from_numpy(cond_cat[start:end]).long().to(_device)
pred_parts.append(sample_flow(model, cc, ck, steps=steps).cpu().numpy())