From 42f578989ee658ce834f55d9d309877b1a968bb8 Mon Sep 17 00:00:00 2001
From: fg-mindee <fg@mindee.co>
Date: Wed, 10 Nov 2021 16:29:17 +0100
Subject: [PATCH 1/9] refactor: Renamed the model argument to arch

---
 references/classification/train_pytorch.py    | 6 +++---
 references/classification/train_tensorflow.py | 6 +++---
 references/detection/train_pytorch.py         | 6 +++---
 references/detection/train_tensorflow.py      | 6 +++---
 references/recognition/train_pytorch.py       | 6 +++---
 references/recognition/train_tensorflow.py    | 6 +++---
 6 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/references/classification/train_pytorch.py b/references/classification/train_pytorch.py
index 9400014edb..d8b4bfaedb 100644
--- a/references/classification/train_pytorch.py
+++ b/references/classification/train_pytorch.py
@@ -104,7 +104,7 @@ def main(args):
     batch_transforms = Normalize(mean=(0.694, 0.695, 0.693), std=(0.299, 0.296, 0.301))
 
     # Load doctr model
-    model = models.__dict__[args.model](pretrained=args.pretrained, num_classes=len(vocab))
+    model = models.__dict__[args.arch](pretrained=args.pretrained, num_classes=len(vocab))
 
     # Resume weights
     if isinstance(args.resume, str):
@@ -163,7 +163,7 @@ def main(args):
 
     # Training monitoring
     current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
-    exp_name = f"{args.model}_{current_time}" if args.name is None else args.name
+    exp_name = f"{args.arch}_{current_time}" if args.name is None else args.name
 
     # W&B
     if args.wb:
@@ -214,7 +214,7 @@ def parse_args():
     parser = argparse.ArgumentParser(description='DocTR training script for character classification (PyTorch)',
                                      formatter_class=argparse.ArgumentDefaultsHelpFormatter)
 
-    parser.add_argument('model', type=str, help='text-recognition model to train')
+    parser.add_argument('arch', type=str, help='text-recognition model to train')
     parser.add_argument('--name', type=str, default=None, help='Name of your training experiment')
     parser.add_argument('--epochs', type=int, default=10, help='number of epochs to train the model on')
     parser.add_argument('-b', '--batch_size', type=int, default=64, help='batch size for training')
diff --git a/references/classification/train_tensorflow.py b/references/classification/train_tensorflow.py
index 236ef65a03..5059af7a01 100644
--- a/references/classification/train_tensorflow.py
+++ b/references/classification/train_tensorflow.py
@@ -100,7 +100,7 @@ def main(args):
           f"{val_loader.num_batches} batches)")
 
     # Load doctr model
-    model = backbones.__dict__[args.model](
+    model = backbones.__dict__[args.arch](
         pretrained=args.pretrained,
         input_shape=(args.input_size, args.input_size, 3),
         num_classes=len(vocab),
@@ -170,7 +170,7 @@ def main(args):
 
     # Tensorboard to monitor training
     current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
-    exp_name = f"{args.model}_{current_time}" if args.name is None else args.name
+    exp_name = f"{args.arch}_{current_time}" if args.name is None else args.name
 
     # W&B
     if args.wb:
@@ -222,7 +222,7 @@ def parse_args():
     parser = argparse.ArgumentParser(description='DocTR training script for character classification (TensorFlow)',
                                      formatter_class=argparse.ArgumentDefaultsHelpFormatter)
 
-    parser.add_argument('model', type=str, help='text-recognition model to train')
+    parser.add_argument('arch', type=str, help='text-recognition model to train')
     parser.add_argument('--name', type=str, default=None, help='Name of your training experiment')
     parser.add_argument('--epochs', type=int, default=10, help='number of epochs to train the model on')
     parser.add_argument('-b', '--batch_size', type=int, default=64, help='batch size for training')
diff --git a/references/detection/train_pytorch.py b/references/detection/train_pytorch.py
index d5ffada605..fa575ba9b1 100644
--- a/references/detection/train_pytorch.py
+++ b/references/detection/train_pytorch.py
@@ -109,7 +109,7 @@ def main(args):
     batch_transforms = Normalize(mean=(0.798, 0.785, 0.772), std=(0.264, 0.2749, 0.287))
 
     # Load doctr model
-    model = detection.__dict__[args.model](pretrained=args.pretrained)
+    model = detection.__dict__[args.arch](pretrained=args.pretrained)
 
     # Resume weights
     if isinstance(args.resume, str):
@@ -190,7 +190,7 @@ def main(args):
 
     # Training monitoring
     current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
-    exp_name = f"{args.model}_{current_time}" if args.name is None else args.name
+    exp_name = f"{args.arch}_{current_time}" if args.name is None else args.name
 
     # W&B
     if args.wb:
@@ -251,7 +251,7 @@ def parse_args():
 
     parser.add_argument('train_path', type=str, help='path to training data folder')
     parser.add_argument('val_path', type=str, help='path to validation data folder')
-    parser.add_argument('model', type=str, help='text-detection model to train')
+    parser.add_argument('arch', type=str, help='text-detection model to train')
     parser.add_argument('--name', type=str, default=None, help='Name of your training experiment')
     parser.add_argument('--epochs', type=int, default=10, help='number of epochs to train the model on')
     parser.add_argument('-b', '--batch_size', type=int, default=2, help='batch size for training')
diff --git a/references/detection/train_tensorflow.py b/references/detection/train_tensorflow.py
index 9287e8eec3..62088b8fa2 100644
--- a/references/detection/train_tensorflow.py
+++ b/references/detection/train_tensorflow.py
@@ -101,7 +101,7 @@ def main(args):
     ])
 
     # Load doctr model
