Repository navigation
simple_trainer: non_blocking batch copies, fewer loss.item() calls, uint8 images from the loader - #1087
Open
kleinicke wants to merge 2 commits into
Open
simple_trainer: non_blocking batch copies, fewer loss.item() calls, uint8 images from the loader#1087kleinicke wants to merge 2 commits into
kleinicke wants to merge 2 commits into
Conversation
The batch comes from a pin_memory loader, so the host-to-device copies can be non_blocking instead of waiting for the previous step's GPU work. loss.item() (and the depth / pose error items) only runs every 10 steps for the progress bar; tensorboard scalars are unaffected. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
The default (uncached) path converted the image to float32 in the DataLoader worker, so the host-to-device copy moved 4x the bytes. Return the decoded image as the cached path of nerfstudio-project#1062 already does; simple_trainer converts to float on the device. Every consumer in the repo divides by 255.0 (simple_trainer, simple_trainer_2dgs, the val loop) or casts back to uint8 (colmap.py __main__). Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>
This branch has not been deployed
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
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.
Two small commits touching
examples/simple_trainer.pyandexamples/datasets/colmap.py.non_blocking=True(the loader already pins memory) andloss.item()is only called every 10 steps for the progress bar. Only the progress-bar text refreshes less often; tensorboard scalars are still written everytb_everysteps.Datasetreturns the decoded uint8 image in the default (uncached) loader too, as the--cache_imagespath of [perf] 3DGS training (simple_trainer.py) on 8x B200: 13.9 ms to 9.8 ms per step, from decoding one image per step #1062 already does; the trainer converts to float on the GPU, so the host-to-device copy is 4x smaller. Pixel values are identical. All consumers in the repo (simple_trainer,simple_trainer_2dgs, the val loop) divide by 255.0, andcolmap.py's__main__casts to uint8; third-party code that assumed float32 fromDatasetwould need.float().garden/4, 7k steps, default flags, seed 42, RTX 5060 Ti, 5 runs each, interleaved:
So commit 1 alone is a small gain (~1.5 %, close to the run-to-run spread) and most of the ~6 % total comes from commit 2. PSNR differences are within run-to-run noise. Rasterization still synchronizes once per step (
Intersect.cppreads the intersection count on the host), so this only removes the blocking copies and the per-stepitem()calls.