damccorm opened a new pull request, #41058:
URL: https://github.com/apache/beam/pull/41058
### Description
Optimizes two hot paths in `apache_beam.ml.inference`:
1. **`PytorchModelHandlerTensor`, `PytorchModelHandlerKeyedTensor`,
`HuggingFaceModelHandlerTensor`, and `HuggingFaceModelHandlerKeyedTensor`
(`get_num_bytes`):**
- Previously, `get_num_bytes` iterated over tensor elements in Python
(`sum(el.element_size() for tensor in batch for el in tensor)`), which both
iterated over the first dimension in Python and undercounted multi-dimensional
tensors (summing `element_size()` per row rather than total elements).
- Replacing with `tensor.element_size() * tensor.nelement()` computes the
exact tensor byte footprint in $O(1)$ per tensor without iterating tensor
slices in Python (**up to 2,028x faster** on `256x256` `float32` tensor
batches: `4.17s` -> `0.0021s` for 1,000 tensors).
2. **`_default_pandas_inference_fn` in `sklearn_inference.py`:**
- When every input `DataFrame` in `batch` is a single row (`len(batch) ==
vectorized_batch.shape[0]`), returns `(predictions, batch)` directly instead of
slicing `vectorized_batch.iloc[[i]]` $N$ times (**35.65x faster** on single-row
DataFrame batches).
### Testing
- Verified with `pytest
sdks/python/apache_beam/ml/inference/pytorch_inference_test.py
sdks/python/apache_beam/ml/inference/sklearn_inference_test.py
sdks/python/apache_beam/ml/inference/huggingface_inference_test.py` (`64
passed`).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
