Use torch.inference_mode vs torch.no_grad for predictions

[thetonus]

🚀 The feature

I think it would be a good idea to create a decorator (similar to the one in Yolov8) that automatically applies the correct context manager for forward passes for code using the torch backend.

Motivation, pitch

Since torch>=1.9.0, there is a new context manager for inference - torch.inference_mode. This would allow for inferences to be faster and more memory conserved than just using torch.no_grad.

From Pytorch Docs:

Code run under this mode gets better performance by disabling view tracking and version counter bumps. Note that unlike some other mechanisms that locally enable or disable grad, entering inference_mode also disables to forward-mode AD.

Alternatives

None. Since Doctr is backwards compatible with older torch versions, there is no way to dynamically use torch.inference_mode vs torch.no_grad.

Additional context

No response

[thetonus]

I am more than happy to create a PR for this feature too.

[felixdittrich92]

Hi @thetonus 👋,

Thanks for the request and sure feel free to open a PR 😊
We can use the decorator directly because docTR requires torch>=1.12 so this should be fine 👍

[felixT2K]

@thetonus Part of this PR would be to add a short benchmark about inference latency and memory usage compared to the current state

Code to modify:

doctr/doctr/models/predictor/pytorch.py

Line 62 in 7ab0ece

@torch.no_grad()

doctr/doctr/models/kie_predictor/pytorch.py

Line 62 in 7ab0ece

@torch.no_grad()

doctr/doctr/models/recognition/predictor/pytorch.py

Line 43 in 7ab0ece

@torch.no_grad()

doctr/doctr/models/detection/predictor/pytorch.py

Line 35 in 7ab0ece

@torch.no_grad()

doctr/doctr/models/classification/predictor/pytorch.py

Line 36 in 7ab0ece

@torch.no_grad()

doctr/references/classification/latency_pytorch.py

Line 22 in 7ab0ece

@torch.no_grad()

doctr/scripts/detect_artefacts.py

Line 40 in 7ab0ece

@torch.no_grad()

doctr/references/detection/evaluate_pytorch.py

Line 29 in 7ab0ece

@torch.no_grad()

doctr/references/detection/latency_pytorch.py

Line 22 in 7ab0ece

@torch.no_grad()

doctr/references/obj_detection/latency_pytorch.py

Line 22 in 7ab0ece

@torch.no_grad()

doctr/references/recognition/evaluate_pytorch.py

Line 25 in 7ab0ece

@torch.no_grad()

doctr/references/recognition/latency_pytorch.py

Line 22 in 7ab0ece

@torch.no_grad()