XiaoHongbo-Hope commented on code in PR #9852:
URL: https://github.com/apache/paimon/pull/9852#discussion_r4023226779
##########
paimon-python/pypaimon/multimodal/lerobot/dataset.py:
##########
@@ -438,19 +445,27 @@ def get_items(self, indices):
})
import torch
+ visual_windows = _stack_visual_windows(
+ plans, converted, [key for key in self._visual_keys
+ if key not in video_windows]
+ ) if plans[0]["windows"] else {}
+ visual_windows.update(video_windows)
Review Comment:
> [P2] Avoid materializing the entire batch of windows before image
transforms
>
> When `image_transforms` returns new storage, such as resizing frames to a
smaller resolution, these eager paths allocate every full-resolution window
before the first transform runs. `visual_windows` then keeps all those input
tensors alive until `get_items()` returns, even after the transformed outputs
have replaced them in each item. Previously, window assembly and transforms ran
one sample at a time.
>
> I reproduced this with 32 samples, two image cameras, 16-frame windows,
and a resize transform: the baseline kept at most 2 original window tensors
alive simultaneously, while this version kept all 64 alive through the last
transform. This significantly increases peak memory in DataLoader workers when
preprocessing reduces the frame size.
>
> Could we retain the per-sample path when `image_transforms` is configured,
or assemble windows in bounded groups? This should cover both
`_decode_video_windows` and `_stack_visual_windows`. Clearing list entries
after consumption alone would not eliminate the initial batch-wide allocation
peak.
Fixed by falling back to per-sample window assembly and transforms when
image_transforms is configured, for both image and video windows. I consider we
can set up a basic benchmark first. Optimizing the `image_transforms` path can
be handled in a follow-up PR.
--
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]
