code: def main(args):
      init_distributed_mode(args)
      print(args)
  
      device = torch.device(args.device)
  
      # 用来保存coco_info的文件
      results_file = "results{}.txt".format(
          datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
      )
  
      # Data loading code
      print("Loading data")
  
      data_transform = {
          "train": transforms.Compose(
              [
                  transforms.ToTensor(),
                  transforms.RandomHorizontalFlip(0.5),
              ]
          ),
          "val": transforms.Compose([transforms.ToTensor()]),
      }
  
      VOC_root = args.data_path
      # check voc root
      if os.path.exists(os.path.join(VOC_root, "VOCdevkit")) is False:
          raise FileNotFoundError(
              "VOCdevkit dose not in path:'{}'.".format(VOC_root)
          )
  
      # load train data set
      # VOCdevkit -> VOC2012 -> ImageSets -> Main -> train.txt
      train_dataset = VOCDataSet(
          VOC_root, "2012", data_transform["train"], "train.txt"
      )
  
      # load validation data set
      # VOCdevkit -> VOC2012 -> ImageSets -> Main -> val.txt
      val_dataset = VOCDataSet(
          VOC_root, "2012", data_transform["val"], "val.txt"
      )
  
      print("Creating data loaders")
      if args.distributed:
          train_sampler = (
              torch.utils.data.distributed.DistributedSampler(
                  train_dataset
              )
          )
          test_sampler = (
              torch.utils.data.distributed.DistributedSampler(
                  val_dataset
              )
          )
      else:
          train_sampler = torch.utils.data.RandomSampler(train_dataset)
          test_sampler = torch.utils.data.SequentialSampler(val_dataset)
  
      if args.aspect_ratio_group_factor >= 0:
          # 统计所有图像比例在bins区间中的位置索引
          group_ids = create_aspect_ratio_groups(
              train_dataset, k=args.aspect_ratio_group_factor
          )
          train_batch_sampler = GroupedBatchSampler(
              train_sampler, group_ids, args.batch_size
          )
      else:
          train_batch_sampler = torch.utils.data.BatchSampler(
              train_sampler, args.batch_size, drop_last=True
          )
  
      data_loader = torch.utils.data.DataLoader(
          train_dataset,
          batch_sampler=train_batch_sampler,
          num_workers=args.workers,
          collate_fn=train_dataset.collate_fn,
      )
  
      data_loader_test = torch.utils.data.DataLoader(
          val_dataset,
          batch_size=1,
          sampler=test_sampler,
          num_workers=args.workers,
          collate_fn=train_dataset.collate_fn,
      )
  
      print("Creating model")
      # create model num_classes equal background + 20 classes
      model = create_model(num_classes=args.num_classes + 1)
      model.to(device)
  
      if args.distributed and args.sync_bn:
          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 = [p for p in model.parameters() if p.requires_grad]
      optimizer = torch.optim.SGD(
          params,
          lr=args.lr,
          momentum=args.momentum,
          weight_decay=args.weight_decay,
      )
  
      scaler = torch.cuda.amp.GradScaler() if args.amp else None
  
      # lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=args.lr_step_size, gamma=args.lr_gamma)
      lr_scheduler = torch.optim.lr_scheduler.MultiStepLR(
          optimizer, milestones=args.lr_steps, gamma=args.lr_gamma
      )
  
      # 如果传入resume参数，即上次训练的权重地址，则接着上次的参数训练
      if args.resume:
          # If map_location is missing, torch.load will first load the module to CPU
          # and then copy each parameter to where it was saved,
          # which would result in all processes on the same machine using the same set of devices.
          checkpoint = torch.load(
              args.resume, map_location="cpu"
          )  # 读取之前保存的权重文件(包括优化器以及学习率策略)
          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 and "scaler" in checkpoint:
              scaler.load_state_dict(checkpoint["scaler"])
  
      if args.test_only:
          utils.evaluate(model, data_loader_test, device=device)
          return
  
      train_loss = []
      learning_rate = []
      val_map = []
  
      print("Start training")
      start_time = time.time()
      for epoch in range(args.start_epoch, args.epochs):
          if args.distributed:
              train_sampler.set_epoch(epoch)
          mean_loss, lr = utils.train_one_epoch(
              model,
              optimizer,
              data_loader,
              device,
              epoch,
              args.print_freq,
              warmup=True,
              scaler=scaler,
          )
          train_loss.append(mean_loss.item())
          learning_rate.append(lr)
  
          # update learning rate
          lr_scheduler.step()
  
          # evaluate after every epoch
          coco_info = utils.evaluate(
              model, data_loader_test, device=device
          )
          val_map.append(coco_info[1])  # pascal mAP
  
          # 只在主进程上进行写操作
          if args.rank in [-1, 0]:
              # write into txt
              with open(results_file, "a") as f:
                  # 写入的数据包括coco指标还有loss和learning rate
                  result_info = [
                      f"{i:.4f}" for i in coco_info + [mean_loss.item()]
                  ] + [f"{lr:.6f}"]
                  txt = "epoch:{} {}".format(
                      epoch, "  ".join(result_info)
                  )
                  f.write(txt + "\n")
  
          if args.output_dir:
              # 只在主节点上执行保存权重操作
              save_files = {
                  "model": model_without_ddp.state_dict(),
                  "optimizer": optimizer.state_dict(),
                  "lr_scheduler": lr_scheduler.state_dict(),
                  "args": args,
                  "epoch": epoch,
              }
              if args.amp:
                  save_files["scaler"] = scaler.state_dict()
              save_on_master(
                  save_files,
                  os.path.join(args.output_dir, f"model_{epoch}.pth"),
              )
  
      total_time = time.time() - start_time
      total_time_str = str(datetime.timedelta(seconds=int(total_time)))
      print("Training time {}".format(total_time_str))
  
      if args.rank in [-1, 0]:
          # plot loss and lr curve
          if len(train_loss) != 0 and len(learning_rate) != 0:
              from plot_curve import plot_loss_and_lr
  
              plot_loss_and_lr(train_loss, learning_rate)
  
          # plot mAP curve
          if len(val_map) != 0:
              from plot_curve import plot_map
  
              plot_map(val_map)
  
file_path: deep-learning-for-image-processing/pytorch_object_detection/faster_rcnn/train_multi_GPU.py
imports: import time
  import os
  import datetime
  import torch
  import transforms
  from my_dataset import VOCDataSet
  from backbone import resnet50_fpn_backbone
  from network_files import FasterRCNN, FastRCNNPredictor
  import train_utils.train_eval_utils
  from train_utils import GroupedBatchSampler, create_aspect_ratio_groups, init_distributed_mode, save_on_master, mkdir
label_desc: Use image node_id 2 for calling a global function with example usage: main(args) without return types