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
  
      st = time.time()
      val_set = DocArtefacts(
          train=False,
          download=True,
          img_transforms=T.Resize((args.input_size, args.input_size)),
      )
      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)"
      )
  
      # Load doctr model
      model = obj_detection.__dict__[args.arch](
          pretrained=args.pretrained, num_classes=5
      )
  
      # 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, target device set to CPU."
          )
      if torch.cuda.is_available():
          torch.cuda.set_device(args.device)
          model = model.cuda()
  
      # Metrics
      metric = DetectionMetric(iou_thresh=0.5)
  
      if args.test_only:
          print("Running evaluation")
          recall, precision, mean_iou = evaluate(
              model, val_loader, metric, amp=args.amp
          )
          print(
              f"Recall: {recall:.2%} | Precision: {precision:.2%} |IoU: {mean_iou:.2%}"
          )
          return
  
      st = time.time()
      # Load train data generators
      train_set = DocArtefacts(
          train=True,
          download=True,
          img_transforms=Compose(
              [
                  T.Resize((args.input_size, args.input_size)),
                  T.RandomApply(T.GaussianNoise(0.0, 0.25), p=0.5),
                  ColorJitter(
                      brightness=0.3,
                      contrast=0.3,
                      saturation=0.3,
                      hue=0.02,
                  ),
                  T.RandomApply(
                      GaussianBlur(kernel_size=(3, 3), sigma=(0.1, 3)),
                      0.3,
                  ),
              ]
          ),
          sample_transforms=T.RandomHorizontalFlip(p=0.5),
      )
      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:
          images, targets = next(iter(train_loader))
          targets = convert_to_abs_coords(targets, images.shape)
          plot_samples(images, targets, train_set.CLASSES)
          return
  
      # Backbone freezing
      if args.freeze_backbone:
          for p in model.backbone.parameters():
              p.reguires_grad_(False)
  
      # Optimizer
      optimizer = optim.SGD(
          [p for p in model.parameters() if p.requires_grad],
          lr=args.lr,
          weight_decay=args.weight_decay,
      )
      # LR Finder
      if args.find_lr:
          lrs, losses = record_lr(
              model, train_loader, optimizer, amp=args.amp
          )
          plot_recorder(lrs, losses)
          return
      # Scheduler
      scheduler = StepLR(optimizer, step_size=8, gamma=0.7)
  
      # 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="object-detection",
              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": "sgd",
                  "framework": "pytorch",
                  "scheduler": "step",
                  "pretrained": args.pretrained,
                  "amp": args.amp,
              },
          )
  
      max_score = 0.0
      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, 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
          )
  
          if f1_score > max_score:
              print(
                  f"Validation metric increased {max_score:.6} --> {f1_score:.6}: saving state..."
              )
              torch.save(model.state_dict(), f"./{exp_name}.pt")
              max_score = f1_score
          log_msg = f"Epoch {epoch + 1}/{args.epochs} - "
          if any(val is None for val in (recall, precision, mean_iou)):
              log_msg += "Undefined metric value, caused by empty GTs or predictions"
          else:
              log_msg += f"Recall: {recall:.2%} | Precision: {precision:.2%} | Mean IoU: {mean_iou:.2%}"
          print(log_msg)
          # W&B
          if args.wb:
              wandb.log(
                  {
                      "recall": recall,
                      "precision": precision,
                      "mean_iou": mean_iou,
                  }
              )
          if args.early_stop and early_stopper.early_stop(f1_score):
              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="obj_detection", run_config=args
          )
  
imports: import os
  import datetime
  import logging
  import multiprocessing
  import time
  import numpy
  import torch
  import torch.optim
  import wandb
  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
  from doctr import transforms
  from doctr.datasets import DocArtefacts
  from doctr.models import login_to_hub, obj_detection, push_to_hf_hub
  from doctr.utils import DetectionMetric
  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