code: def main(args):
      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
  
      # Data loading code
      print("Loading data")
      val_resize_size = tuple(args.val_resize_size)
      val_crop_size = tuple(args.val_crop_size)
      train_resize_size = tuple(args.train_resize_size)
      train_crop_size = tuple(args.train_crop_size)
  
      traindir = os.path.join(args.data_path, "train")
      valdir = os.path.join(args.data_path, "val")
  
      print("Loading training data")
      st = time.time()
      cache_path = _get_cache_path(traindir, args)
      transform_train = presets.VideoClassificationPresetTrain(
          crop_size=train_crop_size, resize_size=train_resize_size
      )
  
      if args.cache_dataset and os.path.exists(cache_path):
          print(f"Loading dataset_train from {cache_path}")
          dataset, _ = torch.load(cache_path, weights_only=True)
          dataset.transform = transform_train
      else:
          if args.distributed:
              print(
                  "It is recommended to pre-compute the dataset cache on a single-gpu first, as it will be faster"
              )
          dataset = datasets.KineticsWithVideoId(
              args.data_path,
              frames_per_clip=args.clip_len,
              num_classes=args.kinetics_version,
              split="train",
              step_between_clips=1,
              transform=transform_train,
              frame_rate=args.frame_rate,
              extensions=(
                  "avi",
                  "mp4",
              ),
              output_format="TCHW",
          )
          if args.cache_dataset:
              print(f"Saving dataset_train to {cache_path}")
              utils.mkdir(os.path.dirname(cache_path))
              utils.save_on_master((dataset, traindir), cache_path)
  
      print("Took", time.time() - st)
  
      print("Loading validation data")
      cache_path = _get_cache_path(valdir, args)
  
      if args.weights and args.test_only:
          weights = torchvision.models.get_weight(args.weights)
          transform_test = weights.transforms()
      else:
          transform_test = presets.VideoClassificationPresetEval(
              crop_size=val_crop_size, resize_size=val_resize_size
          )
  
      if args.cache_dataset and os.path.exists(cache_path):
          print(f"Loading dataset_test from {cache_path}")
          dataset_test, _ = torch.load(cache_path, weights_only=True)
          dataset_test.transform = transform_test
      else:
          if args.distributed:
              print(
                  "It is recommended to pre-compute the dataset cache on a single-gpu first, as it will be faster"
              )
          dataset_test = datasets.KineticsWithVideoId(
              args.data_path,
              frames_per_clip=args.clip_len,
              num_classes=args.kinetics_version,
              split="val",
              step_between_clips=1,
              transform=transform_test,
              frame_rate=args.frame_rate,
              extensions=(
                  "avi",
                  "mp4",
              ),
              output_format="TCHW",
          )
          if args.cache_dataset:
              print(f"Saving dataset_test to {cache_path}")
              utils.mkdir(os.path.dirname(cache_path))
              utils.save_on_master((dataset_test, valdir), cache_path)
  
      print("Creating data loaders")
      train_sampler = RandomClipSampler(
          dataset.video_clips, args.clips_per_video
      )
      test_sampler = UniformClipSampler(
          dataset_test.video_clips, args.clips_per_video
      )
      if args.distributed:
          train_sampler = DistributedSampler(train_sampler)
          test_sampler = DistributedSampler(test_sampler, shuffle=False)
  
      data_loader = torch.utils.data.DataLoader(
          dataset,
          batch_size=args.batch_size,
          sampler=train_sampler,
          num_workers=args.workers,
          pin_memory=True,
          collate_fn=collate_fn,
      )
  
      data_loader_test = torch.utils.data.DataLoader(
          dataset_test,
          batch_size=args.batch_size,
          sampler=test_sampler,
          num_workers=args.workers,
          pin_memory=True,
          collate_fn=collate_fn,
      )
  
      print("Creating model")
      model = torchvision.models.get_model(
          args.model, weights=args.weights
      )
      model.to(device)
      if args.distributed and args.sync_bn:
          model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)
  
      criterion = nn.CrossEntropyLoss()
  
      optimizer = torch.optim.SGD(
          model.parameters(),
          lr=args.lr,
          momentum=args.momentum,
          weight_decay=args.weight_decay,
      )
      scaler = torch.cuda.amp.GradScaler() if args.amp else None
  
      # convert scheduler to be per iteration, not per epoch, for warmup that lasts
      # between different epochs
      iters_per_epoch = len(data_loader)
      lr_milestones = [
          iters_per_epoch * (m - args.lr_warmup_epochs)
          for m in args.lr_milestones
      ]
      main_lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
          optimizer, milestones=lr_milestones, gamma=args.lr_gamma
      )
  
      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
  
      model_without_ddp = model
      if args.distributed:
          model = torch.nn.parallel.DistributedDataParallel(
              model, device_ids=[args.gpu]
          )
          model_without_ddp = model.module
  
      if args.resume:
          checkpoint = torch.load(
              args.resume, map_location="cpu", weights_only=True
          )
          model_without_ddp.load_state_dict(checkpoint["model"])
          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
          evaluate(model, criterion, data_loader_test, device=device)
          return
  
      print("Start training")
      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,
              lr_scheduler,
              data_loader,
              device,
              epoch,
              args.print_freq,
              scaler,
          )
          evaluate(model, criterion, data_loader_test, device=device)
          if args.output_dir:
              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/video_classification/train.py
imports: import datetime
  import os
  import time
  import warnings
  import datasets
  import presets
  import torch
  import torch.utils.data
  import torchvision
  import torchvision.datasets.video_utils
  import utils
  from torch import nn
  from torch.utils.data.dataloader import default_collate
  from torchvision.datasets.samplers import DistributedSampler, RandomClipSampler, UniformClipSampler
label_desc: Use image node_id 5 for calling a global function with example usage: main(args) without return types