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:
+39
-4
@@ -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())
|
||||
|
||||
Reference in New Issue
Block a user