-    model = detection.__dict__[args.model](
+    model = detection.__dict__[args.arch](
         pretrained=args.pretrained,
         input_shape=(args.input_size, args.input_size, 3)
     )
@@ -165,7 +165,7 @@ def main(args):
 
     # Tensorboard to monitor training
     current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
-    exp_name = f"{args.model}_{current_time}" if args.name is None else args.name
+    exp_name = f"{args.arch}_{current_time}" if args.name is None else args.name
 
     # Tensorboard
     tb_writer = None
@@ -245,7 +245,7 @@ def parse_args():
 
     parser.add_argument('train_path', type=str, help='path to training data folder')
     parser.add_argument('val_path', type=str, help='path to validation data folder')
-    parser.add_argument('model', type=str, help='text-detection model to train')
+    parser.add_argument('arch', type=str, help='text-detection model to train')
     parser.add_argument('--name', type=str, default=None, help='Name of your training experiment')
     parser.add_argument('--epochs', type=int, default=10, help='number of epochs to train the model on')
     parser.add_argument('-b', '--batch_size', type=int, default=2, help='batch size for training')
diff --git a/references/recognition/train_pytorch.py b/references/recognition/train_pytorch.py
index 59e9289f54..5b970a2061 100644
--- a/references/recognition/train_pytorch.py
+++ b/references/recognition/train_pytorch.py
@@ -111,7 +111,7 @@ def main(args):
     batch_transforms = Normalize(mean=(0.694, 0.695, 0.693), std=(0.299, 0.296, 0.301))
 
     # Load doctr model
-    model = recognition.__dict__[args.model](pretrained=args.pretrained, vocab=VOCABS[args.vocab])
+    model = recognition.__dict__[args.arch](pretrained=args.pretrained, vocab=VOCABS[args.vocab])
 
     # Resume weights
     if isinstance(args.resume, str):
@@ -193,7 +193,7 @@ def main(args):
 
     # Training monitoring
     current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
-    exp_name = f"{args.model}_{current_time}" if args.name is None else args.name
+    exp_name = f"{args.arch}_{current_time}" if args.name is None else args.name
 
     # W&B
     if args.wb:
@@ -250,7 +250,7 @@ def parse_args():
 
     parser.add_argument('train_path', type=str, help='path to train data folder(s)')
     parser.add_argument('val_path', type=str, help='path to val data folder')
-    parser.add_argument('model', type=str, help='text-recognition model to train')
+    parser.add_argument('arch', type=str, help='text-recognition model to train')
     parser.add_argument('--name', type=str, default=None, help='Name of your training experiment')
     parser.add_argument('--epochs', type=int, default=10, help='number of epochs to train the model on')
     parser.add_argument('-b', '--batch_size', type=int, default=64, help='batch size for training')
diff --git a/references/recognition/train_tensorflow.py b/references/recognition/train_tensorflow.py
index 4aa280b201..47b87f58ec 100644
--- a/references/recognition/train_tensorflow.py
+++ b/references/recognition/train_tensorflow.py
@@ -98,7 +98,7 @@ def main(args):
           f"{val_loader.num_batches} batches)")
 
     # Load doctr model
