Repository: mindee/doctr

Issue:
[references] Add FP16 support for training
The reference training script should have an option to switch from a FP32 training to a FP16 training.
It would raise a question: how to harmonize model loading from both FP?

A few suggestions:
- simply cast the FP16 checkpoint to FP32 before hashing + uploading
- for a given model, each checkpoint is marked with its FP. The user may choose how the model needs to be instantiated: by default the output FP is FP32, but this can be changed by the user. Depending on the checkpoint FP, it will be cast to its counterpart FP for loading.

The second option would be cleaner and avoid unnecessary large checkpoints when they can be kept in FP16

Here is a proposition:
- [x] add support of FP16 to `doctr.datasets` (#367)
- [x] add support of FP16 to `doctr.models` (#382)
- [x] add support of FP16 to `doctr.transforms` (#388)
- [x] add support of FP16 to `references` with PyTorch (#604)
- [x] add support of FP16 to `references` with TensorFlow (#682)


Discussion:

Resolving patch:
diff --git a/references/classification/train_tensorflow.py b/references/classification/train_tensorflow.py
--- a/references/classification/train_tensorflow.py
+++ b/references/classification/train_tensorflow.py
@@ -16,6 +16,7 @@
 import tensorflow as tf
 import wandb
 from fastprogress.fastprogress import master_bar, progress_bar
+from tensorflow.keras import mixed_precision
 
 gpu_devices = tf.config.experimental.list_physical_devices('GPU')
 if any(gpu_devices):
@@ -27,7 +28,7 @@
 from utils import plot_samples
 
 
-def fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb):
+def fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb, amp=False):
     train_iter = iter(train_loader)
     # Iterate over the batches of the dataset
     for _ in progress_bar(range(train_loader.num_batches), parent=mb):
@@ -39,6 +40,8 @@ def fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb):
             out = model(images, training=True)
             train_loss = tf.nn.sparse_softmax_cross_entropy_with_logits(targets, out)
         grads = tape.gradient(train_loss, model.trainable_weights)
+        if amp:
+            grads = optimizer.get_unscaled_gradients(grads)
         optimizer.apply_gradients(zip(grads, model.trainable_weights))
 
         mb.child.comment = f'Training loss: {train_loss.numpy().mean():.6}'
@@ -81,6 +84,10 @@ def main(args):
 
     vocab = VOCABS[args.vocab]
 
+    # AMP
+    if args.amp:
+        mixed_precision.set_global_policy('mixed_float16')
+
     # Load val data generator
     st = time.time()
     val_set = CharacterGenerator(
@@ -108,6 +115,7 @@ def main(args):
         num_classes=len(vocab),
         include_top=True,
     )
+
     # Resume weights
     if isinstance(args.resume, str):
         model.load_weights(args.resume)
@@ -169,6 +177,8 @@ def main(args):
         beta_2=0.99,
         epsilon=1e-6,
     )
+    if args.amp:
+        optimizer = mixed_precision.LossScaleOptimizer(optimizer)
 
     # Tensorboard to monitor training
     current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
@@ -201,7 +211,7 @@ def main(args):
     # Training loop
     mb = master_bar(range(args.epochs))
     for epoch in mb:
-        fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb)
+        fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb, args.amp)
 
         # Validation loop at the end of each epoch
         val_loss, acc = evaluate(model, val_loader, batch_transforms)
@@ -257,6 +267,7 @@ def parse_args():
                         help='Log to Weights & Biases')
     parser.add_argument('--pretrained', dest='pretrained', action='store_true',
                         help='Load pretrained parameters before starting the training')
+    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_tensorflow.py b/references/detection/train_tensorflow.py
--- a/references/detection/train_tensorflow.py
+++ b/references/detection/train_tensorflow.py
@@ -17,6 +17,7 @@
 import tensorflow as tf
 import wandb
 from fastprogress.fastprogress import master_bar, progress_bar
+from tensorflow.keras import mixed_precision
 
 gpu_devices = tf.config.experimental.list_physical_devices('GPU')
 if any(gpu_devices):
@@ -29,7 +30,7 @@
 from utils import plot_samples
 
 
-def fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb):
+def fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb, amp=False):
     train_iter = iter(train_loader)
     # Iterate over the batches of the dataset
     for images, targets in progress_bar(train_iter, parent=mb):
@@ -39,6 +40,8 @@ def fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb):
         with tf.GradientTape() as tape:
             train_loss = model(images, targets, training=True)['loss']
         grads = tape.gradient(train_loss, model.trainable_weights)
+        if amp:
+            grads = optimizer.get_unscaled_gradients(grads)
         optimizer.apply_gradients(zip(grads, model.trainable_weights))
 
         mb.child.comment = f'Training loss: {train_loss.numpy():.6}'
@@ -74,6 +77,10 @@ def main(args):
     if not isinstance(args.workers, int):
         args.workers = min(16, mp.cpu_count())
 
+    # AMP
+    if args.amp:
+        mixed_precision.set_global_policy('mixed_float16')
+
     st = time.time()
     val_set = DetectionDataset(
         img_folder=os.path.join(args.val_path, 'images'),
@@ -151,6 +158,8 @@ def main(args):
         epsilon=1e-6,
         clipnorm=5
     )
+    if args.amp:
+        optimizer = mixed_precision.LossScaleOptimizer(optimizer)
 
     # Tensorboard to monitor training
     current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
@@ -188,7 +197,7 @@ def main(args):
     # Training loop
     mb = master_bar(range(args.epochs))
     for epoch in mb:
-        fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb)
+        fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb, args.amp)
         # Validation loop at the end of each epoch
         val_loss, recall, precision, mean_iou = evaluate(model, val_loader, batch_transforms, val_metric)
         if val_loss < min_loss:
@@ -240,6 +249,7 @@ def parse_args():
                         help='Load pretrained parameters before starting the training')
     parser.add_argument('--rotation', dest='rotation', action='store_true',
                         help='train with rotated bbox')
+    parser.add_argument("--amp", dest="amp", help="Use Automatic Mixed Precision", action="store_true")
     args = parser.parse_args()
 
     return args
diff --git a/references/recognition/train_tensorflow.py b/references/recognition/train_tensorflow.py
--- a/references/recognition/train_tensorflow.py
+++ b/references/recognition/train_tensorflow.py
@@ -18,6 +18,7 @@
 import tensorflow as tf
 import wandb
 from fastprogress.fastprogress import master_bar, progress_bar
+from tensorflow.keras import mixed_precision
 
 gpu_devices = tf.config.experimental.list_physical_devices('GPU')
 if any(gpu_devices):
@@ -30,7 +31,7 @@
 from utils import plot_samples
 
 
-def fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb):
+def fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb, amp=False):
     train_iter = iter(train_loader)
     # Iterate over the batches of the dataset
     for images, targets in progress_bar(train_iter, parent=mb):
@@ -40,6 +41,8 @@ def fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb):
         with tf.GradientTape() as tape:
             train_loss = model(images, targets, training=True)['loss']
         grads = tape.gradient(train_loss, model.trainable_weights)
+        if amp:
+            grads = optimizer.get_unscaled_gradients(grads)
         optimizer.apply_gradients(zip(grads, model.trainable_weights))
 
         mb.child.comment = f'Training loss: {train_loss.numpy().mean():.6}'
@@ -76,6 +79,10 @@ def main(args):
     if not isinstance(args.workers, int):
         args.workers = min(16, mp.cpu_count())
 
+    # AMP
+    if args.amp:
+        mixed_precision.set_global_policy('mixed_float16')
+
     # Load val data generator
     st = time.time()
     val_set = RecognitionDataset(
@@ -162,6 +169,8 @@ def main(args):
         epsilon=1e-6,
         clipnorm=5
     )
+    if args.amp:
+        optimizer = mixed_precision.LossScaleOptimizer(optimizer)
 
     # Tensorboard to monitor training
     current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
@@ -195,7 +204,7 @@ def main(args):
     # Training loop
     mb = master_bar(range(args.epochs))
     for epoch in mb:
-        fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb)
+        fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb, args.amp)
 
         # Validation loop at the end of each epoch
         val_loss, exact_match, partial_match = evaluate(model, val_loader, batch_transforms, val_metric)
@@ -240,6 +249,7 @@ def parse_args():
                         help='Log to Weights & Biases')
     parser.add_argument('--pretrained', dest='pretrained', action='store_true',
                         help='Load pretrained parameters before starting the training')
+    parser.add_argument("--amp", dest="amp", help="Use Automatic Mixed Precision", action="store_true")
     args = parser.parse_args()
 
     return args


Test patch: