feat: ✨ use tqdm instead of fastprogress

pyproject.toml

@@ -179,7 +179,7 @@ preview=true
 
 [tool.ruff.isort]
 known-first-party = ["doctr", "app", "utils"]
-known-third-party = ["tensorflow", "torch", "torchvision", "wandb", "fastprogress", "fastapi"]
+known-third-party = ["tensorflow", "torch", "torchvision", "wandb", "tqdm", "fastapi"]
 
 [tool.ruff.per-file-ignores]
 "doctr/models/**.py" = ["N806", "F841"]

references/classification/train_pytorch.py

@@ -15,7 +15,6 @@
 import numpy as np
 import torch
 import wandb
-from fastprogress.fastprogress import master_bar, progress_bar
 from torch.nn.functional import cross_entropy
 from torch.optim.lr_scheduler import CosineAnnealingLR, MultiplicativeLR, OneCycleLR
 from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
@@ -29,6 +28,7 @@
     RandomPhotometricDistort,
     RandomRotation,
 )
+from tqdm.auto import tqdm, trange
 
 from doctr import transforms as T
 from doctr.datasets import VOCABS, CharacterGenerator
@@ -107,13 +107,14 @@ def record_lr(
     return lr_recorder[: len(loss_recorder)], loss_recorder
 
 
-def fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, mb, amp=False):
+def fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, amp=False):
     if amp:
         scaler = torch.cuda.amp.GradScaler()
 
     model.train()
     # Iterate over the batches of the dataset
-    for images, targets in progress_bar(train_loader, parent=mb):
+    pbar = tqdm(train_loader, position=1)
+    for images, targets in pbar:
         if torch.cuda.is_available():
             images = images.cuda()
             targets = targets.cuda()
@@ -136,7 +137,7 @@ def fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, m
             optimizer.step()
         scheduler.step()
 
-        mb.child.comment = f"Training loss: {train_loss.item():.6}"
+        pbar.set_description(f"Training loss: {train_loss.item():.6}")
 
 
 @torch.no_grad()
@@ -329,9 +330,9 @@ def main(args):
     # Create loss queue
     min_loss = np.inf
     # Training loop
-    mb = master_bar(range(args.epochs))
+    mb = trange(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)
 
         # Validation loop at the end of each epoch
         val_loss, acc = evaluate(model, val_loader, batch_transforms)

references/classification/train_tensorflow.py

@@ -14,8 +14,8 @@
 
 import numpy as np
 import tensorflow as tf
-from fastprogress.fastprogress import master_bar, progress_bar
 from tensorflow.keras import mixed_precision
+from tqdm.auto import tqdm, trange
 
 from doctr.models import login_to_hub, push_to_hf_hub
 
@@ -83,9 +83,10 @@ def record_lr(
     return lr_recorder[: len(loss_recorder)], loss_recorder
 
 
-def fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb, amp=False):
+def fit_one_epoch(model, train_loader, batch_transforms, optimizer, amp=False):
     # Iterate over the batches of the dataset
-    for images, targets in progress_bar(train_loader, parent=mb):
+    pbar = tqdm(train_loader, position=1)
+    for images, targets in pbar:
         images = batch_transforms(images)
 
         with tf.GradientTape() as tape:
@@ -96,7 +97,7 @@ def fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb, amp=Fals
             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}"
+        pbar.set_description(f"Training loss: {train_loss.numpy().mean():.6}")
 
 
 def evaluate(model, val_loader, batch_transforms):
@@ -293,9 +294,9 @@ def main(args):
     min_loss = np.inf
 
     # Training loop
-    mb = master_bar(range(args.epochs))
+    mb = trange(args.epochs)
     for epoch in mb:
-        fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb, args.amp)
+        fit_one_epoch(model, train_loader, batch_transforms, optimizer, args.amp)
 
         # Validation loop at the end of each epoch
         val_loss, acc = evaluate(model, val_loader, batch_transforms)

references/detection/train_pytorch.py

@@ -17,10 +17,10 @@
 import psutil
 import torch
 import wandb
-from fastprogress.fastprogress import master_bar, progress_bar
 from torch.optim.lr_scheduler import CosineAnnealingLR, MultiplicativeLR, OneCycleLR
 from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
 from torchvision.transforms.v2 import Compose, GaussianBlur, Normalize, RandomGrayscale, RandomPhotometricDistort
+from tqdm.auto import tqdm, trange
 
 from doctr import transforms as T
 from doctr.datasets import DetectionDataset
