feat: ✨ PyTorch Recognition Model Multi-GPU support

[odulcy-mindee]

Fix #1146

Highly inspired from:

• https://pytorch.org/tutorials/intermediate/ddp_tutorial.html#comparison-between-dataparallel-and-distributeddataparallel
• https://pytorch.org/docs/stable/notes/ddp.html
• https://pytorch.org/docs/stable/distributed.html#which-backend-to-use

How to use this script

wget https://github.com/mindee/doctr/releases/download/v0.3.1/toy_recogition_set-036a4d80.zip
sudo apt-get update && sudo apt-get install unzip -y
unzip toy_recogition_set-036a4d80.zip -d reco_set

python references/recognition/train_pytorch_ddp.py crnn_mobilenet_v3_small --train_path ./reco_set --val_path ./reco_set -b 2 --epochs 10 --devices 0 1

How to retrieve device id

import torch
devices = [torch.cuda.device(i) for i in range(torch.cuda.device_count())]
device_names = [torch.cuda.get_device_name(d) for d in devices]

Update 28/03/23

We created a Dataset to verify if train_pytorch_ddp.py is correct. Training has been launched today on two GPUs, we'll see if the modifications are correct. Hopefully, this PR will be merged tomorrow if there is no issue.

[codecov]

Codecov Report

Merging #1164 (86a1dd0) into main (9e254fe) will not change coverage.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##             main    #1164   +/-   ##
=======================================
  Coverage   95.01%   95.01%           
=======================================
  Files         150      150           
  Lines        6418     6418           
=======================================
  Hits         6098     6098           
  Misses        320      320           

Flag	Coverage Δ
unittests	95.01% <100.00%> (ø)

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

Impacted Files	Coverage Δ
doctr/models/factory/hub.py	74.69% <100.00%> (ø)

Help us with your feedback. Take ten seconds to tell us how you rate us. Have a feature suggestion? Share it here.

[aminemindee]

LGTM 👍

[felixdittrich92]

Hi @odulcy-mindee thanks for the PR 👍
Could you update the readme in references with an example how to use this script ? And maybe add this also for the detection task ? :)

[odulcy-mindee]

@felixdittrich92 Hello! Yep, I'll update the README first and if possible detection task !