[PyTorch] replace no_grad with inference_mode

[felixdittrich92]

This PR:

• change from no_grad to inference_mode (ref.: https://pytorch.org/docs/stable/generated/torch.inference_mode.html)
• ref: Use torch.inference_mode vs torch.no_grad for predictions #1316

from doctr.io import DocumentFile
from doctr.models import ocr_predictor
import time

doc_path = "/home/felix/Desktop/1.jpg"
start = time.time()
doc = DocumentFile.from_images([doc_path] * 10)
model = ocr_predictor(det_arch="db_resnet50", reco_arch="parseq", pretrained=True, det_bs=8, reco_bs=512).cuda().half()
result = model(doc)
print(time.time() - start)

# with torch.no_grad():
# cuda: VRAM: ~4400MB time: 15.616s
# cuda half: VRAM: ~2400MB time: 9.3s

# with torch.inference_mode():
# cuda: VRAM: ~4246MB time: 15.612s
# cuda half: VRAM: ~2340MB time: 9.14s

[codecov]

Codecov Report

Merging #1323 (6064f57) into main (8245706) will increase coverage by 0.01%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##             main    #1323      +/-   ##
==========================================
+ Coverage   95.75%   95.76%   +0.01%     
==========================================
  Files         154      154              
  Lines        6902     6902              
==========================================
+ Hits         6609     6610       +1     
+ Misses        293      292       -1     

Flag	Coverage Δ
unittests	95.76% <100.00%> (+0.01%)	⬆️

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

Files Changed	Coverage Δ
doctr/models/classification/predictor/pytorch.py	95.45% <100.00%> (ø)
doctr/models/detection/predictor/pytorch.py	95.23% <100.00%> (ø)
doctr/models/kie_predictor/pytorch.py	100.00% <100.00%> (ø)
doctr/models/predictor/pytorch.py	100.00% <100.00%> (ø)
doctr/models/recognition/predictor/pytorch.py	91.66% <100.00%> (ø)

... and 2 files with indirect coverage changes

[odulcy-mindee]

Should the modification be made for the evaluate function in references/classification/train_pytorch.py ? Same question for the others train_pytorch.py ?

[felixdittrich92]

Should the modification be made for the evaluate function in references/classification/train_pytorch.py ? Same question for the others train_pytorch.py ?

No because in the training scripts we switch between train mode and eval mode this works with no_grad but if the model/pipe is once in inference mode we cannot switch back into train mode but that's fine for inference :)

[odulcy-mindee]

Ok, thanks @felixdittrich92 !