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())
  
      vocab = VOCABS[args.vocab]
  
      fonts = args.font.split(",")
  
      # AMP
      if args.amp:
          mixed_precision.set_global_policy("mixed_float16")
  
      # Load val data generator
      st = time.time()
      val_set = CharacterGenerator(
          vocab=vocab,
          num_samples=args.val_samples * len(vocab),
          cache_samples=True,
          img_transforms=T.Compose(
              [
                  T.Resize((args.input_size, args.input_size)),
                  # Ensure we have a 90% split of white-background images
                  T.RandomApply(T.ColorInversion(), 0.9),
              ]
          ),
          font_family=fonts,
      )
      val_loader = DataLoader(
          val_set,
          batch_size=args.batch_size,
          shuffle=False,
          drop_last=False,
          num_workers=args.workers,
          collate_fn=collate_fn,
      )
      print(
          f"Validation set loaded in {time.time() - st:.4}s ({len(val_set)} samples in "
          f"{val_loader.num_batches} batches)"
      )
  
      # Load doctr model
      model = classification.__dict__[args.arch](
          pretrained=args.pretrained,
          input_shape=(args.input_size, args.input_size, 3),
          num_classes=len(vocab),
          classes=list(vocab),
          include_top=True,
      )
  
      # Resume weights
      if isinstance(args.resume, str):
          model.load_weights(args.resume)
  
      batch_transforms = T.Compose(
          [
              T.Normalize(
                  mean=(0.694, 0.695, 0.693), std=(0.299, 0.296, 0.301)
              ),
          ]
      )
  
      if args.test_only:
          print("Running evaluation")
          val_loss, acc = evaluate(model, val_loader, batch_transforms)
          print(f"Validation loss: {val_loss:.6} (Acc: {acc:.2%})")
          return
  
      st = time.time()
  
      # Load train data generator
      train_set = CharacterGenerator(
          vocab=vocab,
          num_samples=args.train_samples * len(vocab),
          cache_samples=True,
          img_transforms=T.Compose(
              [
                  T.Resize((args.input_size, args.input_size)),
                  # Augmentations
                  T.RandomApply(T.ColorInversion(), 0.9),
                  T.RandomApply(T.ToGray(3), 0.1),
                  T.RandomJpegQuality(60),
                  T.RandomSaturation(0.3),
                  T.RandomContrast(0.3),
                  T.RandomBrightness(0.3),
                  # Blur
                  T.RandomApply(
                      T.GaussianBlur(kernel_shape=(3, 3), std=(0.1, 3)),
                      0.3,
                  ),
              ]
          ),
          font_family=fonts,
      )
      train_loader = DataLoader(
          train_set,
          batch_size=args.batch_size,
          shuffle=True,
          drop_last=True,
          num_workers=args.workers,
          collate_fn=collate_fn,
      )
      print(
          f"Train set loaded in {time.time() - st:.4}s ({len(train_set)} samples in "
          f"{train_loader.num_batches} batches)"
      )
  
      if args.show_samples:
          x, target = next(iter(train_loader))
          plot_samples(x, list(map(vocab.__getitem__, target)))
          return
  
      # Optimizer
      scheduler = tf.keras.optimizers.schedules.ExponentialDecay(
          args.lr,
          decay_steps=args.epochs * len(train_loader),
          decay_rate=1 / (1e3),  # final lr as a fraction of initial lr
          staircase=False,
          name="ExponentialDecay",
      )
      optimizer = tf.keras.optimizers.Adam(
          learning_rate=scheduler,
          beta_1=0.95,
          beta_2=0.99,
          epsilon=1e-6,
      )
      if args.amp:
          optimizer = mixed_precision.LossScaleOptimizer(optimizer)
  
      # LR Finder
      if args.find_lr:
          lrs, losses = record_lr(
              model,
              train_loader,
              batch_transforms,
              optimizer,
              amp=args.amp,
          )
          plot_recorder(lrs, losses)
          return
  
      # Tensorboard to monitor training
      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
      )
  
      config = {
          "learning_rate": args.lr,
          "epochs": args.epochs,
          "batch_size": args.batch_size,
          "architecture": args.arch,
          "input_size": args.input_size,
          "optimizer": optimizer.name,
          "framework": "tensorflow",
          "vocab": args.vocab,
          "scheduler": scheduler.name,
          "pretrained": args.pretrained,
      }
  
      # W&B
      if args.wb:
          import wandb
  
          run = wandb.init(
              name=exp_name,
              project="character-classification",
              config=config,
          )
      # ClearML
      if args.clearml:
          from clearml import Task
  
          task = Task.init(
              project_name="docTR/character-classification",
              task_name=exp_name,
              reuse_last_task_id=False,
          )
          task.upload_artifact("config", config)
  
      # 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, args.amp
          )
  
          # Validation loop at the end of each epoch
          val_loss, acc = evaluate(model, val_loader, batch_transforms)
          if val_loss < min_loss:
              print(
                  f"Validation loss decreased {min_loss:.6} --> {val_loss:.6}: saving state..."
              )
              model.save_weights(f"./{exp_name}/weights")
              min_loss = val_loss
          print(
              f"Epoch {epoch + 1}/{args.epochs} - Validation loss: {val_loss:.6} (Acc: {acc:.2%})"
          )
          # W&B
          if args.wb:
              wandb.log(
                  {
                      "val_loss": val_loss,
                      "acc": acc,
                  }
              )
  
          # ClearML
          if args.clearml:
              from clearml import Logger
  
              logger = Logger.current_logger()
              logger.report_scalar(
                  title="Validation Loss",
                  series="val_loss",
                  value=val_loss,
                  iteration=epoch,
              )
              logger.report_scalar(
                  title="Accuracy",
                  series="acc",
                  value=acc,
                  iteration=epoch,
              )
          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="classification", run_config=args
          )
  
      if args.export_onnx:
          print("Exporting model to ONNX...")
          if args.arch == "vit_b":
              # fixed batch size for vit
              dummy_input = [
                  tf.TensorSpec(
                      [1, args.input_size, args.input_size, 3],
                      tf.float32,
                      name="input",
                  )
              ]
          else:
              # dynamic batch size
              dummy_input = [
                  tf.TensorSpec(
                      [None, args.input_size, args.input_size, 3],
                      tf.float32,
                      name="input",
                  )
              ]
          model_path, _ = export_model_to_onnx(
              model, exp_name, dummy_input
          )
          print(f"Exported model saved in {model_path}")
  
file_path: doctr/references/classification/train_tensorflow.py
imports: import os
  import datetime
  import multiprocessing
  import time
  import numpy
  import tensorflow
  from tensorflow.keras import mixed_precision
  from tqdm.auto import tqdm
  from doctr.models import login_to_hub, push_to_hf_hub
  from doctr import transforms
  from doctr.datasets import VOCABS, CharacterGenerator, DataLoader
  from doctr.models import classification
  from doctr.models.utils import export_model_to_onnx
  from utils import EarlyStopper, plot_recorder, plot_samples
label_desc: Use image node_id 5 for calling a global function with example usage: main(args) without return types