Fix device allocation for top-down models#174
Open
C-Achard wants to merge 6 commits into
Open
Conversation
Update PyTorch top-down pose postprocessing to keep tensors on the same device/dtype by using `new_zeros` and `new_tensor` instead of creating CPU tensors. Also replace `torch.cat(poses)` with `torch.stack(..., dim=0)` so batched pose outputs preserve the expected batch dimension.
Refactor bbox creation in `SkipFrames.__call__` to build tensors from `pose` (`new_zeros/new_tensor/new_ones`) so boxes and scores stay on the same device and dtype. The NaN handling for min/max coordinates is now explicit via `nan=` arguments, and bbox limiting uses `torch.clamp` with typed min/max tensors.
Update the PyTorch top-down path to use a correctly structured fallback `offsets_and_scales` value when no detections are found. This keeps the data shape consistent with post-processing expectations and avoids downstream tuple unpacking/type mismatches.
C-Achard
added
bug
Contributor
There was a problem hiding this comment.
Pull request overview
This PR targets the PyTorch pose-estimation runner’s top-down pipeline, aiming to make tensor creation device/dtype-consistent and to fix shape/parameter handling issues during top-down post-processing.
Changes:
- Switched bbox tensor construction to
pose.new_*-style APIs for device/dtype consistency. - Adjusted top-down offset/scale fallback structure and improved post-processing aggregation (
stackvscat). - Improved empty-batch handling in
_postprocess_top_downviabatch_pose.new_zeros(...).
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+98
to
+101
| self._detections = { | ||
| "boxes": bboxes, | ||
| "scores": pose.new_ones(num_det), | ||
| } |
Improve the PyTorch runner’s frame handling by validating HWC input shape and allowed channel counts, correctly unpacking dimensions, and dropping alpha for 4-channel frames. Also fix the top-down fallback `offsets_and_scales` format to use paired float scales instead of a scalar, preventing downstream shape/format issues.
In `PyTorchRunner`, transformed pose tensors are now merged with `torch.cat(..., dim=0)` instead of `torch.stack`. This aligns aggregation with tensors that already carry a leading batch dimension and avoids introducing an extra axis.
Introduce a reusable `DebugTimingStats` utility in `dlclive.utils` and instrument `PyTorchRunner.get_pose` with stage-level timing (input conversion, detector, pose transform/inference/decode, dynamic crop/update, postprocess, and CPU transfer). The runner now synchronizes CUDA around timed inference sections for accurate measurements and periodically logs aggregated performance stats. This also fixes top-down fallback metadata shape (`offsets_and_scales`) and restores the missing return in `_postprocess_top_down`.
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.
Fixes parts of the pose estimation runner to improve tensor handling and ensure consistency in the use of PyTorch tensor operations by:
Tensor handling improvements:
torch.zerosandtorch.oneswithpose.new_zerosandpose.new_onesto ensure tensors are created on the correct device and with the correct dtype. Also, usedpose.new_tensorfor creating tensors from lists, improving consistency._postprocess_top_down, replacedtorch.zeroswithbatch_pose.new_zerosfor empty batch handling, and ensured all scale and offset operations use tensors created with the same device and dtype as the pose input.Bug fixes and API consistency:
offsets_and_scalesto always be a pair of tuples, ensuring the expected formattorch.cattotorch.stackto preserve batch dimension and avoid shape errors