-    model = recognition.__dict__[args.model](
+    model = recognition.__dict__[args.arch](
         pretrained=args.pretrained,
         input_shape=(args.input_size, 4 * args.input_size, 3),
         vocab=VOCABS[args.vocab]
@@ -174,7 +174,7 @@ def main(args):
 
     # Tensorboard to monitor training
     current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
-    exp_name = f"{args.model}_{current_time}" if args.name is None else args.name
+    exp_name = f"{args.arch}_{current_time}" if args.name is None else args.name
 
     # Tensorboard
     tb_writer = None
@@ -245,7 +245,7 @@ def parse_args():
 
     parser.add_argument('train_path', type=str, help='path to train data folder(s)')
     parser.add_argument('val_path', type=str, help='path to val data folder')
-    parser.add_argument('model', type=str, help='text-recognition model to train')
+    parser.add_argument('arch', type=str, help='text-recognition model to train')
     parser.add_argument('--name', type=str, default=None, help='Name of your training experiment')
     parser.add_argument('--epochs', type=int, default=10, help='number of epochs to train the model on')
     parser.add_argument('-b', '--batch_size', type=int, default=64, help='batch size for training')

From f0986e7348f058a72cde5b2b5d5a986b39cdf0b0 Mon Sep 17 00:00:00 2001
From: fg-mindee <fg@mindee.co>
Date: Wed, 10 Nov 2021 16:30:22 +0100
Subject: [PATCH 2/9] refactor: Removed epoch index logging

---
 references/classification/train_pytorch.py    | 1 -
 references/classification/train_tensorflow.py | 1 -
 references/detection/train_pytorch.py         | 1 -
 references/detection/train_tensorflow.py      | 1 -
 references/recognition/train_pytorch.py       | 1 -
 references/recognition/train_tensorflow.py    | 1 -
 6 files changed, 6 deletions(-)

diff --git a/references/classification/train_pytorch.py b/references/classification/train_pytorch.py
index d8b4bfaedb..8641e42487 100644
--- a/references/classification/train_pytorch.py
+++ b/references/classification/train_pytorch.py
@@ -200,7 +200,6 @@ def main(args):
         # W&B
         if args.wb:
             wandb.log({
-                'epochs': epoch + 1,
                 'val_loss': val_loss,
                 'acc': acc,
             })
diff --git a/references/classification/train_tensorflow.py b/references/classification/train_tensorflow.py
index 5059af7a01..1bb85b2ee3 100644
--- a/references/classification/train_tensorflow.py
+++ b/references/classification/train_tensorflow.py
@@ -208,7 +208,6 @@ def main(args):
         # W&B
         if args.wb:
             wandb.log({
-                'epochs': epoch + 1,
                 'val_loss': val_loss,
                 'acc': acc,
             })
diff --git a/references/detection/train_pytorch.py b/references/detection/train_pytorch.py
index fa575ba9b1..0bcc74b47b 100644
--- a/references/detection/train_pytorch.py
+++ b/references/detection/train_pytorch.py
@@ -231,7 +231,6 @@ def main(args):
         # W&B
         if args.wb:
             wandb.log({
-                'epochs': epoch + 1,
                 'val_loss': val_loss,
                 'recall': recall,
                 'precision': precision,
diff --git a/references/detection/train_tensorflow.py b/references/detection/train_tensorflow.py
index 62088b8fa2..bec6dda0ad 100644
--- a/references/detection/train_tensorflow.py
+++ b/references/detection/train_tensorflow.py
@@ -225,7 +225,6 @@ def main(args):
         # W&B
         if args.wb:
             wandb.log({
-                'epochs': epoch + 1,
                 'val_loss': val_loss,
                 'recall': recall,
                 'precision': precision,
diff --git a/references/recognition/train_pytorch.py b/references/recognition/train_pytorch.py
index 5b970a2061..e15bd990e1 100644
--- a/references/recognition/train_pytorch.py
+++ b/references/recognition/train_pytorch.py
@@ -231,7 +231,6 @@ def main(args):
         # W&B
         if args.wb:
             wandb.log({
-                'epochs': epoch + 1,
                 'val_loss': val_loss,
                 'exact_match': exact_match,
                 'partial_match': partial_match,
diff --git a/references/recognition/train_tensorflow.py b/references/recognition/train_tensorflow.py
index 47b87f58ec..bcc86872e9 100644
--- a/references/recognition/train_tensorflow.py
+++ b/references/recognition/train_tensorflow.py
@@ -226,7 +226,6 @@ def main(args):
         # W&B
         if args.wb:
             wandb.log({
-                'epochs': epoch + 1,
                 'val_loss': val_loss,
                 'exact_match': exact_match,
                 'partial_match': partial_match,

From 0251045e777daa22fd2c42a586c6afbdda45835b Mon Sep 17 00:00:00 2001
From: fg-mindee <fg@mindee.co>
Date: Wed, 10 Nov 2021 16:32:43 +0100
Subject: [PATCH 3/9] feat: Added dataset hash logging

---
 references/detection/train_pytorch.py      | 5 +++++
 references/detection/train_tensorflow.py   | 5 +++++
 references/recognition/train_pytorch.py    | 5 +++++
 references/recognition/train_tensorflow.py | 5 +++++
 4 files changed, 20 insertions(+)

diff --git a/references/detection/train_pytorch.py b/references/detection/train_pytorch.py
index 0bcc74b47b..88eccad177 100644
--- a/references/detection/train_pytorch.py
+++ b/references/detection/train_pytorch.py
@@ -8,6 +8,7 @@
 os.environ['USE_TORCH'] = '1'
 
 import datetime
+import hashlib
 import logging
 import multiprocessing as mp
 import time
@@ -105,6 +106,8 @@ def main(args):
     )
     print(f"Validation set loaded in {time.time() - st:.4}s ({len(val_set)} samples in "
           f"{len(val_loader)} batches)")
+    with open(os.path.join(args.val_path, 'labels.json'), 'rb') as f:
+        val_hash = hashlib.sha256(f.read()).hexdigest()
 
     batch_transforms = Normalize(mean=(0.798, 0.785, 0.772), std=(0.264, 0.2749, 0.287))
 
@@ -167,6 +170,8 @@ def main(args):
     )
     print(f"Train set loaded in {time.time() - st:.4}s ({len(train_set)} samples in "
           f"{len(train_loader)} batches)")
+    with open(os.path.join(args.train_path, 'labels.json'), 'rb') as f:
+        train_hash = hashlib.sha256(f.read()).hexdigest()
 
     if args.show_samples:
         x, target = next(iter(train_loader))
diff --git a/references/detection/train_tensorflow.py b/references/detection/train_tensorflow.py
index bec6dda0ad..7e4818622a 100644
--- a/references/detection/train_tensorflow.py
+++ b/references/detection/train_tensorflow.py
@@ -9,6 +9,7 @@
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
 
 import datetime
+import hashlib
 import time
 from collections import deque
 from pathlib import Path
@@ -95,6 +96,8 @@ def main(args):
     val_loader = DataLoader(val_set, batch_size=args.batch_size, shuffle=False, drop_last=False, workers=args.workers)
     print(f"Validation set loaded in {time.time() - st:.4}s ({len(val_set)} samples in "
           f"{val_loader.num_batches} batches)")
+    with open(os.path.join(args.val_path, 'labels.json'), 'rb') as f:
+        val_hash = hashlib.sha256(f.read()).hexdigest()
 
     batch_transforms = T.Compose([
         T.Normalize(mean=(0.798, 0.785, 0.772), std=(0.264, 0.2749, 0.287)),
@@ -142,6 +145,8 @@ def main(args):
     train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, drop_last=True, workers=args.workers)
     print(f"Train set loaded in {time.time() - st:.4}s ({len(train_set)} samples in "
           f"{train_loader.num_batches} batches)")
+    with open(os.path.join(args.train_path, 'labels.json'), 'rb') as f:
+        train_hash = hashlib.sha256(f.read()).hexdigest()
 
     if args.show_samples:
         x, target = next(iter(train_loader))
diff --git a/references/recognition/train_pytorch.py b/references/recognition/train_pytorch.py
index e15bd990e1..0c32db4e30 100644
--- a/references/recognition/train_pytorch.py
+++ b/references/recognition/train_pytorch.py
@@ -8,6 +8,7 @@
 os.environ['USE_TORCH'] = '1'
 
 import datetime
+import hashlib
 import logging
 import multiprocessing as mp
 import time
@@ -107,6 +108,8 @@ def main(args):
     )
     print(f"Validation set loaded in {time.time() - st:.4}s ({len(val_set)} samples in "
           f"{len(val_loader)} batches)")
+    with open(os.path.join(args.val_path, 'labels.json'), 'rb') as f:
+        val_hash = hashlib.sha256(f.read()).hexdigest()
 
     batch_transforms = Normalize(mean=(0.694, 0.695, 0.693), std=(0.299, 0.296, 0.301))
 
@@ -175,6 +178,8 @@ def main(args):
     )
     print(f"Train set loaded in {time.time() - st:.4}s ({len(train_set)} samples in "
           f"{len(train_loader)} batches)")
+    with open(parts[0].joinpath('labels.json'), 'rb') as f:
+        train_hash = hashlib.sha256(f.read()).hexdigest()
 
     if args.show_samples:
         x, target = next(iter(train_loader))
diff --git a/references/recognition/train_tensorflow.py b/references/recognition/train_tensorflow.py
index bcc86872e9..6b841d9abf 100644
--- a/references/recognition/train_tensorflow.py
+++ b/references/recognition/train_tensorflow.py
@@ -9,6 +9,7 @@
 os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"
 
 import datetime
+import hashlib
 import time
 from collections import deque
 from pathlib import Path
@@ -96,6 +97,8 @@ def main(args):
     val_loader = DataLoader(val_set, batch_size=args.batch_size, shuffle=False, drop_last=False, workers=args.workers)
     print(f"Validation set loaded in {time.time() - st:.4}s ({len(val_set)} samples in "
           f"{val_loader.num_batches} batches)")
+    with open(os.path.join(args.val_path, 'labels.json'), 'rb') as f:
+        val_hash = hashlib.sha256(f.read()).hexdigest()
 
     # Load doctr model
     model = recognition.__dict__[args.arch](
@@ -151,6 +154,8 @@ def main(args):
     train_loader = DataLoader(train_set, batch_size=args.batch_size, shuffle=True, drop_last=True, workers=args.workers)
     print(f"Train set loaded in {time.time() - st:.4}s ({len(train_set)} samples in "
           f"{train_loader.num_batches} batches)")
+    with open(parts[0].joinpath('labels.json'), 'rb') as f:
+        train_hash = hashlib.sha256(f.read()).hexdigest()
 
     if args.show_samples:
         x, target = next(iter(train_loader))

From e17dbd2d978d01db2503274f98c9e94d09f5084e Mon Sep 17 00:00:00 2001
From: fg-mindee <fg@mindee.co>
Date: Wed, 10 Nov 2021 16:34:57 +0100
Subject: [PATCH 4/9] feat: Added elements to the payload logged on W&B at each
 epoch

---
 references/classification/train_pytorch.py    |  7 +++++--
 references/classification/train_tensorflow.py |  7 +++++--
 references/detection/train_pytorch.py         |  9 +++++++--
 references/detection/train_tensorflow.py      |  9 +++++++--
 references/recognition/train_pytorch.py       |  9 +++++++--
 references/recognition/train_tensorflow.py    | 10 ++++++++--
 6 files changed, 39 insertions(+), 12 deletions(-)

diff --git a/references/classification/train_pytorch.py b/references/classification/train_pytorch.py
index 8641e42487..65af66c200 100644
--- a/references/classification/train_pytorch.py
+++ b/references/classification/train_pytorch.py
@@ -174,12 +174,15 @@ def main(args):
             config={
                 "learning_rate": args.lr,
                 "epochs": args.epochs,
+                "weight_decay": args.weight_decay,
                 "batch_size": args.batch_size,
-                "architecture": args.model,
+                "architecture": args.arch,
                 "input_size": args.input_size,
                 "optimizer": "adam",
-                "exp_type": "character-classification",
                 "framework": "pytorch",
+                "vocab": args.vocab,
+                "scheduler": args.sched,
+                "pretrained": args.pretrained,
             }
         )
 
diff --git a/references/classification/train_tensorflow.py b/references/classification/train_tensorflow.py
index 1bb85b2ee3..77e465ce06 100644
--- a/references/classification/train_tensorflow.py
+++ b/references/classification/train_tensorflow.py
@@ -181,12 +181,15 @@ def main(args):
             config={
                 "learning_rate": args.lr,
                 "epochs": args.epochs,
+                "weight_decay": args.weight_decay,
                 "batch_size": args.batch_size,
-                "architecture": args.model,
+                "architecture": args.arch,
                 "input_size": args.input_size,
                 "optimizer": "adam",
-                "exp_type": "character-classification",
                 "framework": "tensorflow",
+                "vocab": args.vocab,
+                "scheduler": args.sched,
+                "pretrained": args.pretrained,
             }
         )
 
diff --git a/references/detection/train_pytorch.py b/references/detection/train_pytorch.py
index 88eccad177..bb9f725fd9 100644
--- a/references/detection/train_pytorch.py
+++ b/references/detection/train_pytorch.py
@@ -206,12 +206,17 @@ def main(args):
             config={
                 "learning_rate": args.lr,
                 "epochs": args.epochs,
+                "weight_decay": args.weight_decay,
                 "batch_size": args.batch_size,
-                "architecture": args.model,
+                "architecture": args.arch,
                 "input_size": args.input_size,
                 "optimizer": "adam",
-                "exp_type": "text-detection",
                 "framework": "pytorch",
+                "scheduler": args.sched,
+                "train_hash": train_hash,
+                "val_hash": val_hash,
+                "pretrained": args.pretrained,
+                "rotation": args.rotation,
             }
         )
 
diff --git a/references/detection/train_tensorflow.py b/references/detection/train_tensorflow.py
index 7e4818622a..636e7141dd 100644
--- a/references/detection/train_tensorflow.py
+++ b/references/detection/train_tensorflow.py
@@ -188,12 +188,17 @@ def main(args):
             config={
                 "learning_rate": args.lr,
                 "epochs": args.epochs,
+                "weight_decay": args.weight_decay,
                 "batch_size": args.batch_size,
-                "architecture": args.model,
+                "architecture": args.arch,
                 "input_size": args.input_size,
                 "optimizer": "adam",
-                "exp_type": "text-detection",
                 "framework": "tensorflow",
+                "scheduler": args.sched,
+                "train_hash": train_hash,
+                "val_hash": val_hash,
+                "pretrained": args.pretrained,
+                "rotation": args.rotation,
             }
         )
 
diff --git a/references/recognition/train_pytorch.py b/references/recognition/train_pytorch.py
index 0c32db4e30..ae0f31de00 100644
--- a/references/recognition/train_pytorch.py
+++ b/references/recognition/train_pytorch.py
@@ -209,12 +209,17 @@ def main(args):
             config={
                 "learning_rate": args.lr,
                 "epochs": args.epochs,
+                "weight_decay": args.weight_decay,
                 "batch_size": args.batch_size,
-                "architecture": args.model,
+                "architecture": args.arch,
                 "input_size": args.input_size,
                 "optimizer": "adam",
-                "exp_type": "text-recognition",
                 "framework": "pytorch",
+                "scheduler": args.sched,
+                "vocab": args.vocab,
+                "train_hash": train_hash,
+                "val_hash": val_hash,
+                "pretrained": args.pretrained,
             }
         )
 
diff --git a/references/recognition/train_tensorflow.py b/references/recognition/train_tensorflow.py
index 6b841d9abf..034b870eea 100644
--- a/references/recognition/train_tensorflow.py
+++ b/references/recognition/train_tensorflow.py
@@ -197,11 +197,17 @@ def main(args):
             config={
                 "learning_rate": args.lr,
                 "epochs": args.epochs,
+                "weight_decay": args.weight_decay,
                 "batch_size": args.batch_size,
-                "architecture": args.model,
+                "architecture": args.arch,
                 "input_size": args.input_size,
                 "optimizer": "adam",
-                "exp_type": "text-recognition",
+                "framework": "tensorflow",
+                "scheduler": args.sched,
+                "vocab": args.vocab,
+                "train_hash": train_hash,
+                "val_hash": val_hash,
+                "pretrained": args.pretrained,
             }
         )
 

From 5a6be5faf921e163a765ac944fbb447fc66f7d80 Mon Sep 17 00:00:00 2001
From: fg-mindee <fg@mindee.co>
Date: Wed, 10 Nov 2021 16:36:20 +0100
Subject: [PATCH 5/9] refactor: Removed unused arguments

---
 references/detection/train_pytorch.py    | 2 +-
 references/detection/train_tensorflow.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/references/detection/train_pytorch.py b/references/detection/train_pytorch.py
index bb9f725fd9..313a3efd83 100644
--- a/references/detection/train_pytorch.py
+++ b/references/detection/train_pytorch.py
@@ -175,7 +175,7 @@ def main(args):
 
     if args.show_samples:
         x, target = next(iter(train_loader))
-        plot_samples(x, target, rotation=args.rotation)
+        plot_samples(x, target)
         return
 
     # Backbone freezing
diff --git a/references/detection/train_tensorflow.py b/references/detection/train_tensorflow.py
index 636e7141dd..5f2f2bd0d2 100644
--- a/references/detection/train_tensorflow.py
+++ b/references/detection/train_tensorflow.py
@@ -150,7 +150,7 @@ def main(args):
 
     if args.show_samples:
         x, target = next(iter(train_loader))
-        plot_samples(x, target, rotation=args.rotation)
+        plot_samples(x, target)
         return
 
     # Optimizer

From 6edca88a08d482b7e13f6d48148bb81dc609046c Mon Sep 17 00:00:00 2001
From: fg-mindee <fg@mindee.co>
Date: Wed, 10 Nov 2021 17:06:36 +0100
Subject: [PATCH 6/9] feat: Cleaned sample plotting

---
 references/classification/utils.py | 15 ++++++++++-----
 references/detection/utils.py      |  5 +++--
 references/recognition/utils.py    | 15 ++++++++++-----
 3 files changed, 23 insertions(+), 12 deletions(-)

diff --git a/references/classification/utils.py b/references/classification/utils.py
index 7e172cc583..fe776ad22a 100644
--- a/references/classification/utils.py
+++ b/references/classification/utils.py
@@ -12,8 +12,8 @@
 def plot_samples(images, targets):
     # Unnormalize image
     num_samples = min(len(images), 12)
-    num_rows = min(len(images), 3)
-    num_cols = int(math.ceil(num_samples / num_rows))
+    num_cols = min(len(images), 8)
+    num_rows = int(math.ceil(num_samples / num_cols))
     _, axes = plt.subplots(num_rows, num_cols, figsize=(20, 5))
     for idx in range(num_samples):
         img = (255 * images[idx].numpy()).round().clip(0, 255).astype(np.uint8)
@@ -23,7 +23,12 @@ def plot_samples(images, targets):
         row_idx = idx // num_cols
         col_idx = idx % num_cols
 
-        axes[row_idx][col_idx].imshow(img)
-        axes[row_idx][col_idx].axis('off')
-        axes[row_idx][col_idx].set_title(targets[idx])
+        ax = axes[row_idx] if num_rows > 1 else axes
+        ax = ax[col_idx] if num_cols > 1 else ax
+
+        ax.imshow(img)
+        ax.set_title(targets[idx])
+    # Disable axis
+    for ax in axes.ravel():
+        ax.axis('off')
     plt.show()
diff --git a/references/detection/utils.py b/references/detection/utils.py
index 05b10a0c8e..0c8d306fec 100644
--- a/references/detection/utils.py
+++ b/references/detection/utils.py
@@ -34,7 +34,8 @@ def plot_samples(images, targets: List[Dict[str, np.ndarray]]) -> None:
                 target[int(box[1]): int(box[3]) + 1, int(box[0]): int(box[2]) + 1] = 1
 
         axes[0][idx].imshow(img)
-        axes[0][idx].axis('off')
         axes[1][idx].imshow(target.astype(bool))
-        axes[1][idx].axis('off')
+    # Disable axis
+    for ax in axes.ravel():
+        ax.axis('off')
     plt.show()
diff --git a/references/recognition/utils.py b/references/recognition/utils.py
index 7e172cc583..979ce5175e 100644
--- a/references/recognition/utils.py
+++ b/references/recognition/utils.py
@@ -12,8 +12,8 @@
 def plot_samples(images, targets):
     # Unnormalize image
     num_samples = min(len(images), 12)
-    num_rows = min(len(images), 3)
-    num_cols = int(math.ceil(num_samples / num_rows))
+    num_cols = min(len(images), 4)
+    num_rows = int(math.ceil(num_samples / num_cols))
     _, axes = plt.subplots(num_rows, num_cols, figsize=(20, 5))
     for idx in range(num_samples):
         img = (255 * images[idx].numpy()).round().clip(0, 255).astype(np.uint8)
@@ -22,8 +22,13 @@ def plot_samples(images, targets):
 
         row_idx = idx // num_cols
         col_idx = idx % num_cols
+        ax = axes[row_idx] if num_rows > 1 else axes
+        ax = ax[col_idx] if num_cols > 1 else ax
+
+        ax.imshow(img)
+        ax.set_title(targets[idx])
+    # Disable axis
+    for ax in axes.ravel():
+        ax.axis('off')
 
-        axes[row_idx][col_idx].imshow(img)
-        axes[row_idx][col_idx].axis('off')
-        axes[row_idx][col_idx].set_title(targets[idx])
     plt.show()

From 0f0fcaf6a4ebd88ac31435ebbd9c26d5ab2fe934 Mon Sep 17 00:00:00 2001
From: fg-mindee <fg@mindee.co>
Date: Wed, 10 Nov 2021 17:07:34 +0100
Subject: [PATCH 7/9] feat: Added AMP support to text detection training

---
 references/detection/train_pytorch.py | 43 +++++++++++++++++++++------
 1 file changed, 34 insertions(+), 9 deletions(-)

diff --git a/references/detection/train_pytorch.py b/references/detection/train_pytorch.py
index 313a3efd83..33c5cb195c 100644
--- a/references/detection/train_pytorch.py
+++ b/references/detection/train_pytorch.py
@@ -29,7 +29,11 @@
 from doctr.utils.metrics import LocalizationConfusion
 
 
-def fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, mb):
+def fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, mb, amp=False):
+
+    if amp:
+        scaler = torch.cuda.amp.GradScaler()
+
     model.train()
     train_iter = iter(train_loader)
     # Iterate over the batches of the dataset
