ML pipelines: RunInference - OSS Image Object detection, OSS Image Captioning, OSS Image Classification#37186
ML pipelines: RunInference - OSS Image Object detection, OSS Image Captioning, OSS Image Classification#37186Amar3tto wants to merge 34 commits into
Conversation
Summary of ChangesHello @Amar3tto, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances Apache Beam's machine learning capabilities by integrating a new PyTorch-based image object detection pipeline. The pipeline leverages the RunInference transform for efficient batched GPU inference with open-source TorchVision models, processing images from cloud storage and outputting structured detection results to BigQuery. This addition is complemented by a new performance benchmark and corresponding documentation, ensuring that the pipeline's efficiency and resource usage can be consistently monitored and evaluated. Highlights
🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console. Ignored Files
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
Amar3tto
changed the title
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #37186 +/- ##
=============================================
- Coverage 55.28% 36.33% -18.96%
Complexity 1676 1676
=============================================
Files 1067 1069 +2
Lines 167148 167178 +30
Branches 1208 1208
=============================================
- Hits 92415 60737 -31678
- Misses 72551 104259 +31708
Partials 2182 2182
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
Amar3tto
changed the title
|
Assigning reviewers: R: @claudevdm for label python. Note: If you would like to opt out of this review, comment Available commands:
The PR bot will only process comments in the main thread (not review comments). |
|
@Abacn Could you please help with review? |
|
/gemini review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces three new ML inference pipelines for image classification, object detection, and image captioning using PyTorch, along with their corresponding benchmarks and documentation. The pipelines are well-structured and showcase advanced Beam features like RunInference with custom model handlers and stateful DoFns. My review focuses on improving scalability, robustness, and maintainability. I've identified a few key areas for improvement, including a scalability bottleneck in the data loading pipelines, several instances of broad exception handling that could mask errors, some potentially buggy logic, and a few copy-paste errors in the new documentation pages. Overall, this is a valuable contribution, and the suggested changes aim to make these examples more robust and easier to understand.
|
Reminder, please take a look at this pr: @claudevdm @liferoad @shunping |
|
Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment R: @jrmccluskey for label python. Available commands:
|
|
Reminder, please take a look at this pr: @jrmccluskey @damccorm |
|
Assigning new set of reviewers because Pr has gone too long without review. If you would like to opt out of this review, comment R: @shunping for label python. Available commands:
|
|
waiting on author |
Amar3tto
force-pushed
the
oss-image-detection
branch
from
a7ba9db to
439b5aa
Compare
|
Could you please fix the formatting failures? Also, please avoid rebasing when possible to avoid breaking GitHub's review features |
|
/gemini review |
There was a problem hiding this comment.
Code Review
This pull request introduces several new PyTorch-based inference pipelines (image captioning, object detection, and image classification with adaptive batch sizing) along with their corresponding benchmarks, requirements, and performance documentation. The review feedback highlights critical bugs regarding shape mismatches during batching in the classification and object detection pipelines, a performance bottleneck from redundant image encoding in the CLIP model, an architectural flaw where model warmup runs on the submission client instead of the workers, and a robustness issue concerning guaranteed cleanup of Pub/Sub resources.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| def decode_and_preprocess(image_bytes: bytes, size: int = 224) -> torch.Tensor: | ||
| """Decode bytes->RGB PIL->resize/crop->tensor->normalize.""" | ||
| with PILImage.open(io.BytesIO(image_bytes)) as img: | ||
| img = img.convert("RGB") | ||
| img.thumbnail((256, 256)) | ||
| w, h = img.size | ||
| left = (w - size) // 2 | ||
| top = (h - size) // 2 | ||
| img = img.crop( | ||
| (max(0, left), max(0, top), min(w, left + size), min(h, top + size))) | ||
|
|
||
| # To tensor [0..1] | ||
| import numpy as np | ||
| mean = np.array(IMAGENET_MEAN, dtype=np.float32) | ||
| std = np.array(IMAGENET_STD, dtype=np.float32) | ||
| arr = np.asarray(img).astype("float32") / 255.0 # H,W,3 | ||
| # Normalize | ||
| arr = (arr - mean) / std | ||
| # HWC -> CHW | ||
| arr = np.transpose(arr, (2, 0, 1)).astype("float32") | ||
| return torch.from_numpy(arr).float() # float32, shape (3,224,224) |
There was a problem hiding this comment.
Critical Bug: Shape Mismatch during Batching
Using img.thumbnail((256, 256)) preserves the aspect ratio of the image. If the input image is not square, one of its dimensions will be smaller than 256. When you subsequently crop it using (w - size) // 2 where size = 224, the resulting cropped image will have a dimension smaller than 224 (e.g., 224 x 128).
Because different images in the dataset have different aspect ratios, they will yield tensors of different shapes (e.g., [3, 224, 128] and [3, 128, 224]). When RunInference tries to batch these tensors using torch.stack, it will crash with a shape mismatch error.
To fix this, resize the shorter side of the image to 256 while preserving the aspect ratio, and then perform a center crop to exactly size x size.
def decode_and_preprocess(image_bytes: bytes, size: int = 224) -> torch.Tensor:
"""Decode bytes->RGB PIL->resize/crop->tensor->normalize."""
with PILImage.open(io.BytesIO(image_bytes)) as img:
img = img.convert("RGB")
w, h = img.size
if w < h:
new_w = 256
new_h = int(h * (256 / w))
else:
new_h = 256
new_w = int(w * (256 / h))
img = img.resize((new_w, new_h), PILImage.BILINEAR)
left = (new_w - size) // 2
top = (new_h - size) // 2
img = img.crop((left, top, left + size, top + size))
# To tensor [0..1]
import numpy as np
mean = np.array(IMAGENET_MEAN, dtype=np.float32)
std = np.array(IMAGENET_STD, dtype=np.float32)
arr = np.asarray(img).astype("float32") / 255.0 # H,W,3
# Normalize
arr = (arr - mean) / std
# HWC -> CHW
arr = np.transpose(arr, (2, 0, 1)).astype("float32")
return torch.from_numpy(arr).float() # float32, shape (3,224,224)| def decode_to_tens( | ||
| image_bytes: bytes, | ||
| resize_shorter_side: Optional[int] = None) -> torch.Tensor: | ||
| """Decode bytes -> RGB PIL -> optional resize -> float tensor [0..1], CHW. | ||
|
|
||
| Note: TorchVision detection models apply their own normalization internally. | ||
| """ | ||
| with PILImage.open(io.BytesIO(image_bytes)) as img: | ||
| img = img.convert("RGB") | ||
|
|
||
| if resize_shorter_side and resize_shorter_side > 0: | ||
| w, h = img.size | ||
| # Resize so that shorter side == resize_shorter_side, keep aspect ratio. | ||
| if w < h: | ||
| new_w = resize_shorter_side | ||
| new_h = int(h * (resize_shorter_side / float(w))) | ||
| else: | ||
| new_h = resize_shorter_side | ||
| new_w = int(w * (resize_shorter_side / float(h))) | ||
| img = img.resize((new_w, new_h)) | ||
|
|
||
| import numpy as np | ||
| arr = np.asarray(img).astype("float32") / 255.0 # H,W,3 in [0..1] | ||
| arr = np.transpose(arr, (2, 0, 1)) # CHW | ||
| return torch.from_numpy(arr) |
There was a problem hiding this comment.
Critical Bug: Varying Image Shapes cause Stacking Crash
PytorchModelHandlerTensor internally calls torch.stack(batch) to create a batched tensor. If the images in the batch have different aspect ratios or sizes, torch.stack will raise a RuntimeError: stack expects each tensor to be equal size and crash the pipeline.
To make the pipeline robust for general datasets, resize all images to a fixed square size (e.g., 640x640 or 800x800) during preprocessing.
def decode_to_tens(
image_bytes: bytes,
size: Tuple[int, int] = (640, 640)) -> torch.Tensor:
"""Decode bytes -> RGB PIL -> resize to fixed size -> float tensor [0..1], CHW."""
with PILImage.open(io.BytesIO(image_bytes)) as img:
img = img.convert("RGB")
img = img.resize(size, PILImage.BILINEAR)
import numpy as np
arr = np.asarray(img).astype("float32") / 255.0 # H,W,3 in [0..1]
arr = np.transpose(arr, (2, 0, 1)) # CHW
return torch.from_numpy(arr)| inputs = processor( | ||
| text=texts, | ||
| images=images, | ||
| return_tensors="pt", | ||
| padding=True, | ||
| truncation=True, | ||
| ) | ||
| inputs = { | ||
| k: (v.to(self.device) if torch.is_tensor(v) else v) | ||
| for k, v in inputs.items() | ||
| } | ||
|
|
||
| # avoid NxN logits inside CLIPModel.forward() | ||
| img = model.get_image_features( | ||
| pixel_values=inputs["pixel_values"]) # [N, D] | ||
| txt = model.get_text_features( | ||
| input_ids=inputs["input_ids"], | ||
| attention_mask=inputs.get("attention_mask"), | ||
| ) # [N, D] | ||
|
|
||
| img = img / img.norm(dim=-1, keepdim=True) | ||
| txt = txt / txt.norm(dim=-1, keepdim=True) | ||
|
|
||
| logit_scale = model.logit_scale.exp() # scalar tensor | ||
| pair_scores = (img * txt).sum(dim=-1) * logit_scale # [N] | ||
| pair_scores_cpu = pair_scores.detach().cpu().tolist() |
There was a problem hiding this comment.
Performance Bottleneck: Redundant Image Encoding in CLIP
In the current implementation, the same image is duplicated num_captions (default: 5) times in the images list and passed to processor and model.get_image_features. Since the image encoder (ViT) is computationally heavy, encoding the exact same image multiple times per batch element is highly inefficient.
We can optimize this by encoding only the unique active images once, and then repeating/aligning their features to match the flat candidate text features before computing the cosine similarity.
with torch.no_grad():
# Extract unique images to avoid redundant encoding
unique_images = []
for x in batch:
if x.get("candidates"):
unique_images.append(x["image"])
image_inputs = processor(images=unique_images, return_tensors="pt")
image_inputs = {k: v.to(self.device) for k, v in image_inputs.items()}
text_inputs = processor(
text=texts,
return_tensors="pt",
padding=True,
truncation=True,
)
text_inputs = {k: v.to(self.device) for k, v in text_inputs.items()}
img_features = model.get_image_features(**image_inputs) # [B_active, D]
txt_features = model.get_text_features(**text_inputs) # [total_pairs, D]
img_features = img_features / img_features.norm(dim=-1, keepdim=True)
txt_features = txt_features / txt_features.norm(dim=-1, keepdim=True)
logit_scale = model.logit_scale.exp() # scalar tensor
# Align image features with text features
repeated_img_features = []
active_idx = 0
for start_i, end_i in offsets:
if start_i != end_i:
num_candidates = end_i - start_i
repeated_img_features.append(img_features[active_idx].repeat(num_candidates, 1))
active_idx += 1
if repeated_img_features:
repeated_img_features = torch.cat(repeated_img_features, dim=0)
pair_scores = (repeated_img_features * txt_features).sum(dim=-1) * logit_scale
pair_scores_cpu = pair_scores.detach().cpu().tolist()
else:
pair_scores_cpu = []| bs_ok = None | ||
| last_err = None | ||
| for bs in tried: | ||
| try: | ||
| model_handler = PytorchModelHandlerTensor( | ||
| model_class=lambda: create_timm_m(known_args.pretrained_model_name), | ||
| model_params={}, | ||
| state_dict_path=known_args.model_state_dict_path, | ||
| device=device, | ||
| inference_batch_size=bs | ||
| if bs is not None else 64, # start guess for warmup | ||
| ) | ||
| # quick warmup to validate memory (single dummy tensor) | ||
| dummy = torch.zeros((3, known_args.image_size, known_args.image_size), | ||
| dtype=torch.float32) | ||
| _ = model_handler.load_model() # ensures weights on device | ||
| with torch.no_grad(): | ||
| mdl = model_handler._model | ||
| mdl(torch.unsqueeze(dummy, 0)) | ||
| bs_ok = bs if bs is not None else 64 | ||
| break | ||
| except RuntimeError as e: | ||
| last_err = e | ||
| logging.warning("Batch size %s failed during warmup: %s", bs, e) | ||
| continue | ||
|
|
||
| if bs_ok is None: | ||
| logging.warning( | ||
| "Falling back to batch_size=8 due to previous errors: %s", last_err) | ||
| bs_ok = 8 | ||
| model_handler = PytorchModelHandlerTensor( | ||
| model_class=lambda: create_timm_m(known_args.pretrained_model_name), | ||
| model_params={}, | ||
| state_dict_path=known_args.model_state_dict_path, | ||
| device=device, | ||
| inference_batch_size=bs_ok, | ||
| ) | ||
|
|
There was a problem hiding this comment.
Architectural Flaw: Warmup / Right-fitting Executed on Submission Client
Currently, the warmup loop to determine the optimal batch size runs directly inside the run() function on the submission client (driver) before the pipeline is sent to the runner. This has two major issues:
- If the client machine does not have a GPU (which is typical for CI/CD runners or local submission environments), loading the model with
device='GPU'will fail immediately during pipeline submission. - The warmup selects a batch size based on the client's hardware instead of the actual worker's hardware (e.g., Tesla T4 on Google Cloud).
To fix this, subclass PytorchModelHandlerTensor and perform the warmup/right-fitting logic inside the load_model method, which executes directly on the workers.
class RightFitPytorchModelHandler(PytorchModelHandlerTensor):
def __init__(self, tried_batch_sizes, image_size, *args, **kwargs):
super().__init__(*args, **kwargs)
self.tried_batch_sizes = tried_batch_sizes
self.image_size = image_size
def load_model(self):
last_err = None
for bs in self.tried_batch_sizes:
try:
self._inference_batch_size = bs
model = super().load_model()
# Warmup to validate memory on the actual worker
dummy = torch.zeros((3, self.image_size, self.image_size), dtype=torch.float32).to(self.device)
with torch.no_grad():
model(torch.unsqueeze(dummy, 0))
return model
except RuntimeError as e:
last_err = e
logging.warning("Batch size %s failed during warmup: %s", bs, e)
continue
# Fallback
self._inference_batch_size = 8
return super().load_model()| result = pipeline.run() | ||
| result.wait_until_finish(duration=1800000) # 30 min | ||
| try: | ||
| result.cancel() | ||
| except Exception: | ||
| pass | ||
|
|
||
| if known_args.mode == 'streaming': | ||
| cleanup_pubsub_resources( | ||
| project=known_args.project, | ||
| topic_path=known_args.pubsub_topic, | ||
| subscription_path=known_args.pubsub_subscription) | ||
|
|
There was a problem hiding this comment.
Robustness: Ensure Pub/Sub Resources are Cleaned Up
If the pipeline fails or is cancelled due to a timeout, the cleanup code at the end of the run function might not be reached, leaving orphaned Pub/Sub topics and subscriptions. Wrap the pipeline execution and waiting in a try...finally block to guarantee cleanup.
| result = pipeline.run() | |
| result.wait_until_finish(duration=1800000) # 30 min | |
| try: | |
| result.cancel() | |
| except Exception: | |
| pass | |
| if known_args.mode == 'streaming': | |
| cleanup_pubsub_resources( | |
| project=known_args.project, | |
| topic_path=known_args.pubsub_topic, | |
| subscription_path=known_args.pubsub_subscription) | |
| try: | |
| result = pipeline.run() | |
| result.wait_until_finish(duration=1800000) # 30 min | |
| finally: | |
| try: | |
| result.cancel() | |
| except Exception: | |
| pass | |
| if known_args.mode == 'streaming': | |
| cleanup_pubsub_resources( | |
| project=known_args.project, | |
| topic_path=known_args.pubsub_topic, | |
| subscription_path=known_args.pubsub_subscription) |
| del inference_args | ||
| del model_id |
2 participants
Please add a meaningful description for your change here
Thank you for your contribution! Follow this checklist to help us incorporate your contribution quickly and easily:
addresses #123), if applicable. This will automatically add a link to the pull request in the issue. If you would like the issue to automatically close on merging the pull request, commentfixes #<ISSUE NUMBER>instead.CHANGES.mdwith noteworthy changes.See the Contributor Guide for more tips on how to make review process smoother.
To check the build health, please visit https://github.com/apache/beam/blob/master/.test-infra/BUILD_STATUS.md
GitHub Actions Tests Status (on master branch)
See CI.md for more information about GitHub Actions CI or the workflows README to see a list of phrases to trigger workflows.