feat: Added support of GPU for predictors in PyTorch

[fg-mindee]

This PR simply adds a dynamic device selection for predictor inference.
This snippet:

import os

os.environ['USE_TORCH'] = '1'
from doctr.io import DocumentFile
from doctr.models import ocr_predictor

doc = DocumentFile.from_pdf('/path/to/sample.pdf').as_images()
predictor = ocr_predictor(pretrained=True).cuda()

out = predictor(doc)

used to yield:

RuntimeError: Input type (torch.FloatTensor) and weight type (torch.cuda.FloatTensor) should be the same or input should be a MKLDNN tensor and weight is a dense tensor

because the model is on cuda, while the inputs are never moved to GPU.
And now, the snippet executes perfectly (& much faster than on CPU)

Since the predictions returned by each postprocessor are in numpy (automatically moved to CPU), this doesn't break any previous behaviour :)

Any feedback is welcome!

[codecov]

Codecov Report

Merging #808 (26d5a49) into main (45d4bfd) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main     #808   +/-   ##
=======================================
  Coverage   96.01%   96.01%           
=======================================
  Files         131      131           
  Lines        4942     4944    +2     
=======================================
+ Hits         4745     4747    +2     
  Misses        197      197           

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

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

Impacted Files	Coverage Δ
doctr/models/detection/predictor/pytorch.py	94.73% <100.00%> (+0.29%)	⬆️
doctr/models/recognition/predictor/pytorch.py	91.17% <100.00%> (+0.26%)	⬆️
doctr/transforms/modules/base.py	94.59% <0.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 45d4bfd...26d5a49. Read the comment docs.

[charlesmindee]

Thanks for the fix

[gganes3]

When using a pretrained model, using map_location="gpu" throws an error
RuntimeError: don't know how to restore data location of torch.FloatStorage (tagged with gpu)
det_params = torch.load(r"C:\Projects\db_resnet50-ac60cadc.pt", map_location="gpu")
reco_params = torch.load(r"C:\Projects\crnn_vgg16_bn-9762b0b0.pt", map_location="gpu")
# det_model.load_state_dict()
det_model.load_state_dict(det_params)
reco_model.load_state_dict(reco_params)

    # Ask the preprocessor of each task to resize and normalize similarly to your training
    # cf. https://github.com/mindee/doctr/blob/main/references/detection/train_pytorch.py#L94 & https://github.com/mindee/doctr/blob/main/references/detection/train_pytorch.py#L109
    det_predictor = DetectionPredictor(
        PreProcessor((128, 1024), batch_size=1, mean=(0.798, 0.785, 0.772), std=(0.264, 0.2749, 0.287)), det_model)
    # cf. https://github.com/mindee/doctr/blob/main/references/recognition/train_pytorch.py#L97 & https://github.com/mindee/doctr/blob/main/references/recognition/train_pytorch.py#L111
    reco_predictor = RecognitionPredictor(
        PreProcessor((32, 128), preserve_aspect_ratio=True, batch_size=32, mean=(0.694, 0.695, 0.693),
                     std=(0.299, 0.296, 0.301)), reco_model)


    model = OCRPredictor(det_predictor, reco_predictor).cuda()

[fg-mindee]

@gganes3 you're loading params on GPU, into modules that are on CPU so this is expected

when instantiating your model, I'd suggest to create it on CPU, load your params using map_location="cpu" and then move the whole module to GPU 👍