code: def main(args):
      if torch.cuda.is_available() is False:
          raise EnvironmentError("not find GPU device for training.")
  
      # 初始化各进程环境
      init_distributed_mode(args=args)
  
      rank = args.rank
      device = torch.device(args.device)
      batch_size = args.batch_size
      num_classes = args.num_classes
      weights_path = args.weights
      args.lr *= args.world_size  # 学习率要根据并行GPU的数量进行倍增
  
      if rank == 0:  # 在第一个进程中打印信息，并实例化tensorboard
          print(args)
          print(
              'Start Tensorboard with "tensorboard --logdir=runs", view at http://localhost:6006/'
          )
          tb_writer = SummaryWriter()
          if os.path.exists("./weights") is False:
              os.makedirs("./weights")
  
      data_transform = {
          "train": transforms.Compose(
              [
                  transforms.RandomResizedCrop(224),
                  transforms.RandomHorizontalFlip(),
                  transforms.ToTensor(),
                  transforms.Normalize(
                      [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
                  ),
              ]
          ),
          "val": transforms.Compose(
              [
                  transforms.Resize(256),
                  transforms.CenterCrop(224),
                  transforms.ToTensor(),
                  transforms.Normalize(
                      [0.485, 0.456, 0.406], [0.229, 0.224, 0.225]
                  ),
              ]
          ),
      }
  
      data_root = args.data_path
      json_path = "./classes_name.json"
      # 实例化训练数据集
      train_dataset = MyDataSet(
          root_dir=data_root,
          csv_name="new_train.csv",
          json_path=json_path,
          transform=data_transform["train"],
      )
  
      # check num_classes
      if args.num_classes != len(train_dataset.labels):
          raise ValueError(
              "dataset have {} classes, but input {}".format(
                  len(train_dataset.labels), args.num_classes
              )
          )
  
      # 实例化验证数据集
      val_dataset = MyDataSet(
          root_dir=data_root,
          csv_name="new_val.csv",
          json_path=json_path,
          transform=data_transform["val"],
      )
  
      # 给每个rank对应的进程分配训练的样本索引
      train_sampler = torch.utils.data.distributed.DistributedSampler(
          train_dataset
      )
      val_sampler = torch.utils.data.distributed.DistributedSampler(
          val_dataset
      )
  
      # 将样本索引每batch_size个元素组成一个list
      train_batch_sampler = torch.utils.data.BatchSampler(
          train_sampler, batch_size, drop_last=True
      )
  
      nw = min(
          [os.cpu_count(), batch_size if batch_size > 1 else 0, 8]
      )  # number of workers
      if rank == 0:
          print("Using {} dataloader workers every process".format(nw))
  
      train_loader = torch.utils.data.DataLoader(
          train_dataset,
          batch_sampler=train_batch_sampler,
          pin_memory=True,
          num_workers=nw,
          collate_fn=train_dataset.collate_fn,
      )
  
      val_loader = torch.utils.data.DataLoader(
          val_dataset,
          batch_size=1,
          sampler=val_sampler,
          pin_memory=True,
          num_workers=nw,
          collate_fn=val_dataset.collate_fn,
      )
      # 实例化模型
      model = shufflenet_v2_x1_0(num_classes=num_classes).to(device)
  
      # 如果存在预训练权重则载入
      if os.path.exists(weights_path):
          weights_dict = torch.load(weights_path, map_location=device)
          load_weights_dict = {
              k: v
              for k, v in weights_dict.items()
              if model.state_dict()[k].numel() == v.numel()
          }
          model.load_state_dict(load_weights_dict, strict=False)
      else:
          checkpoint_path = os.path.join(
              tempfile.gettempdir(), "initial_weights.pt"
          )
          # 如果不存在预训练权重，需要将第一个进程中的权重保存，然后其他进程载入，保持初始化权重一致
          if rank == 0:
              torch.save(model.state_dict(), checkpoint_path)
  
          dist.barrier()
          # 这里注意，一定要指定map_location参数，否则会导致第一块GPU占用更多资源
          model.load_state_dict(
              torch.load(checkpoint_path, map_location=device)
          )
  
      # 是否冻结权重
      if args.freeze_layers:
          for name, para in model.named_parameters():
              # 除最后的全连接层外，其他权重全部冻结
              if "fc" not in name:
                  para.requires_grad_(False)
      else:
          # 只有训练带有BN结构的网络时使用SyncBatchNorm采用意义
          if args.syncBN:
              # 使用SyncBatchNorm后训练会更耗时
              model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(
                  model
              ).to(device)
  
      # 转为DDP模型
      model = torch.nn.parallel.DistributedDataParallel(
          model, device_ids=[args.gpu]
      )
  
      # optimizer
      pg = [p for p in model.parameters() if p.requires_grad]
      optimizer = optim.SGD(
          pg, lr=args.lr, momentum=0.9, weight_decay=4e-5
      )
      # Scheduler https://arxiv.org/pdf/1812.01187.pdf
      lf = (
          lambda x: ((1 + math.cos(x * math.pi / args.epochs)) / 2)
          * (1 - args.lrf)
          + args.lrf
      )  # cosine
      scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf)
  
      for epoch in range(args.epochs):
          train_sampler.set_epoch(epoch)
  
          mean_loss = train_one_epoch(
              model=model,
              optimizer=optimizer,
              data_loader=train_loader,
              device=device,
              epoch=epoch,
          )
  
          scheduler.step()
  
          acc = evaluate(
              model=model, data_loader=val_loader, device=device
          )
  
          if rank == 0:
              print(
                  "[epoch {}] accuracy: {}".format(epoch, round(acc, 3))
              )
              tags = ["loss", "accuracy", "learning_rate"]
              tb_writer.add_scalar(tags[0], mean_loss, epoch)
              tb_writer.add_scalar(tags[1], acc, epoch)
              tb_writer.add_scalar(
                  tags[2], optimizer.param_groups[0]["lr"], epoch
              )
  
              torch.save(
                  model.module.state_dict(),
                  "./weights/model-{}.pth".format(epoch),
              )
  
      # 删除临时缓存文件
      if rank == 0:
          if os.path.exists(checkpoint_path) is True:
              os.remove(checkpoint_path)
  
      cleanup()
  
file_path: deep-learning-for-image-processing/pytorch_classification/mini_imagenet/train_multi_gpu_using_launch.py
imports: import os
  import math
  import tempfile
  import argparse
  import torch
  import torch.optim
  import torch.optim.lr_scheduler
  from torch.utils.tensorboard import SummaryWriter
  from torchvision import transforms
  from model import shufflenet_v2_x1_0
  from my_dataset import MyDataSet
  from multi_train_utils import train_one_epoch, evaluate, init_distributed_mode, dist, cleanup
label_desc: Use image node_id 1 for calling a global function with example usage: main(args) without return types