Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[Feat] Add torch.compile support #1791

Open
wants to merge 6 commits into
base: main
Choose a base branch
from

Conversation

felixdittrich92
Copy link
Contributor

@felixdittrich92 felixdittrich92 commented Nov 21, 2024

This PR:

  • Add torch.compile basic support (different backends / fullgraph is up to the user - we support only the basic compile compatibility with inductor (default) backend)
  • Add corresponding tests and a documentation section
  • No benchmarks added here because this scales with the users hardware

Any feedback is welcome 🤗

Closes: #1684 #1690

@felixdittrich92 felixdittrich92 added topic: documentation Improvements or additions to documentation module: models Related to doctr.models framework: pytorch Related to PyTorch backend topic: text detection Related to the task of text detection topic: text recognition Related to the task of text recognition topic: character classification Related to the task of character classification type: new feature New feature ext: docs Related to docs folder labels Nov 21, 2024
@felixdittrich92 felixdittrich92 added this to the 0.11.0 milestone Nov 21, 2024
@felixdittrich92 felixdittrich92 self-assigned this Nov 21, 2024
Copy link

codecov bot commented Nov 21, 2024

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 96.59%. Comparing base (b53517a) to head (f4e8545).

Additional details and impacted files
@@            Coverage Diff            @@
##           main    #1791       +/-   ##
=========================================
+ Coverage      0   96.59%   +96.59%     
=========================================
  Files         0      165      +165     
  Lines         0     7938     +7938     
=========================================
+ Hits          0     7668     +7668     
- Misses        0      270      +270     
Flag Coverage Δ
unittests 96.59% <100.00%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.


🚨 Try these New Features:

@felixdittrich92
Copy link
Contributor Author

@odulcy-mindee This slows down our pytorch test CI extremly (before ~9min | now ~20min) because the compilation takes some time :/

@@ -195,11 +195,16 @@ def forward(
out["out_map"] = prob_map

if target is None or return_preds:
# Disable for torch.compile compatibility
@torch.compiler.disable # type: ignore[attr-defined]
Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

post proc can't be compiled so model runs fully compiled and all the other parts keept as is

@felixdittrich92
Copy link
Contributor Author

felixdittrich92 commented Nov 21, 2024

Then i think rdy to review :)

I would say we shouldn't add a benchmark to the docs here because this is really user hardware depending or how do you see it ?

At the end the Onnx way is still to prefer ... compile brings ~5-10% boost ..Onnx boosts by +50% (on CPU) ^^

@felixdittrich92 felixdittrich92 marked this pull request as ready for review November 21, 2024 11:13
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
ext: docs Related to docs folder framework: pytorch Related to PyTorch backend module: models Related to doctr.models topic: character classification Related to the task of character classification topic: documentation Improvements or additions to documentation topic: text detection Related to the task of text detection topic: text recognition Related to the task of text recognition type: new feature New feature
Projects
None yet
Development

Successfully merging this pull request may close these issues.

Add support for torch.compile
2 participants