@@ -40,19 +44,34 @@ def fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, m
             images = images.cuda()
         images = batch_transforms(images)
 
-        train_loss = model(images, targets)['loss']
+        if amp:
+            with torch.cuda.amp.autocast():
+                train_loss = model(images, targets)['loss']
+        else:
+            train_loss = model(images, targets)['loss']
 
         optimizer.zero_grad()
-        train_loss.backward()
-        torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
-        optimizer.step()
+        if amp:
+            with torch.cuda.amp.autocast():
+                scaler.scale(train_loss).backward()
+                # Gradient clipping
+                scaler.unscale_(optimizer)
+                torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
+                # Update the params
+                scaler.step(optimizer)
+                scaler.update()
+        else:
+
+            train_loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
+            optimizer.step()
         scheduler.step()
 
         mb.child.comment = f'Training loss: {train_loss.item():.6}'
 
 
 @torch.no_grad()
-def evaluate(model, val_loader, batch_transforms, val_metric):
+def evaluate(model, val_loader, batch_transforms, val_metric, amp=False):
     # Model in eval mode
     model.eval()
     # Reset val metric
@@ -64,7 +83,11 @@ def evaluate(model, val_loader, batch_transforms, val_metric):
         if torch.cuda.is_available():
             images = images.cuda()
         images = batch_transforms(images)
