Repository: mindee/doctr

Issue:
[training] Pytorch training of recognition models is very slow
Pytorch trainings are very slow for recognition models, far more slower than tensorflow trainings with the same dataset.
It could either be:
- A torch dataloading issue (dataset/dataloader)
- A torch model issue
- A hardware issue (CUDA/GPU/CPU lack of compatibility)


Discussion:
We need to address this indeed!
Would you mind timing this the evaluation loop with both frameworks please? (since it's sequentially sampled it means we will process the exact same samples between Pytorch & TF)
As you mentioned, I think we need to time the dataloading, the switch to GPU + additional transforms, the model inference, and the backward pass. Spotting difference between TF & pytorch on those should help us to tackle this efficiently :) 

Resolving patch:
diff --git a/references/detection/train_pytorch.py b/references/detection/train_pytorch.py
--- a/references/detection/train_pytorch.py
+++ b/references/detection/train_pytorch.py
@@ -9,6 +9,7 @@
 
 import time
 import datetime
+import logging
 import multiprocessing as mp
 import numpy as np
 from fastprogress.fastprogress import master_bar, progress_bar
@@ -122,6 +123,12 @@ def main(args):
             raise AssertionError("PyTorch cannot access your GPU. Please investigate!")
         if args.device >= torch.cuda.device_count():
             raise ValueError("Invalid device index")
+    # Silent default switch to GPU if available
+    elif torch.cuda.is_available():
+        args.device = 0
+    else:
+        logging.warning("No accessible GPU, targe device set to CPU.")
+    if torch.cuda.is_available():
         torch.cuda.set_device(args.device)
         model = model.cuda()
 
diff --git a/references/recognition/train_pytorch.py b/references/recognition/train_pytorch.py
--- a/references/recognition/train_pytorch.py
+++ b/references/recognition/train_pytorch.py
@@ -7,6 +7,7 @@
 
 os.environ['USE_TORCH'] = '1'
 
+import logging
 import time
 import datetime
 import multiprocessing as mp
@@ -35,6 +36,8 @@ def fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, m
     for _ in progress_bar(range(len(train_loader)), parent=mb):
         images, targets = next(train_iter)
 
+        if torch.cuda.is_available():
+            images = images.cuda()
         images = batch_transforms(images)
 
         train_loss = model(images, targets)['loss']
@@ -58,6 +61,8 @@ def evaluate(model, val_loader, batch_transforms, val_metric):
     val_loss, batch_cnt = 0, 0
     val_iter = iter(val_loader)
     for images, targets in val_iter:
+        if torch.cuda.is_available():
+            images = images.cuda()
         images = batch_transforms(images)
         out = model(images, targets, return_preds=True)
         # Compute metric
@@ -114,6 +119,21 @@ def main(args):
         checkpoint = torch.load(args.resume, map_location='cpu')
         model.load_state_dict(checkpoint)
 
+    # GPU
+    if isinstance(args.device, int):
+        if not torch.cuda.is_available():
+            raise AssertionError("PyTorch cannot access your GPU. Please investigate!")
+        if args.device >= torch.cuda.device_count():
+            raise ValueError("Invalid device index")
+    # Silent default switch to GPU if available
+    elif torch.cuda.is_available():
+        args.device = 0
+    else:
+        logging.warning("No accessible GPU, targe device set to CPU.")
+    if torch.cuda.is_available():
+        torch.cuda.set_device(args.device)
+        model = model.cuda()
+
     # Metrics
     val_metric = TextMatch()
 


Test patch: