Enable GPU inference for transformers models #4587
Conversation
WalkthroughThe recent changes to the Changes
Sequence Diagram(s)sequenceDiagram
participant User
participant FiftyOneTransformer
User->>FiftyOneTransformer: Initialize with config
FiftyOneTransformer->>FiftyOneTransformer: Set device (CPU/GPU)
User->>FiftyOneTransformer: Call _embed with inputs
FiftyOneTransformer->>FiftyOneTransformer: Move inputs to device
FiftyOneTransformer->>FiftyOneTransformer: Process inputs
User->>FiftyOneTransformer: Call _predict with inputs
FiftyOneTransformer->>FiftyOneTransformer: Move inputs to device
FiftyOneTransformer->>FiftyOneTransformer: Generate predictions
Poem
Thank you for using CodeRabbit. We offer it for free to the OSS community and would appreciate your support in helping us grow. If you find it useful, would you consider giving us a shout-out on your favorite social media? TipsChatThere are 3 ways to chat with CodeRabbit:
Note: Be mindful of the bot's finite context window. It's strongly recommended to break down tasks such as reading entire modules into smaller chunks. For a focused discussion, use review comments to chat about specific files and their changes, instead of using the PR comments. CodeRabbit Commands (invoked as PR comments)
Additionally, you can add CodeRabbit Configuration File (
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 0
Review details
Configuration used: .coderabbit.yaml
Review profile: CHILL
Files selected for processing (1)
- fiftyone/utils/transformers.py (13 hunks)
Additional comments not posted (11)
fiftyone/utils/transformers.py (11)
454-456: LGTM! Device initialization is correct.The
self.deviceattribute is correctly initialized based on whether the model parameters are on CUDA or CPU.
501-503: LGTM! Device initialization is correct.The
self.deviceattribute is correctly initialized based on whether the model parameters are on CUDA or CPU.
378-378: LGTM! Device management is correct.The input tensors are correctly moved to the specified device before processing.
399-401: LGTM! Device management is correct.The input tensors are correctly moved to the specified device before processing.
438-440: LGTM! Device management is correct.The input tensors are correctly moved to the specified device before processing.
704-704: LGTM! Device management is correct.The input tensors are correctly moved to the specified device before processing.
603-603: LGTM! Device management is correct.The input tensors are correctly moved to the specified device before processing.
785-785: LGTM! Device management is correct.The input tensors are correctly moved to the specified device before processing.
832-832: LGTM! Device management is correct.The input tensors are correctly moved to the specified device before processing.
890-890: LGTM! Device management is correct.The input tensors are correctly moved to the specified device before processing.
939-939: LGTM! Device management is correct.The input tensors are correctly moved to the specified device before processing.
What changes are proposed in this pull request?
Automatically detect if transformers model is on the GPU and if so perform inference on same device.
How is this patch tested? If it is not, please explain why.
Tested in collab with and without GPU, results are within precision difference of each other and no changes on CPU.
Release Notes
Is this a user-facing change that should be mentioned in the release notes?
notes for FiftyOne users.
(Details in 1-2 sentences. You can just refer to another PR with a description
if this PR is part of a larger change.)
What areas of FiftyOne does this PR affect?
fiftyonePython library changesSummary by CodeRabbit