-        out = model(images, targets, return_boxes=True)
+        if amp:
+            with torch.cuda.amp.autocast():
+                out = model(images, targets, return_boxes=True)
+        else:
+            out = model(images, targets, return_boxes=True)
         # Compute metric
         loc_preds, _ = out['preds']
         for boxes_gt, boxes_pred in zip(targets, loc_preds):
@@ -217,6 +240,7 @@ def main(args):
                 "val_hash": val_hash,
                 "pretrained": args.pretrained,
                 "rotation": args.rotation,
+                "amp": args.amp,
             }
         )
 
@@ -226,9 +250,9 @@ def main(args):
     # Training loop
     mb = master_bar(range(args.epochs))
     for epoch in mb:
-        fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, mb)
+        fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, mb, amp=args.amp)
         # Validation loop at the end of each epoch
-        val_loss, recall, precision, mean_iou = evaluate(model, val_loader, batch_transforms, val_metric)
+        val_loss, recall, precision, mean_iou = evaluate(model, val_loader, batch_transforms, val_metric, amp=args.amp)
         if val_loss < min_loss:
             print(f"Validation loss decreased {min_loss:.6} --> {val_loss:.6}: saving state...")
             torch.save(model.state_dict(), f"./{exp_name}.pt")
@@ -282,6 +306,7 @@ def parse_args():
     parser.add_argument('--rotation', dest='rotation', action='store_true',
                         help='train with rotated bbox')
     parser.add_argument('--sched', type=str, default='cosine', help='scheduler to use')
