code: def main(args):
      device = torch.device(
          args.device if torch.cuda.is_available() else "cpu"
      )
      print("Using {} device training.".format(device.type))
  
      # 用来保存coco_info的文件
      now = datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
      det_results_file = f"det_results{now}.txt"
      seg_results_file = f"seg_results{now}.txt"
  
      data_transform = {
          "train": transforms.Compose(
              [
                  transforms.ToTensor(),
                  transforms.RandomHorizontalFlip(0.5),
              ]
          ),
          "val": transforms.Compose([transforms.ToTensor()]),
      }
  
      data_root = args.data_path
  
      # load train data set
      # coco2017 -> annotations -> instances_train2017.json
      train_dataset = CocoDetection(
          data_root, "train", data_transform["train"]
      )
      # VOCdevkit -> VOC2012 -> ImageSets -> Main -> train.txt
      # train_dataset = VOCInstances(data_root, year="2012", txt_name="train.txt", transforms=data_transform["train"])
      train_sampler = None
  
      # 是否按图片相似高宽比采样图片组成batch
      # 使用的话能够减小训练时所需GPU显存，默认使用
      if args.aspect_ratio_group_factor >= 0:
          train_sampler = torch.utils.data.RandomSampler(train_dataset)
          # 统计所有图像高宽比例在bins区间中的位置索引
          group_ids = create_aspect_ratio_groups(
              train_dataset, k=args.aspect_ratio_group_factor
          )
          # 每个batch图片从同一高宽比例区间中取
          train_batch_sampler = GroupedBatchSampler(
              train_sampler, group_ids, args.batch_size
          )
  
      # 注意这里的collate_fn是自定义的，因为读取的数据包括image和targets，不能直接使用默认的方法合成batch
      batch_size = args.batch_size
      nw = min(
          [os.cpu_count(), batch_size if batch_size > 1 else 0, 8]
      )  # number of workers
      print("Using %g dataloader workers" % nw)
  
      if train_sampler:
          # 如果按照图片高宽比采样图片，dataloader中需要使用batch_sampler
          train_data_loader = torch.utils.data.DataLoader(
              train_dataset,
              batch_sampler=train_batch_sampler,
              pin_memory=True,
              num_workers=nw,
              collate_fn=train_dataset.collate_fn,
          )
      else:
          train_data_loader = torch.utils.data.DataLoader(
              train_dataset,
              batch_size=batch_size,
              shuffle=True,
              pin_memory=True,
              num_workers=nw,
              collate_fn=train_dataset.collate_fn,
          )
  
      # load validation data set
      # coco2017 -> annotations -> instances_val2017.json
      val_dataset = CocoDetection(
          data_root, "val", data_transform["val"]
      )
      # VOCdevkit -> VOC2012 -> ImageSets -> Main -> val.txt
      # val_dataset = VOCInstances(data_root, year="2012", txt_name="val.txt", transforms=data_transform["val"])
      val_data_loader = torch.utils.data.DataLoader(
          val_dataset,
          batch_size=1,
          shuffle=False,
          pin_memory=True,
          num_workers=nw,
          collate_fn=train_dataset.collate_fn,
      )
  
      # create model num_classes equal background + classes
      model = create_model(
          num_classes=args.num_classes + 1,
          load_pretrain_weights=args.pretrain,
      )
      model.to(device)
  
      train_loss = []
      learning_rate = []
      val_map = []
  
      # define optimizer
      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
  
      # learning rate scheduler
      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.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"])
  
      for epoch in range(args.start_epoch, args.epochs):
          # train for one epoch, printing every 50 iterations
          mean_loss, lr = utils.train_one_epoch(
              model,
              optimizer,
              train_data_loader,
              device,
              epoch,
              print_freq=50,
              warmup=True,
              scaler=scaler,
          )
          train_loss.append(mean_loss.item())
          learning_rate.append(lr)
  
          # update the learning rate
          lr_scheduler.step()
  
          # evaluate on the test dataset
          det_info, seg_info = utils.evaluate(
              model, val_data_loader, device=device
          )
  
          # write detection into txt
          with open(det_results_file, "a") as f:
              # 写入的数据包括coco指标还有loss和learning rate
              result_info = [
                  f"{i:.4f}" for i in det_info + [mean_loss.item()]
              ] + [f"{lr:.6f}"]
              txt = "epoch:{} {}".format(epoch, "  ".join(result_info))
              f.write(txt + "\n")
  
          # write seg into txt
          with open(seg_results_file, "a") as f:
              # 写入的数据包括coco指标还有loss和learning rate
              result_info = [
                  f"{i:.4f}" for i in seg_info + [mean_loss.item()]
              ] + [f"{lr:.6f}"]
              txt = "epoch:{} {}".format(epoch, "  ".join(result_info))
              f.write(txt + "\n")
  
          val_map.append(det_info[1])  # pascal mAP
  
          # save weights
          save_files = {
              "model": model.state_dict(),
              "optimizer": optimizer.state_dict(),
              "lr_scheduler": lr_scheduler.state_dict(),
              "epoch": epoch,
          }
          if args.amp:
              save_files["scaler"] = scaler.state_dict()
          torch.save(
              save_files, "./save_weights/model_{}.pth".format(epoch)
          )
  
      # 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/mask_rcnn/train.py
imports: import os
  import datetime
  import torch
  from torchvision.ops.misc import FrozenBatchNorm2d
  import transforms
  from network_files import MaskRCNN
  from backbone import resnet50_fpn_backbone
  from my_dataset_coco import CocoDetection
  from my_dataset_voc import VOCInstances
  from train_utils import train_eval_utils
  from train_utils import GroupedBatchSampler, create_aspect_ratio_groups
label_desc: Use image node_id 2 for calling a global function with example usage: main(args) without return types