code: def main(args):
      init_distributed_mode(args)
      print(args)
  
      device = torch.device(args.device)
      # segmentation nun_classes + background
      num_classes = args.num_classes + 1
  
      # 用来保存coco_info的文件
      results_file = "results{}.txt".format(
          datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
      )
  
      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 -> Segmentation -> train.txt
      train_dataset = VOCSegmentation(
          args.data_path,
          year="2012",
          transforms=get_transform(train=True),
          txt_name="train.txt",
      )
      # load validation data set
      # VOCdevkit -> VOC2012 -> ImageSets -> Segmentation -> val.txt
      val_dataset = VOCSegmentation(
          args.data_path,
          year="2012",
          transforms=get_transform(train=False),
          txt_name="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)
  
      train_data_loader = torch.utils.data.DataLoader(
          train_dataset,
          batch_size=args.batch_size,
          sampler=train_sampler,
          num_workers=args.workers,
          collate_fn=train_dataset.collate_fn,
          drop_last=True,
      )
  
      val_data_loader = 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(aux=args.aux, num_classes=num_classes)
      model.to(device)
  
      if 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_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:
          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
  
      # 创建学习率更新策略，这里是每个step更新一次(不是每个epoch)
      lr_scheduler = create_lr_scheduler(
          optimizer, len(train_data_loader), args.epochs, warmup=True
      )
  
      # 如果传入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:
              scaler.load_state_dict(checkpoint["scaler"])
  
      if args.test_only:
          confmat = evaluate(
              model,
              val_data_loader,
              device=device,
              num_classes=num_classes,
          )
          val_info = str(confmat)
          print(val_info)
          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)
          mean_loss, lr = train_one_epoch(
              model,
              optimizer,
              train_data_loader,
              device,
              epoch,
              lr_scheduler=lr_scheduler,
              print_freq=args.print_freq,
              scaler=scaler,
          )
  
          confmat = evaluate(
              model,
              val_data_loader,
              device=device,
              num_classes=num_classes,
          )
          val_info = str(confmat)
          print(val_info)
  
          # 只在主进程上进行写操作
          if args.rank in [-1, 0]:
              # write into txt
              with open(results_file, "a") as f:
                  # 记录每个epoch对应的train_loss、lr以及验证集各指标
                  train_info = (
                      f"[epoch: {epoch}]\n"
                      f"train_loss: {mean_loss:.4f}\n"
                      f"lr: {lr:.6f}\n"
                  )
                  f.write(train_info + val_info + "\n\n")
  
          if args.output_dir:
              # 只在主节点上执行保存权重操作
              save_file = {
                  "model": model_without_ddp.state_dict(),
                  "optimizer": optimizer.state_dict(),
                  "lr_scheduler": lr_scheduler.state_dict(),
                  "args": args,
                  "epoch": epoch,
              }
              if args.amp:
                  save_file["scaler"] = scaler.state_dict()
              save_on_master(
                  save_file,
                  os.path.join(
                      args.output_dir, "model_{}.pth".format(epoch)
                  ),
              )
  
      total_time = time.time() - start_time
      total_time_str = str(datetime.timedelta(seconds=int(total_time)))
      print("Training time {}".format(total_time_str))
  
file_path: deep-learning-for-image-processing/pytorch_segmentation/fcn/train_multi_GPU.py
imports: import time
  import os
  import datetime
  import torch
  from src import fcn_resnet50
  from train_utils import train_one_epoch, evaluate, create_lr_scheduler, init_distributed_mode, save_on_master, mkdir
  from my_dataset import VOCSegmentation
  import transforms
label_desc: Use image node_id 3 for calling a global function with example usage: main(args) without return types