+    parser.add_argument("--amp", dest="amp", help="Use Automatic Mixed Precision", action="store_true")
     args = parser.parse_args()
 
     return args

From 1af2e736f3c626748aea18d180a372f33ebe46fb Mon Sep 17 00:00:00 2001
From: fg-mindee <fg@mindee.co>
Date: Wed, 10 Nov 2021 17:21:27 +0100
Subject: [PATCH 8/9] feat: Added AMP for pytorch trainings

---
 references/classification/train_pytorch.py | 36 ++++++++++++++++------
 references/detection/train_pytorch.py      | 22 ++++++-------
 references/recognition/train_pytorch.py    | 34 ++++++++++++++++----
 3 files changed, 64 insertions(+), 28 deletions(-)

diff --git a/references/classification/train_pytorch.py b/references/classification/train_pytorch.py
index 65af66c200..c22ec1fdf4 100644
--- a/references/classification/train_pytorch.py
+++ b/references/classification/train_pytorch.py
@@ -27,7 +27,11 @@
 from doctr.datasets import VOCABS, CharacterGenerator
 
 
-def fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, mb):
+def fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, mb, amp=False):
+
+    if amp:
+        scaler = torch.cuda.amp.GradScaler()
+
     model.train()
     train_iter = iter(train_loader)
     # Iterate over the batches of the dataset
