code: def main(args):
      print(args)
  
      if args.push_to_hub:
          login_to_hub()
  
      if not isinstance(args.workers, int):
          args.workers = min(16, mp.cpu_count())
  
      torch.backends.cudnn.benchmark = True
  
      vocab = VOCABS[args.vocab]
      fonts = args.font.split(",")
  
      # Load val data generator
      st = time.time()
      if isinstance(args.val_path, str):
          with open(
              os.path.join(args.val_path, "labels.json"), "rb"
          ) as f:
              val_hash = hashlib.sha256(f.read()).hexdigest()
  
          val_set = RecognitionDataset(
              img_folder=os.path.join(args.val_path, "images"),
              labels_path=os.path.join(args.val_path, "labels.json"),
              img_transforms=T.Resize(
                  (args.input_size, 4 * args.input_size),
                  preserve_aspect_ratio=True,
              ),
          )
      else:
          val_hash = None
          # Load synthetic data generator
          val_set = WordGenerator(
              vocab=vocab,
              min_chars=args.min_chars,
              max_chars=args.max_chars,
              num_samples=args.val_samples * len(vocab),
              font_family=fonts,
              img_transforms=Compose(
                  [
                      T.Resize(
                          (args.input_size, 4 * args.input_size),
                          preserve_aspect_ratio=True,
                      ),
                      # Ensure we have a 90% split of white-background images
                      T.RandomApply(T.ColorInversion(), 0.9),
                  ]
              ),
          )
  
      val_loader = DataLoader(
          val_set,
          batch_size=args.batch_size,
          drop_last=False,
          num_workers=args.workers,
          sampler=SequentialSampler(val_set),
          pin_memory=torch.cuda.is_available(),
          collate_fn=val_set.collate_fn,
      )
      print(
          f"Validation set loaded in {time.time() - st:.4}s ({len(val_set)} samples in "
          f"{len(val_loader)} batches)"
      )
  
      batch_transforms = Normalize(
          mean=(0.694, 0.695, 0.693), std=(0.299, 0.296, 0.301)
      )
  
      # Load doctr model
      model = recognition.__dict__[args.arch](
          pretrained=args.pretrained, vocab=vocab
      )
  
      # Resume weights
      if isinstance(args.resume, str):
          print(f"Resuming {args.resume}")
          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()
  
      if args.test_only:
          print("Running evaluation")
          val_loss, exact_match, partial_match = evaluate(
              model,
              val_loader,
              batch_transforms,
              val_metric,
              amp=args.amp,
          )
          print(
              f"Validation loss: {val_loss:.6} (Exact: {exact_match:.2%} | Partial: {partial_match:.2%})"
          )
          return
  
      st = time.time()
  
      if isinstance(args.train_path, str):
          # Load train data generator
          base_path = Path(args.train_path)
          parts = (
              [base_path]
              if base_path.joinpath("labels.json").is_file()
              else [
                  base_path.joinpath(sub)
                  for sub in os.listdir(base_path)
              ]
          )
          with open(parts[0].joinpath("labels.json"), "rb") as f:
              train_hash = hashlib.sha256(f.read()).hexdigest()
  
          train_set = RecognitionDataset(
              parts[0].joinpath("images"),
              parts[0].joinpath("labels.json"),
              img_transforms=Compose(
                  [
                      T.Resize(
                          (args.input_size, 4 * args.input_size),
                          preserve_aspect_ratio=True,
                      ),
                      # Augmentations
                      T.RandomApply(T.ColorInversion(), 0.1),
                      RandomGrayscale(p=0.1),
                      RandomPhotometricDistort(p=0.1),
                      T.RandomApply(T.RandomShadow(), p=0.4),
                      T.RandomApply(
                          T.GaussianNoise(mean=0, std=0.1), 0.1
                      ),
                      T.RandomApply(GaussianBlur(3), 0.3),
                      RandomPerspective(distortion_scale=0.2, p=0.3),
                  ]
              ),
          )
          if len(parts) > 1:
              for subfolder in parts[1:]:
                  train_set.merge_dataset(
                      RecognitionDataset(
                          subfolder.joinpath("images"),
                          subfolder.joinpath("labels.json"),
                      )
                  )
      else:
          train_hash = None
          # Load synthetic data generator
          train_set = WordGenerator(
              vocab=vocab,
              min_chars=args.min_chars,
              max_chars=args.max_chars,
              num_samples=args.train_samples * len(vocab),
              font_family=fonts,
              img_transforms=Compose(
                  [
                      T.Resize(
                          (args.input_size, 4 * args.input_size),
                          preserve_aspect_ratio=True,
                      ),
                      # Ensure we have a 90% split of white-background images
                      T.RandomApply(T.ColorInversion(), 0.9),
                      RandomGrayscale(p=0.1),
                      RandomPhotometricDistort(p=0.1),
                      T.RandomApply(T.RandomShadow(), p=0.4),
                      T.RandomApply(
                          T.GaussianNoise(mean=0, std=0.1), 0.1
                      ),
                      T.RandomApply(GaussianBlur(3), 0.3),
                      RandomPerspective(distortion_scale=0.2, p=0.3),
                  ]
              ),
          )
  
      train_loader = DataLoader(
          train_set,
          batch_size=args.batch_size,
          drop_last=True,
          num_workers=args.workers,
          sampler=RandomSampler(train_set),
          pin_memory=torch.cuda.is_available(),
          collate_fn=train_set.collate_fn,
      )
      print(
          f"Train set loaded in {time.time() - st:.4}s ({len(train_set)} samples in "
          f"{len(train_loader)} batches)"
      )
  
      if args.show_samples:
          x, target = next(iter(train_loader))
          plot_samples(x, target)
          return
  
      # Backbone freezing
      if args.freeze_backbone:
          for p in model.feat_extractor.parameters():
              p.requires_grad = False
  
      # Optimizer
      optimizer = torch.optim.Adam(
          [p for p in model.parameters() if p.requires_grad],
          args.lr,
          betas=(0.95, 0.99),
          eps=1e-6,
          weight_decay=args.weight_decay,
      )
      # LR Finder
      if args.find_lr:
          lrs, losses = record_lr(
              model,
              train_loader,
              batch_transforms,
              optimizer,
              amp=args.amp,
          )
          plot_recorder(lrs, losses)
          return
      # Scheduler
      if args.sched == "cosine":
          scheduler = CosineAnnealingLR(
              optimizer,
              args.epochs * len(train_loader),
              eta_min=args.lr / 25e4,
          )
      elif args.sched == "onecycle":
          scheduler = OneCycleLR(
              optimizer, args.lr, args.epochs * len(train_loader)
          )
  
      # Training monitoring
      current_time = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
      exp_name = (
          f"{args.arch}_{current_time}"
          if args.name is None
          else args.name
      )
  
      # W&B
      if args.wb:
          run = wandb.init(
              name=exp_name,
              project="text-recognition",
              config={
                  "learning_rate": args.lr,
                  "epochs": args.epochs,
                  "weight_decay": args.weight_decay,
                  "batch_size": args.batch_size,
                  "architecture": args.arch,
                  "input_size": args.input_size,
                  "optimizer": "adam",
                  "framework": "pytorch",
                  "scheduler": args.sched,
                  "vocab": args.vocab,
                  "train_hash": train_hash,
                  "val_hash": val_hash,
                  "pretrained": args.pretrained,
              },
          )
  
      # Create loss queue
      min_loss = np.inf
      # Training loop
      if args.early_stop:
          early_stopper = EarlyStopper(
              patience=args.early_stop_epochs,
              min_delta=args.early_stop_delta,
          )
      for epoch in range(args.epochs):
          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,
          )
          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")
              min_loss = val_loss
          print(
              f"Epoch {epoch + 1}/{args.epochs} - Validation loss: {val_loss:.6} "
              f"(Exact: {exact_match:.2%} | Partial: {partial_match:.2%})"
          )
          # W&B
          if args.wb:
              wandb.log(
                  {
                      "val_loss": val_loss,
                      "exact_match": exact_match,
                      "partial_match": partial_match,
                  }
              )
          if args.early_stop and early_stopper.early_stop(val_loss):
              print(
                  "Training halted early due to reaching patience limit."
              )
              break
      if args.wb:
          run.finish()
  
      if args.push_to_hub:
          push_to_hf_hub(
              model, exp_name, task="recognition", run_config=args
          )
  
imports: import os
  import datetime
  import hashlib
  import logging
  import multiprocessing
  import time
  from pathlib import Path
  import numpy
  import torch
  import wandb
  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, RandomPerspective, RandomPhotometricDistort
  from tqdm.auto import tqdm
  from doctr import transforms
  from doctr.datasets import VOCABS, RecognitionDataset, WordGenerator
  from doctr.models import login_to_hub, push_to_hf_hub, recognition
  from doctr.utils.metrics import TextMatch
  from utils import EarlyStopper, plot_recorder, plot_samples
label_desc: Use image node_id 4 for calling a global function with example usage: main(args) without return types