Remove PIL and tensor round-trips from the inference image path#727
Open
maxwbuckley wants to merge 3 commits into
Open
Remove PIL and tensor round-trips from the inference image path#727maxwbuckley wants to merge 3 commits into
maxwbuckley wants to merge 3 commits into
Conversation
The per-control-step eval image path did: numpy frames -> per-frame
np.array copy -> albumentations -> per-frame torch.from_numpy().permute
-> stack views -> .numpy() -> per-frame CHW->HWC transpose -> PIL ->
Qwen3VL processor. The checkpoint ships the fast torchvision image
processor (Qwen2VLImageProcessorFast for the N1.7 / Cosmos-Reason2-2B
config), which accepts numpy HWC / torch CHW frames directly and
produces bitwise-identical outputs to PIL input, so the PIL conversion
was pure overhead: at inference the frames are now handed to it raw.
- _get_vlm_inputs gains an eval-only fast path (_eval_transform_frames)
that applies the identical albumentations/torchvision transform calls
per frame but drops the numpy copies, the HWC->CHW->HWC torch round
trip, the .numpy() and the numpy->PIL conversions. Gated on eval
mode, no masks, and no replay-style transform; the training /
augmentation path is byte-for-byte untouched (guarded by tests), as
is process_observation.
- New opt-in flag use_fast_image_processor (default None = checkpoint
default, no change): True/False force the fast (torchvision) / slow
(PIL) HF image processor. NOT bitwise: measured fast-vs-slow
divergence max-abs 2/255 (mean-abs ~1.5e-5) on 320x240 where a real
resize occurs, ~6e-8 on already-aligned 256x256. Documented as
requiring closed-loop validation before use.
Bitwise guarantee scope: for the default configuration (any
use_fast_image_processor=None checkpoint), the final model inputs
(pixel_values, image_grid_thw, input_ids, attention_mask) of the eval
path are torch.equal to the original implementation. Gated by
tests/gr00t/model/test_image_preprocessing_parity.py, which reimplements
the original path verbatim and asserts equality on {256x256, 320x240} x
{1, 3 cameras} x T=2 for both albumentations and torchvision flavors,
plus a direct numpy/torch-vs-PIL processor-input parity test.
Profile (tools/perf/profile_preprocessing.py, real N1.7 checkpoint
processor, CPU, WSL2, interleaved A/B in one process, 150 iters,
median / mean ms per step, before -> after):
1 cam x 2 frames @256x256: 2.07 / 2.92 -> 0.99 / 1.42 (Welch p=8.6e-09)
3 cam x 2 frames @256x256: 6.25 / 7.14 -> 2.47 / 2.98 (Welch p=8.8e-61)
1 cam x 2 frames @240x320: 3.16 / 4.31 -> 1.72 / 2.34 (Welch p=7.5e-06)
Sub-stages removed (3-cam medians): PIL conversion 2.6 ms, np/torch
conversions 0.18 ms; the HF image-processor call itself is also
faster on raw numpy frames than on PIL (median 1.29 -> 0.79 ms).
Run-to-run stdev decreases as well (e.g. 3-cam totals 1.9 -> 1.3 ms).
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Reduce the change to its essence: stop converting CHW uint8 frames to PIL before the Qwen3VL processor (which accepts tensors directly and produces bitwise-identical outputs) and drop the now-unneeded .numpy() conversion. 9 insertions / 9 deletions in the core file. Interleaved timing (150 iters, real processor): 1-cam 1.84 -> 0.69 ms, 3-cam 5.93 -> 1.76 ms per step — the full previously-measured win, at a fraction of the diff. The eval fast-path restructuring, the use_fast_image_processor override, and the profiling script were removed from this PR (they remain available on the development fork). Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Author
|
Slimmed this PR to the minimal deadweight removal (9+/9− in the core file): the initial push also included an eval-path restructuring, an image-processor override flag, and a profiling script, but interleaved measurement showed the PIL round-trip removal alone accounts for ~95% of the win — so everything else was dropped to keep the diff reviewable. The description above reflects the current state. |
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
Removes deadweight from the inference image path — 9 insertions / 9 deletions in one file:
_apply_vlm_processingconverted every CHW uint8 frame to PIL before handing it to the Qwen3VL processor. The processor (Qwen2VLImageProcessorFast, which the checkpoint config already resolves to) accepts CHW tensors directly and produces identical outputs, so the PIL conversion was a redundant per-frame encode/decode. Frames are now passed as tensors (torch.as_tensor— zero-copy for numpy inputs)..numpy()("processor expects numpy array" no longer holds).Bitwise guarantee
tests/gr00t/model/test_vlm_image_input_parity.pyreimplements the old PIL path verbatim as a reference oracle and assertstorch.equalon every processor output (pixel_values,image_grid_thw,input_ids,attention_mask) across {256×256, 240×320 non-square, 243×243 non-patch-multiple} × {2, 6 frames}, on the real Cosmos-Reason2-2B processor (configs only, no weights; skips when the gated repo is unreachable).Measured
Interleaved old/new timing on the real processor (150 iters, medians, per control step):
On an RTX 5090 end-to-end (
Gr00tPolicy.get_action, WSL2), this class of change took the CPU-preprocessing stage from 35.6 ms to ~5 ms (27% → ~4% of E2E), with RNG-seeded policy actions bitwise-identical to unmodifiedmain.Tests
HF_TOKENwith Cosmos-Reason2-2B access, matching existing suite convention; skips otherwise).tests/gr00t/model/ + tests/gr00t/policy/ -m "not gpu": no regressions vsmain.pre-commitclean.🤖 Generated with Claude Code