From 8f0d7c66c19bfa84c738ca666ad173eb85260a04 Mon Sep 17 00:00:00 2001 From: Danny McCormick Date: Fri, 6 Nov 2026 18:51:15 +0000 Subject: [PATCH] Optimize RunInference PyTorch/HF get_num_bytes and Sklearn Pandas batching --- .../apache_beam/ml/inference/huggingface_inference.py | 10 +++++++--- .../apache_beam/ml/inference/pytorch_inference.py | 5 +++-- .../apache_beam/ml/inference/sklearn_inference.py | 2 ++ 3 files changed, 12 insertions(+), 5 deletions(-) diff --git a/sdks/python/apache_beam/ml/inference/huggingface_inference.py b/sdks/python/apache_beam/ml/inference/huggingface_inference.py index 87c9f203ba69..8ad2168e87e4 100644 --- a/sdks/python/apache_beam/ml/inference/huggingface_inference.py +++ b/sdks/python/apache_beam/ml/inference/huggingface_inference.py @@ -404,7 +404,8 @@ def get_num_bytes( return sum(sys.getsizeof(element) for element in batch) else: return sum( - (el.element_size() for tensor in batch for el in tensor.values())) + el.element_size() * el.nelement() for tensor in batch + for el in tensor.values()) def get_metrics_namespace(self) -> str: """ @@ -607,9 +608,12 @@ def get_num_bytes( """ if self._framework == "tf": return sum(sys.getsizeof(element) for element in batch) - else: + elif batch and isinstance(batch[0], dict): return sum( - (el.element_size() for tensor in batch for el in tensor.values())) + el.element_size() * el.nelement() for tensor in batch + for el in tensor.values()) + else: + return sum(tensor.element_size() * tensor.nelement() for tensor in batch) def get_metrics_namespace(self) -> str: """ diff --git a/sdks/python/apache_beam/ml/inference/pytorch_inference.py b/sdks/python/apache_beam/ml/inference/pytorch_inference.py index 81234cb5c4f8..669797e7ad4d 100644 --- a/sdks/python/apache_beam/ml/inference/pytorch_inference.py +++ b/sdks/python/apache_beam/ml/inference/pytorch_inference.py @@ -360,7 +360,7 @@ def get_num_bytes(self, batch: Sequence[torch.Tensor]) -> int: Returns: The number of bytes of data for a batch of Tensors. """ - return sum((el.element_size() for tensor in batch for el in tensor)) + return sum(tensor.element_size() * tensor.nelement() for tensor in batch) def get_metrics_namespace(self) -> str: """ @@ -609,7 +609,8 @@ def get_num_bytes(self, batch: Sequence[torch.Tensor]) -> int: """ # If elements in `batch` are provided as a dictionaries from key to Tensors return sum( - (el.element_size() for tensor in batch for el in tensor.values())) + el.element_size() * el.nelement() for tensor in batch + for el in tensor.values()) def get_metrics_namespace(self) -> str: """ diff --git a/sdks/python/apache_beam/ml/inference/sklearn_inference.py b/sdks/python/apache_beam/ml/inference/sklearn_inference.py index 65220689283e..e88ad437b1a5 100644 --- a/sdks/python/apache_beam/ml/inference/sklearn_inference.py +++ b/sdks/python/apache_beam/ml/inference/sklearn_inference.py @@ -212,6 +212,8 @@ def _default_pandas_inference_fn( # vectorize data for better performance vectorized_batch = pandas.concat(batch, axis=0) predictions = model.predict(vectorized_batch) + if len(batch) == vectorized_batch.shape[0]: + return predictions, batch splits = [ vectorized_batch.iloc[[i]] for i in range(vectorized_batch.shape[0]) ]