Batch predictions when test set is large #125
Labels
enhancement
New feature or request
Comments
LeoGrin
added
enhancement
|
I agree it would be nice to support batching the predictions, but I removed this feature during the refactoring because the user should handle it. Batching the test predictions highly depends on the deployment hardware, input data, and inference requirements. In other words, finding a catch-all solution to determining the size of each batch that perfectly trades-off time with memory is tough. IMO, we could instead add an example for this to tabpfn-extensions and avoid supporting it in the interface here. Moreover, I am not aware of a standard from scikit-learn about this. Here is the very unsophisticated code I used before. The problem are how to set test_batch_threshold = 20_000
batch_predict_call = X.shape[0] > test_batch_threshold
if batch_predict_call:
warnings.warn(
f"Using batch_predict_call! This is a workaround to avoid VRAM memory issues that splits predictions "
f"with more than {test_batch_threshold} test instances into three predict-wise batches. "
stacklevel=2,
)
split_point_a = len(X) // 3
split_point_b = split_point_a * 2
input_batches = [
(
X[:split_point_a],
None if additional_y is None else additional_y[:split_point_a],
),
(
X[split_point_a:split_point_b],
None
if additional_y is None
else additional_y[split_point_a:split_point_b],
),
(
X[split_point_b:],
None if additional_y is None else additional_y[split_point_b:],
),
]
# Handle predictions per batch
for batch_X, batch_additional_y in input_batches:
# Predict here...
output_per_batch.append(
(
batch_prediction,
batch_additional_outputs,
)
)
if not batch_predict_call:
prediction, additional_outputs = output_per_batch[0]
else:
prediction = np.vstack([batch[0] for batch in output_per_batch])
additional_outputs = {}
for _, _batch_additional_outputs in output_per_batch:
for k in _batch_additional_outputs:
if k not in additional_outputs:
additional_outputs[k] = _batch_additional_outputs[k]
else:
additional_outputs[k] = torch.cat(
[additional_outputs[k], _batch_additional_outputs[k]], dim=1
) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Copying a discussion from discord:
Seems that it makes sense to batch when the test set is too big, and maybe what too big means depend on different things including whether we're using flash attention and memory-efficient SDPA.
The text was updated successfully, but these errors were encountered: