PyTorch: Improve inference batching speed#3099
Merged
AlexEMG merged 1 commit intoDeepLabCut:mainfrom Sep 21, 2025
Merged
Conversation
Member
|
This is excellent @arashsm79 -- thanks for the contribution! |
maximpavliv
approved these changes
Sep 18, 2025
Contributor
There was a problem hiding this comment.
Thanks for fixing the batching mechanism! 🚀
The fix significantly improves performance for larger batch sizes, which is a big win.
Not directly related to this PR, but I realized that the CTDInferenceRunner is lacking the multithreading scheme (preprocessing and batching preformed by a producer thread, prediction preformed by a consumer thread). Let's address this in a future PR.
Member
|
Looks great, let's fix merge conflicts and the merge this @arashsm79 |
AlexEMG
approved these changes
Sep 18, 2025
Member
AlexEMG
left a comment
There was a problem hiding this comment.
Fantastic @arashsm79 -- let's fix the conflicts and we're ready for rc13!
Use list accumulation for inference batches to eliminate O(n^2) torch.cat - Replaced incremental tensor _batch with list _batch_list - Stack only at batch processing time - Updated sequential and async inference paths
85402e3 to
61c5d68
Compare
2 tasks
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
4 participants
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.
Summary
This PR replaces incremental tensor concatenation ($O(n^2)$ ) during inference with a list-based accumulation. Final stacking now happens only when forming a batch, avoiding repeated reallocation and copy.
(Depends on #3094)
Main points:
Details
The previous pattern for batching images during inference:
caused$O(n^2)$ total memory movement for $n$ appended images. This was a bottleneck for larger batches, causing allocator churn and CPU overhead.
Profiling
This was confirmed by benchmarking the inference procedure with Torch Profile and Scalene.
Torch Profile
We can see for large batch sizes, the GPU is stalled and waits for the producer CPU thread to finish preprocessing the images.
Scalene
Statistical analysis shows that the concatenation operation is one of the main hot spots of the inference procedure due to reallocation and copying of the immutable tensor:
Results
We can see that these changes fix the problem with large batch sizes being inefficient (thanks @maximpavliv for running the benchmark):