@@ -36,19 +40,27 @@ def fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, m
 
         images = batch_transforms(images)
 
-        out = model(images)
-        train_loss = cross_entropy(out, targets)
-
         optimizer.zero_grad()
-        train_loss.backward()
-        optimizer.step()
+        if amp:
+            with torch.cuda.amp.autocast():
+                out = model(images)
+                train_loss = cross_entropy(out, targets)
+            scaler.scale(train_loss).backward()
+            # Update the params
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            out = model(images)
+            train_loss = cross_entropy(out, targets)
+            train_loss.backward()
+            optimizer.step()
         scheduler.step()
 
         mb.child.comment = f'Training loss: {train_loss.item():.6}'
 
 
 @torch.no_grad()
-def evaluate(model, val_loader, batch_transforms):
+def evaluate(model, val_loader, batch_transforms, amp=False):
     # Model in eval mode
     model.eval()
     # Validation loop
@@ -56,8 +68,13 @@ def evaluate(model, val_loader, batch_transforms):
     val_iter = iter(val_loader)
     for images, targets in val_iter:
         images = batch_transforms(images)
-        out = model(images)
-        loss = cross_entropy(out, targets)
+        if amp:
+            with torch.cuda.amp.autocast():
+                out = model(images)
+                loss = cross_entropy(out, targets)
+        else:
+            out = model(images)
+            loss = cross_entropy(out, targets)
         # Compute metric
         correct += (out.argmax(dim=1) == targets).sum().item()
 
@@ -250,6 +267,7 @@ def parse_args():
     parser.add_argument('--pretrained', dest='pretrained', action='store_true',
                         help='Load pretrained parameters before starting the training')
     parser.add_argument('--sched', type=str, default='cosine', help='scheduler to use')
+    parser.add_argument("--amp", dest="amp", help="Use Automatic Mixed Precision", action="store_true")
     args = parser.parse_args()
 
     return args
diff --git a/references/detection/train_pytorch.py b/references/detection/train_pytorch.py
index 33c5cb195c..64972156e6 100644
--- a/references/detection/train_pytorch.py
+++ b/references/detection/train_pytorch.py
@@ -44,27 +44,23 @@ def fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, m
             images = images.cuda()
         images = batch_transforms(images)
 