@@ -100,13 +100,14 @@ def record_lr(
     return lr_recorder[: len(loss_recorder)], loss_recorder
 
 
-def fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, mb, amp=False):
+def fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, amp=False):
     if amp:
         scaler = torch.cuda.amp.GradScaler()
 
     model.train()
     # Iterate over the batches of the dataset
-    for images, targets in progress_bar(train_loader, parent=mb):
+    pbar = tqdm(train_loader, position=1)
+    for images, targets in pbar:
         if torch.cuda.is_available():
             images = images.cuda()
         images = batch_transforms(images)
@@ -130,7 +131,7 @@ def fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, m
 
         scheduler.step()
 
-        mb.child.comment = f"Training loss: {train_loss.item():.6}"
+        pbar.set_description(f"Training loss: {train_loss.item():.6}")
 
 
 @torch.no_grad()
@@ -368,9 +369,9 @@ def main(args):
     min_loss = np.inf
 
     # Training loop
-    mb = master_bar(range(args.epochs))
+    mb = trange(args.epochs)
     for epoch in mb:
-        fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, mb, amp=args.amp)
+        fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, 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, amp=args.amp)
         if val_loss < min_loss:

references/detection/train_tensorflow.py

@@ -16,8 +16,8 @@
 import numpy as np
 import psutil
 import tensorflow as tf
-from fastprogress.fastprogress import master_bar, progress_bar
 from tensorflow.keras import mixed_precision
+from tqdm.auto import tqdm, trange
 
 from doctr.models import login_to_hub, push_to_hf_hub
 
@@ -84,10 +84,11 @@ def record_lr(
     return lr_recorder[: len(loss_recorder)], loss_recorder
 
 
-def fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb, amp=False):
+def fit_one_epoch(model, train_loader, batch_transforms, optimizer, amp=False):
     train_iter = iter(train_loader)
     # Iterate over the batches of the dataset
-    for images, targets in progress_bar(train_iter, parent=mb):
+    pbar = tqdm(train_iter, position=1)
+    for images, targets in pbar:
         images = batch_transforms(images)
 
         with tf.GradientTape() as tape:
@@ -97,7 +98,7 @@ def fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb, amp=Fals
             grads = optimizer.get_unscaled_gradients(grads)
         optimizer.apply_gradients(zip(grads, model.trainable_weights))
 
-        mb.child.comment = f"Training loss: {train_loss.numpy():.6}"
+        pbar.set_description(f"Training loss: {train_loss.numpy():.6}")
 
 
 def evaluate(model, val_loader, batch_transforms, val_metric):
@@ -327,9 +328,9 @@ def main(args):
     min_loss = np.inf
 
     # Training loop
-    mb = master_bar(range(args.epochs))
+    mb = trange(args.epochs)
     for epoch in mb:
-        fit_one_epoch(model, train_loader, batch_transforms, optimizer, mb, args.amp)
+        fit_one_epoch(model, train_loader, batch_transforms, optimizer, 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:

references/obj_detection/train_pytorch.py

@@ -16,10 +16,10 @@
 import torch
 import torch.optim as optim
 import wandb
-from fastprogress.fastprogress import master_bar, progress_bar
 from torch.optim.lr_scheduler import MultiplicativeLR, StepLR
 from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
 from torchvision.transforms import ColorJitter, Compose, GaussianBlur
+from tqdm.auto import tqdm, trange
 
 from doctr import transforms as T
 from doctr.datasets import DocArtefacts
@@ -113,13 +113,14 @@ def convert_to_abs_coords(targets, img_shape):
     return targets
 
 
-def fit_one_epoch(model, train_loader, optimizer, scheduler, mb, amp=False):
+def fit_one_epoch(model, train_loader, optimizer, scheduler, amp=False):
     if amp:
         scaler = torch.cuda.amp.GradScaler()
 
     model.train()
     # Iterate over the batches of the dataset
-    for images, targets in progress_bar(train_loader, parent=mb):
+    pbar = tqdm(train_loader, position=1)
+    for images, targets in pbar:
         targets = convert_to_abs_coords(targets, images.shape)
         if torch.cuda.is_available():
             images = images.cuda()
@@ -140,7 +141,7 @@ def fit_one_epoch(model, train_loader, optimizer, scheduler, mb, amp=False):
             loss.backward()
             optimizer.step()
 
-        mb.child.comment = f"Training loss: {loss.item()}"
+        pbar.set_description(f"Training loss: {loss.item()}")
     scheduler.step()
 
 
@@ -303,11 +304,11 @@ def main(args):
             },
         )
 
-    mb = master_bar(range(args.epochs))
+    mb = trange(args.epochs)
     max_score = 0.0
 
     for epoch in mb:
-        fit_one_epoch(model, train_loader, optimizer, scheduler, mb, amp=args.amp)
+        fit_one_epoch(model, train_loader, optimizer, scheduler, amp=args.amp)
         # Validation loop at the end of each epoch
         recall, precision, mean_iou = evaluate(model, val_loader, metric, amp=args.amp)
         f1_score = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0.0

references/recognition/train_pytorch.py

@@ -17,7 +17,6 @@
 import numpy as np
 import torch
 import wandb
-from fastprogress.fastprogress import master_bar, progress_bar
 from torch.optim.lr_scheduler import CosineAnnealingLR, MultiplicativeLR, OneCycleLR
 from torch.utils.data import DataLoader, RandomSampler, SequentialSampler
 from torchvision.transforms.v2 import (
@@ -28,6 +27,7 @@
     RandomPerspective,
     RandomPhotometricDistort,
 )
+from tqdm.auto import tqdm, trange
 
 from doctr import transforms as T
 from doctr.datasets import VOCABS, RecognitionDataset, WordGenerator
@@ -107,13 +107,14 @@ def record_lr(
     return lr_recorder[: len(loss_recorder)], loss_recorder
 
 
-def fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, mb, amp=False):
+def fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, amp=False):
     if amp:
         scaler = torch.cuda.amp.GradScaler()
 
     model.train()
     # Iterate over the batches of the dataset
-    for images, targets in progress_bar(train_loader, parent=mb):
+    pbar = tqdm(train_loader, position=1)
+    for images, targets in pbar:
         if torch.cuda.is_available():
             images = images.cuda()
         images = batch_transforms(images)
@@ -139,7 +140,7 @@ def fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, m
 
         scheduler.step()
 
-        mb.child.comment = f"Training loss: {train_loss.item():.6}"
+        pbar.set_description(f"Training loss: {train_loss.item():.6}")
 
 
 @torch.no_grad()
@@ -391,9 +392,9 @@ def main(args):
     # Create loss queue
     min_loss = np.inf
     # Training loop
-    mb = master_bar(range(args.epochs))
+    mb = trange(args.epochs)
     for epoch in mb:
-        fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, mb, amp=args.amp)
+        fit_one_epoch(model, train_loader, batch_transforms, optimizer, scheduler, amp=args.amp)
 
         # Validation loop at the end of each epoch
         val_loss, exact_match, partial_match = evaluate(model, val_loader, batch_transforms, val_metric, amp=args.amp)

references/recognition/train_pytorch_ddp.py

@@ -18,7 +18,6 @@
 import torch.distributed as dist
 import torch.multiprocessing as mp
 import wandb
-from fastprogress.fastprogress import master_bar, progress_bar
 from torch.nn.parallel import DistributedDataParallel as DDP
 from torch.optim.lr_scheduler import CosineAnnealingLR, OneCycleLR
 from torch.utils.data import DataLoader, SequentialSampler
@@ -31,6 +30,7 @@
     RandomPerspective,
     RandomPhotometricDistort,
 )
+from tqdm.auto import tqdm, trange
 
 from doctr import transforms as T
 from doctr.datasets import VOCABS, RecognitionDataset, WordGenerator
@@ -39,13 +39,14 @@
 from utils import plot_samples
 
 
-def fit_one_epoch(model, device, train_loader, batch_transforms, optimizer, scheduler, mb, amp=False):
+def fit_one_epoch(model, device, train_loader, batch_transforms, optimizer, scheduler, amp=False):
     if amp:
         scaler = torch.cuda.amp.GradScaler()
 
     model.train()
     # Iterate over the batches of the dataset
-    for images, targets in progress_bar(train_loader, parent=mb):
+    pbar = tqdm(train_loader, position=1)
+    for images, targets in pbar:
         images = images.to(device)
         images = batch_transforms(images)
 
@@ -70,7 +71,7 @@ def fit_one_epoch(model, device, train_loader, batch_transforms, optimizer, sche
 
         scheduler.step()
 
-        mb.child.comment = f"Training loss: {train_loss.item():.6}"
+        pbar.set_description(f"Training loss: {train_loss.item():.6}")
 
 
 @torch.no_grad()
@@ -321,9 +322,9 @@ def main(rank: int, world_size: int, args):
     # Create loss queue
     min_loss = np.inf
     # Training loop
-    mb = master_bar(range(args.epochs))
+    mb = trange(args.epochs)
     for epoch in mb:
-        fit_one_epoch(model, device, train_loader, batch_transforms, optimizer, scheduler, mb, amp=args.amp)
+        fit_one_epoch(model, device, train_loader, batch_transforms, optimizer, scheduler, amp=args.amp)
 
         if rank == 0:
             # Validation loop at the end of each epoch