Amar3tto commented on code in PR #37186:
URL: https://github.com/apache/beam/pull/37186#discussion_r3595412612
##########
sdks/python/apache_beam/examples/inference/pytorch_imagenet_rightfit.py:
##########
@@ -332,6 +337,43 @@ def pick_batch_size(arg: str) -> Optional[int]:
return None
+class RightFittingPytorchModelHandlerTensor(PytorchModelHandlerTensor):
+ def __init__(self, batch_sizes_to_try, image_size, *args, **kwargs):
+ self._batch_sizes_to_try = batch_sizes_to_try
+ self._rightfit_image_size = image_size
+ super().__init__(*args, **kwargs)
+
+ def load_model(self):
+ model = super().load_model()
+ last_err = None
+
+ for bs in self._batch_sizes_to_try:
+ try:
+ model_device = next(model.parameters()).device
+ dummy = torch.zeros(
+ (bs, 3, self._rightfit_image_size, self._rightfit_image_size),
+ dtype=torch.float32,
+ device=model_device)
+
+ with torch.no_grad():
+ model(dummy)
+
+ self._batch_size = bs
+ self._inference_batch_size = bs
Review Comment:
When the user provides an explicit batch size, it is set up front through
`inference_batch_size`. The runtime selection is only used for auto, because
the submission client cannot know which batch size fits on the actual Dataflow
worker GPU. The selection itself is deterministic: it tries the configured
sizes in order and picks the first one that fits.
--
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]
