code: def main(args):
      device = torch.device(
          args.device if torch.cuda.is_available() else "cpu"
      )
      batch_size = args.batch_size
      # segmentation nun_classes + background
      num_classes = args.num_classes + 1
  
      # 用来保存训练以及验证过程中信息
      results_file = "results{}.txt".format(
          datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
      )
  
      # VOCdevkit -> VOC2012 -> ImageSets -> Segmentation -> train.txt
      train_dataset = VOCSegmentation(
          args.data_path,
          year="2012",
          transforms=get_transform(train=True),
          txt_name="train.txt",
      )
  
      # VOCdevkit -> VOC2012 -> ImageSets -> Segmentation -> val.txt
      val_dataset = VOCSegmentation(
          args.data_path,
          year="2012",
          transforms=get_transform(train=False),
          txt_name="val.txt",
      )
  
      num_workers = min(
          [os.cpu_count(), batch_size if batch_size > 1 else 0, 8]
      )
      train_loader = torch.utils.data.DataLoader(
          train_dataset,
          batch_size=batch_size,
          num_workers=num_workers,
          shuffle=True,
          pin_memory=True,
          collate_fn=train_dataset.collate_fn,
      )
  
      val_loader = torch.utils.data.DataLoader(
          val_dataset,
          batch_size=1,
          num_workers=num_workers,
          pin_memory=True,
          collate_fn=val_dataset.collate_fn,
      )
  
      model = create_model(aux=args.aux, num_classes=num_classes)
      model.to(device)
  
      params_to_optimize = [
          {
              "params": [
                  p
                  for p in model.backbone.parameters()
                  if p.requires_grad
              ]
          },
          {
              "params": [
                  p
                  for p in model.classifier.parameters()
                  if p.requires_grad
              ]
          },
      ]
  
      if args.aux:
          params = [
              p
              for p in model.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_loader), args.epochs, warmup=True
      )
  
      if args.resume:
          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:
              scaler.load_state_dict(checkpoint["scaler"])
  
      start_time = time.time()
      for epoch in range(args.start_epoch, args.epochs):
          mean_loss, lr = train_one_epoch(
              model,
              optimizer,
              train_loader,
              device,
              epoch,
              lr_scheduler=lr_scheduler,
              print_freq=args.print_freq,
              scaler=scaler,
          )
  
          confmat = evaluate(
              model, val_loader, device=device, num_classes=num_classes
          )
          val_info = str(confmat)
          print(val_info)
          # 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")
  
          save_file = {
              "model": model.state_dict(),
              "optimizer": optimizer.state_dict(),
              "lr_scheduler": lr_scheduler.state_dict(),
              "epoch": epoch,
              "args": args,
          }
          if args.amp:
              save_file["scaler"] = scaler.state_dict()
          torch.save(
              save_file, "save_weights/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.py
imports: import os
  import time
  import datetime
  import torch
  from src import fcn_resnet50
  from train_utils import train_one_epoch, evaluate, create_lr_scheduler
  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