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的文件
      results_file = "results{}.txt".format(
          datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
      )
  
      with open(args.keypoints_path, "r") as f:
          person_kps_info = json.load(f)
  
      fixed_size = args.fixed_size
      heatmap_hw = (args.fixed_size[0] // 4, args.fixed_size[1] // 4)
      kps_weights = np.array(
          person_kps_info["kps_weights"], dtype=np.float32
      ).reshape((args.num_joints,))
      data_transform = {
          "train": transforms.Compose(
              [
                  transforms.HalfBody(
                      0.3,
                      person_kps_info["upper_body_ids"],
                      person_kps_info["lower_body_ids"],
                  ),
                  transforms.AffineTransform(
                      scale=(0.65, 1.35),
                      rotation=(-45, 45),
                      fixed_size=fixed_size,
                  ),
                  transforms.RandomHorizontalFlip(
                      0.5, person_kps_info["flip_pairs"]
                  ),
                  transforms.KeypointToHeatMap(
                      heatmap_hw=heatmap_hw,
                      gaussian_sigma=2,
                      keypoints_weights=kps_weights,
                  ),
                  transforms.ToTensor(),
                  transforms.Normalize(
                      mean=[0.485, 0.456, 0.406],
                      std=[0.229, 0.224, 0.225],
                  ),
              ]
          ),
          "val": transforms.Compose(
              [
                  transforms.AffineTransform(
                      scale=(1.25, 1.25), fixed_size=fixed_size
                  ),
                  transforms.ToTensor(),
                  transforms.Normalize(
                      mean=[0.485, 0.456, 0.406],
                      std=[0.229, 0.224, 0.225],
                  ),
              ]
          ),
      }
  
      data_root = args.data_path
  
      # load train data set
      # coco2017 -> annotations -> person_keypoints_train2017.json
      train_dataset = CocoKeypoint(
          data_root,
          "train",
          transforms=data_transform["train"],
          fixed_size=args.fixed_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)
  
      train_data_loader = 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 -> person_keypoints_val2017.json
      val_dataset = CocoKeypoint(
          data_root,
          "val",
          transforms=data_transform["val"],
          fixed_size=args.fixed_size,
          det_json_path=args.person_det,
      )
      val_data_loader = data.DataLoader(
          val_dataset,
          batch_size=batch_size,
          shuffle=False,
          pin_memory=True,
          num_workers=nw,
          collate_fn=val_dataset.collate_fn,
      )
  
      # create model
      model = create_model(num_joints=args.num_joints)
      # print(model)
  
      model.to(device)
  
      # define optimizer
      params = [p for p in model.parameters() if p.requires_grad]
      optimizer = torch.optim.AdamW(
          params, lr=args.lr, 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
      )
  
      # 如果指定了上次训练保存的权重文件地址，则接着上次结果接着训练
      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 and "scaler" in checkpoint:
              scaler.load_state_dict(checkpoint["scaler"])
          print(
              "the training process from epoch{}...".format(
                  args.start_epoch
              )
          )
  
      train_loss = []
      learning_rate = []
      val_map = []
  
      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=device,
              epoch=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
          coco_info = utils.evaluate(
              model,
              val_data_loader,
              device=device,
              flip=True,
              flip_pairs=person_kps_info["flip_pairs"],
          )
  
          # 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")
  
          val_map.append(coco_info[1])  # @0.5 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_keypoint/HRNet/train.py
imports: import json
  import os
  import datetime
  import torch
  from torch.utils import data
  import numpy
  import transforms
  from model import HighResolutionNet
  from my_dataset_coco import CocoKeypoint
  from train_utils import train_eval_utils
label_desc: Use image node_id 2 for calling a global function with example usage: main(args) without return types