Skip to content

Commit

Permalink
Use mini-batches when evaluating candidates within optimize_acqf_disc…
Browse files Browse the repository at this point in the history
…rete_local_search

Summary: If `X_loc` is large (which happens if there are many potential neighbors) this could lead to OOM. Breaking evaluations into mini batches should help.

Reviewed By: dme65

Differential Revision: D68354268
  • Loading branch information
saitcakmak authored and facebook-github-bot committed Jan 17, 2025
1 parent 7c5c7cb commit d08426f
Showing 1 changed file with 5 additions and 1 deletion.
6 changes: 5 additions & 1 deletion botorch/optim/optimize.py
Original file line number Diff line number Diff line change
Expand Up @@ -1466,7 +1466,11 @@ def optimize_acqf_discrete_local_search(
if len(X_loc) == 0:
break
with torch.no_grad():
acqval_loc = acq_function(X_loc.unsqueeze(1))
acqval_loc = _split_batch_eval_acqf(
acq_function=acq_function,
X=X_loc.unsqueeze(1),
max_batch_size=max_batch_size,
)
# break if no neighbor is better than the current point (local optimum)
if acqval_loc.max() <= curr_acqval:
break
Expand Down

0 comments on commit d08426f

Please sign in to comment.