code: def main(args):
      if args.backend.lower() != "pil" and not args.use_v2:
          # TODO: Support tensor backend in V1?
          raise ValueError(
              "Use --use-v2 if you want to use the tv_tensor or tensor backend."
          )
      if args.use_v2 and args.dataset != "coco":
          raise ValueError(
              "v2 is only support supported for coco dataset for now."
          )
  
      if args.output_dir:
          utils.mkdir(args.output_dir)
  
      utils.init_distributed_mode(args)
      print(args)
  
      device = torch.device(args.device)
  
      if args.use_deterministic_algorithms:
          torch.backends.cudnn.benchmark = False
          torch.use_deterministic_algorithms(True)
      else:
          torch.backends.cudnn.benchmark = True
  
      dataset, num_classes = get_dataset(args, is_train=True)
      dataset_test, _ = get_dataset(args, is_train=False)
  
      if args.distributed:
          train_sampler = (
              torch.utils.data.distributed.DistributedSampler(dataset)
          )
          test_sampler = (
              torch.utils.data.distributed.DistributedSampler(
                  dataset_test, shuffle=False
              )
          )
      else:
          train_sampler = torch.utils.data.RandomSampler(dataset)
          test_sampler = torch.utils.data.SequentialSampler(
              dataset_test
          )
  
      data_loader = torch.utils.data.DataLoader(
          dataset,
          batch_size=args.batch_size,
          sampler=train_sampler,
          num_workers=args.workers,
          collate_fn=utils.collate_fn,
          drop_last=True,
      )
  
      data_loader_test = torch.utils.data.DataLoader(
          dataset_test,
          batch_size=1,
          sampler=test_sampler,
          num_workers=args.workers,
          collate_fn=utils.collate_fn,
      )
  
      model = torchvision.models.get_model(
          args.model,
          weights=args.weights,
          weights_backbone=args.weights_backbone,
          num_classes=num_classes,
          aux_loss=args.aux_loss,
      )
      model.to(device)
      if args.distributed:
          model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
  
      model_without_ddp = model
      if args.distributed:
          model = torch.nn.parallel.DistributedDataParallel(
              model, device_ids=[args.gpu]
          )
          model_without_ddp = model.module
  
      params_to_optimize = [
          {
              "params": [
                  p
                  for p in model_without_ddp.backbone.parameters()
                  if p.requires_grad
              ]
          },
          {
              "params": [
                  p
                  for p in model_without_ddp.classifier.parameters()
                  if p.requires_grad
              ]
          },
      ]
      if args.aux_loss:
          params = [
              p
              for p in model_without_ddp.aux_classifier.parameters()
              if p.requires_grad
          ]
          params_to_optimize.append(
              {"params": params, "lr": args.lr * 10}
          )
      optimizer = torch.optim.SGD(
          params_to_optimize,
          lr=args.lr,
          momentum=args.momentum,
          weight_decay=args.weight_decay,
      )
  
      scaler = torch.cuda.amp.GradScaler() if args.amp else None
  
      iters_per_epoch = len(data_loader)
      main_lr_scheduler = PolynomialLR(
          optimizer,
          total_iters=iters_per_epoch
          * (args.epochs - args.lr_warmup_epochs),
          power=0.9,
      )
  
      if args.lr_warmup_epochs > 0:
          warmup_iters = iters_per_epoch * args.lr_warmup_epochs
          args.lr_warmup_method = args.lr_warmup_method.lower()
          if args.lr_warmup_method == "linear":
              warmup_lr_scheduler = torch.optim.lr_scheduler.LinearLR(
                  optimizer,
                  start_factor=args.lr_warmup_decay,
                  total_iters=warmup_iters,
              )
          elif args.lr_warmup_method == "constant":
              warmup_lr_scheduler = torch.optim.lr_scheduler.ConstantLR(
                  optimizer,
                  factor=args.lr_warmup_decay,
                  total_iters=warmup_iters,
              )
          else:
              raise RuntimeError(
                  f"Invalid warmup lr method '{args.lr_warmup_method}'. Only linear and constant are supported."
              )
          lr_scheduler = torch.optim.lr_scheduler.SequentialLR(
              optimizer,
              schedulers=[warmup_lr_scheduler, main_lr_scheduler],
              milestones=[warmup_iters],
          )
      else:
          lr_scheduler = main_lr_scheduler
  
      if args.resume:
          checkpoint = torch.load(
              args.resume, map_location="cpu", weights_only=True
          )
          model_without_ddp.load_state_dict(
              checkpoint["model"], strict=not args.test_only
          )
          if not args.test_only:
              optimizer.load_state_dict(checkpoint["optimizer"])
              lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
              args.start_epoch = checkpoint["epoch"] + 1
              if args.amp:
                  scaler.load_state_dict(checkpoint["scaler"])
  
      if args.test_only:
          # We disable the cudnn benchmarking because it can noticeably affect the accuracy
          torch.backends.cudnn.benchmark = False
          torch.backends.cudnn.deterministic = True
          confmat = evaluate(
              model,
              data_loader_test,
              device=device,
              num_classes=num_classes,
          )
          print(confmat)
          return
  
      start_time = time.time()
      for epoch in range(args.start_epoch, args.epochs):
          if args.distributed:
              train_sampler.set_epoch(epoch)
          train_one_epoch(
              model,
              criterion,
              optimizer,
              data_loader,
              lr_scheduler,
              device,
              epoch,
              args.print_freq,
              scaler,
          )
          confmat = evaluate(
              model,
              data_loader_test,
              device=device,
              num_classes=num_classes,
          )
          print(confmat)
          checkpoint = {
              "model": model_without_ddp.state_dict(),
              "optimizer": optimizer.state_dict(),
              "lr_scheduler": lr_scheduler.state_dict(),
              "epoch": epoch,
              "args": args,
          }
          if args.amp:
              checkpoint["scaler"] = scaler.state_dict()
          utils.save_on_master(
              checkpoint,
              os.path.join(args.output_dir, f"model_{epoch}.pth"),
          )
          utils.save_on_master(
              checkpoint,
              os.path.join(args.output_dir, "checkpoint.pth"),
          )
  
      total_time = time.time() - start_time
      total_time_str = str(datetime.timedelta(seconds=int(total_time)))
      print(f"Training time {total_time_str}")
  
file_path: vision/references/segmentation/train.py
imports: import datetime
  import os
  import time
  import warnings
  import presets
  import torch
  import torch.utils.data
  import torchvision
  import utils
  from coco_utils import get_coco
  from torch import nn
  from torch.optim.lr_scheduler import PolynomialLR
  from torchvision.transforms import functional, InterpolationMode
label_desc: Use image node_id 6 for calling a global function with example usage: main(args) without return types