damccorm commented on code in PR #37186:
URL: https://github.com/apache/beam/pull/37186#discussion_r3596499218
##########
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:
Ideally we just wouldn't need this for any of the examples. The downside
here is that it adds complexity and it doesn't feed into the batching done by
BatchElements (which seems like a better place for this)
--
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]