+        optimizer.zero_grad()
         if amp:
             with torch.cuda.amp.autocast():
                 train_loss = model(images, targets)['loss']
+            scaler.scale(train_loss).backward()
+            # Gradient clipping
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
+            # Update the params
+            scaler.step(optimizer)
+            scaler.update()
         else:
             train_loss = model(images, targets)['loss']
-
-        optimizer.zero_grad()
-        if amp:
-            with torch.cuda.amp.autocast():
-                scaler.scale(train_loss).backward()
-                # Gradient clipping
-                scaler.unscale_(optimizer)
-                torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
-                # Update the params
-                scaler.step(optimizer)
-                scaler.update()
-        else:
-
             train_loss.backward()
             torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
             optimizer.step()
+
         scheduler.step()
 
         mb.child.comment = f'Training loss: {train_loss.item():.6}'
diff --git a/references/recognition/train_pytorch.py b/references/recognition/train_pytorch.py
index ae0f31de00..5594ef02da 100644
--- a/references/recognition/train_pytorch.py
+++ b/references/recognition/train_pytorch.py
@@ -30,7 +30,11 @@
 from doctr.utils.metrics import TextMatch
 
 
-def fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, mb):
+def fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, mb, amp=False):
+
+    if amp:
+        scaler = torch.cuda.amp.GradScaler()
+
     model.train()
     train_iter = iter(train_loader)
     # Iterate over the batches of the dataset
@@ -44,16 +48,29 @@ def fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, m
         train_loss = model(images, targets)['loss']
 
         optimizer.zero_grad()
-        train_loss.backward()
-        optimizer.step()
-        torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
+        if amp:
+            with torch.cuda.amp.autocast():
+                train_loss = model(images, targets)['loss']
+            scaler.scale(train_loss).backward()
+            # Gradient clipping
+            scaler.unscale_(optimizer)
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
+            # Update the params
+            scaler.step(optimizer)
+            scaler.update()
+        else:
+            train_loss = model(images, targets)['loss']
+            train_loss.backward()
+            torch.nn.utils.clip_grad_norm_(model.parameters(), 5)
+            optimizer.step()
+
         scheduler.step()
 
         mb.child.comment = f'Training loss: {train_loss.item():.6}'
 
 
 @torch.no_grad()
-def evaluate(model, val_loader, batch_transforms, val_metric):
+def evaluate(model, val_loader, batch_transforms, val_metric, amp=False):
     # Model in eval mode
     model.eval()
     # Reset val metric
@@ -65,7 +82,11 @@ def evaluate(model, val_loader, batch_transforms, val_metric):
         if torch.cuda.is_available():
             images = images.cuda()
         images = batch_transforms(images)
-        out = model(images, targets, return_preds=True)
+        if amp:
+            with torch.cuda.amp.autocast():
+                out = model(images, targets, return_preds=True)
+        else:
+            out = model(images, targets, return_preds=True)
         # Compute metric
         if len(out['preds']):
             words, _ = zip(*out['preds'])
@@ -278,6 +299,7 @@ def parse_args():
     parser.add_argument('--pretrained', dest='pretrained', action='store_true',
                         help='Load pretrained parameters before starting the training')
     parser.add_argument('--sched', type=str, default='cosine', help='scheduler to use')
+    parser.add_argument("--amp", dest="amp", help="Use Automatic Mixed Precision", action="store_true")
     args = parser.parse_args()
 
     return args

From a4de6c45d2daa8972a33d2a5c38cd2609d0292a5 Mon Sep 17 00:00:00 2001
From: fg-mindee <fg@mindee.co>
Date: Wed, 10 Nov 2021 17:24:54 +0100
Subject: [PATCH 9/9] fix: Fixed compatibility with FP16

---
 doctr/models/detection/core.py | 2 +-
 doctr/utils/geometry.py        | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/doctr/models/detection/core.py b/doctr/models/detection/core.py
index 8df10fc064..973d4ffd34 100644
--- a/doctr/models/detection/core.py
+++ b/doctr/models/detection/core.py
@@ -98,7 +98,7 @@ def __call__(
 
         for p_, bitmap_ in zip(proba_map, bitmap):
             # Perform opening (erosion + dilatation)
-            bitmap_ = cv2.morphologyEx(bitmap_, cv2.MORPH_OPEN, kernel)
+            bitmap_ = cv2.morphologyEx(bitmap_.astype(np.float32), cv2.MORPH_OPEN, kernel)
             # Rotate bitmap and proba_map
             angle = get_bitmap_angle(bitmap_)
             angles_batch.append(angle)
diff --git a/doctr/utils/geometry.py b/doctr/utils/geometry.py
index dc5558db84..de79091709 100644
--- a/doctr/utils/geometry.py
+++ b/doctr/utils/geometry.py
@@ -208,7 +208,7 @@ def rotate_image(
 
     height, width = exp_img.shape[:2]
     rot_mat = cv2.getRotationMatrix2D((width / 2, height / 2), angle, 1.0)
-    rot_img = cv2.warpAffine(exp_img, rot_mat, (width, height))
+    rot_img = cv2.warpAffine(exp_img.astype(np.float32), rot_mat, (width, height))
     if expand:
         # Pad to get the same aspect ratio
         if (image.shape[0] / image.shape[1]) != (rot_img.shape[0] / rot_img.shape[1]):