Add torch.compile support for pytorch 2.4 #1690

Fabioomega · 2024-08-08T18:07:40Z

Added support for torch.compile only for version 2.4 or higher of pytorch. Included support for all the detection models and a recognition model (parseq).

Unfortunately, triton support is only available on linux plataforms. WSL seems to work fine tough, so it may be used that way.

Example use:

Enable the feature by setting the enviroment variable: USE_TRITON = YES.
Try the following code:

from doctr.models import ocr_predictor
from doctr.io import Document
from cv2 import imread
from time import time

t1 = time()
reader = ocr_predictor(det_arch='db_resnet50', pretrained=True)
print('Loading time:', time() - t1)

img = imread("<img>")
t2 = time()
d: Document = reader([img])
print('Document of the first try ', d)
print('Processing time of the first try', time() - t2)
t3 = time()
d: Document = reader([img])
print('Processing time of the second try:', time() - t3)

felixdittrich92 · 2024-08-09T06:26:12Z

doctr/models/recognition/parseq/pytorch.py

@@ -266,7 +266,8 @@ def decode_autoregressive(self, features: torch.Tensor, max_len: Optional[int] =
        ).int()

        pos_logits = []
-        for i in range(max_length):
+        i = 0
+        while i < max_length:


I remember there was a issue with while loops by exporting to onnx so we have to be careful here (needs to be checked)

I changed it because it was some unecessary complication related to breaks in torch.compile. Changing to a while loop and changing the logic a bit helped. Hopefully it works for the onnx also

doctr/models/detection/_utils/pytorch_compile.py

…on for compatibility with future backends

…rch_backend_available

felixdittrich92

Hi @Fabioomega 👋

Tests looks already good to me 👍
Docs section missing here: https://github.com/mindee/doctr/blob/main/docs/source/using_doctr/using_model_export.rst

As mentioned a table would be great :)
For the classification models it would be enough to add both orientation models to the table (i don't think we should blow up the table by adding all the backbone models)

Todo's:

comments
unittests
docs

Left some comments to revert unrequired parts :)

As mentioned for follow up PR's we can focus on fixes for the models which does not work yet out of the box 👍

felixdittrich92 · 2024-09-02T06:16:25Z

doctr/file_utils.py

@@ -76,6 +75,18 @@
        " is installed and that either USE_TF or USE_TORCH is enabled."
    )

+if _torch_available:


@Fabioomega We can remove this.

2 options:

We pin the lower boundary to >= 2.0.0 here

doctr/pyproject.toml

Line 61 in 9045dcf

"torch>=1.12.0,<3.0.0",

doctr/pyproject.toml

Line 105 in 9045dcf

"torch>=1.12.0,<3.0.0",

and torchvision>=0.15.0

or we mention in the docs that this requires >= 2.0.0 for compile and >=2.4.0 for compile + fullgraph

@odulcy-mindee wdyt ?
We are already at 2.4.0 so i would prefer the >=2.0.0 pin (in this case only to mention >=2.4.0 for fullgraph (triton) support)

felixdittrich92 · 2024-09-02T06:16:33Z

doctr/file_utils.py

@@ -104,3 +115,11 @@ def is_torch_available():
 def is_tf_available():
    """Whether TensorFlow is installed."""
    return _tf_available
+
+def does_torch_have_compile_capability():


felixdittrich92 · 2024-09-02T06:16:48Z

doctr/models/classification/zoo.py

can be reverted complete

doctr/models/detection/_utils/__init__.py

felixdittrich92 · 2024-09-02T06:17:14Z

doctr/models/detection/fast/base.py

felixdittrich92 · 2024-09-02T06:17:33Z

doctr/models/detection/zoo.py

felixdittrich92 · 2024-09-02T06:17:48Z

doctr/models/recognition/parseq/pytorch.py

felixdittrich92 · 2024-09-02T06:18:06Z

doctr/models/recognition/zoo.py

felixdittrich92 · 2024-09-02T06:18:17Z

doctr/models/zoo.py

I hame some questions about that! Wasn't the original ideia to add a new argument to enable compilation? Did I misunderstood?

That was the first thought as your code looked like changes to the pipeline/models were needed. However, we then saw that these were not needed.
Which is why we only add tests here and a section on how to use it. The compilation therefore remains on the user side, which is at the same time much more flexible. :)
Additional this avoids to add a arg which at the end only does -> model = torch.compile(model, ..) and is backend depending (PyTorch).

A full sample would look like then for example:

import requests import torch from doctr.models import ocr_predictor, parseq, fast_base from doctr.io import DocumentFile bytes_data = requests.get( "https://i1.rgstatic.net/publication/231831562_Another_Boring_Day_in_Paradise_Rock_and_Roll_and_the_Empowerment_of_Everyday_Life/links/57d02a2408ae601b39a05636/largepreview.png" ).content doc = DocumentFile.from_images([bytes_data]) rec_model = torch.compile(parseq(pretrained=True)) det_model = torch.compile(fast_base(pretrained=True)) predictor = ocr_predictor(det_arch=det_model, reco_arch=rec_model, pretrained=True) res = predictor(doc) res.show()

The only required change here would be to allow also:
torch._dynamo.eval_frame.OptimizedModule in

doctr/doctr/models/recognition/zoo.py

Line 39 in 9045dcf

arch, (recognition.CRNN, recognition.SAR, recognition.MASTER, recognition.ViTSTR, recognition.PARSeq)

and

doctr/doctr/models/detection/zoo.py

Line 59 in 9045dcf

if not isinstance(arch, (detection.DBNet, detection.LinkNet, detection.FAST)):

and

doctr/doctr/models/classification/zoo.py

Line 45 in 9045dcf

if not isinstance(arch, classification.MobileNetV3):

felixdittrich92 · 2024-09-02T06:20:22Z

tests/pytorch/test_models_detection_pt.py

@@ -186,3 +186,46 @@ def test_models_onnx_export(arch_name, input_shape, output_size):
        assert np.allclose(pt_logits, ort_outs[0], atol=1e-4)
    except AssertionError:
        pytest.skip(f"Output of {arch_name}:\nMax element-wise difference: {np.max(np.abs(pt_logits - ort_outs[0]))}")
+
+@pytest.mark.skipif(not does_torch_have_compile_capability(), reason="requires pytorch >= 2.0.0")
+@pytest.mark.skipif(not is_pytorch_backend_available(), reason="requires pytorch backend to be available")


remove the first two skipif Ci runs always on latest pytorch - same for the other tests 👍

felixdittrich92 · 2024-09-18T09:36:16Z

Hi @Fabioomega :)

Are you still interested to work on it ? :)

Fabioomega added 5 commits August 5, 2024 16:15

Added support for pytorch.compile on >=Pytorch 2.4 version

1bc21f8

Added support for controlling compiling behavior via predictors

479c897

Added fullgraph support for parseq

f8a5ee6

Fixed typo

60ca8eb

Fixed typo again

22dd3d9

Fabioomega marked this pull request as draft August 8, 2024 18:13

Removed compile kwarg

38e5a2d

felixdittrich92 reviewed Aug 9, 2024

View reviewed changes

doctr/models/detection/_utils/pytorch_compile.py Outdated Show resolved Hide resolved

Fabioomega added 5 commits August 13, 2024 17:35

Fixed crash in boundingRect when assume_straight_pages=True

5822d4a

Added compile and compile_kwargs and removed explicit mention of trit…

0ded837

…on for compatibility with future backends

Fixed some remaining name changes from is_triton_available to is_pyto…

f7ae74d

…rch_backend_available

Removed changes to the postprocessing step and reverted back

1e6869b

Added test cases for classification, detection and recognition

9850455

felixdittrich92 requested changes Sep 2, 2024

View reviewed changes

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Add torch.compile support for pytorch 2.4 #1690

Add torch.compile support for pytorch 2.4 #1690

Fabioomega commented Aug 8, 2024

felixdittrich92 Aug 9, 2024

Fabioomega Aug 9, 2024

felixdittrich92 left a comment •

edited

Loading

felixdittrich92 Sep 2, 2024 •

edited

Loading

felixdittrich92 Sep 2, 2024

felixdittrich92 Sep 2, 2024

felixdittrich92 Sep 2, 2024

felixdittrich92 Sep 2, 2024

felixdittrich92 Sep 2, 2024

felixdittrich92 Sep 2, 2024

felixdittrich92 Sep 2, 2024

Fabioomega Sep 18, 2024

felixT2K Sep 18, 2024 •

edited

Loading

felixT2K Sep 18, 2024

felixdittrich92 Sep 2, 2024

felixdittrich92 commented Sep 18, 2024

Add torch.compile support for pytorch 2.4 #1690

Are you sure you want to change the base?

Add torch.compile support for pytorch 2.4 #1690

Conversation

Fabioomega commented Aug 8, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

felixdittrich92 left a comment • edited Loading

Choose a reason for hiding this comment

felixdittrich92 Sep 2, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

felixT2K Sep 18, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

felixdittrich92 commented Sep 18, 2024

felixdittrich92 left a comment •

edited

Loading

felixdittrich92 Sep 2, 2024 •

edited

Loading

felixT2K Sep 18, 2024 •

edited

Loading