code: def main(parser_data):
      device = torch.device(
          parser_data.device if torch.cuda.is_available() else "cpu"
      )
      print("Using {} device training.".format(device.type))
  
      if not os.path.exists("save_weights"):
          os.mkdir("save_weights")
  
      results_file = "results{}.txt".format(
          datetime.datetime.now().strftime("%Y%m%d-%H%M%S")
      )
  
      data_transform = {
          "train": transforms.Compose(
              [
                  transforms.SSDCropping(),
                  transforms.Resize(),
                  transforms.ColorJitter(),
                  transforms.ToTensor(),
                  transforms.RandomHorizontalFlip(),
                  transforms.Normalization(),
                  transforms.AssignGTtoDefaultBox(),
              ]
          ),
          "val": transforms.Compose(
              [
                  transforms.Resize(),
                  transforms.ToTensor(),
                  transforms.Normalization(),
              ]
          ),
      }
  
      VOC_root = parser_data.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)
          )
  
      # VOCdevkit -> VOC2012 -> ImageSets -> Main -> train.txt
      train_dataset = VOCDataSet(
          VOC_root,
          "2012",
          data_transform["train"],
          train_set="train.txt",
      )
      # 注意训练时，batch_size必须大于1
      batch_size = parser_data.batch_size
      assert batch_size > 1, "batch size must be greater than 1"
      # 防止最后一个batch_size=1，如果最后一个batch_size=1就舍去
      drop_last = (
          True if len(train_dataset) % batch_size == 1 else False
      )
      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 = torch.utils.data.DataLoader(
          train_dataset,
          batch_size=batch_size,
          shuffle=True,
          num_workers=nw,
          collate_fn=train_dataset.collate_fn,
          drop_last=drop_last,
      )
  
      # VOCdevkit -> VOC2012 -> ImageSets -> Main -> val.txt
      val_dataset = VOCDataSet(
          VOC_root, "2012", data_transform["val"], train_set="val.txt"
      )
      val_data_loader = torch.utils.data.DataLoader(
          val_dataset,
          batch_size=batch_size,
          shuffle=False,
          num_workers=nw,
          collate_fn=train_dataset.collate_fn,
      )
  
      model = create_model(num_classes=args.num_classes + 1)
      model.to(device)
  
      # define optimizer
      params = [p for p in model.parameters() if p.requires_grad]
      optimizer = torch.optim.SGD(
          params, lr=0.0005, momentum=0.9, weight_decay=0.0005
      )
      # learning rate scheduler
      lr_scheduler = torch.optim.lr_scheduler.StepLR(
          optimizer, step_size=5, gamma=0.3
      )
  
      # 如果指定了上次训练保存的权重文件地址，则接着上次结果接着训练
      if parser_data.resume != "":
          checkpoint = torch.load(
              parser_data.resume, map_location="cpu"
          )
          model.load_state_dict(checkpoint["model"])
          optimizer.load_state_dict(checkpoint["optimizer"])
          lr_scheduler.load_state_dict(checkpoint["lr_scheduler"])
          parser_data.start_epoch = checkpoint["epoch"] + 1
          print(
              "the training process from epoch{}...".format(
                  parser_data.start_epoch
              )
          )
  
      train_loss = []
      learning_rate = []
      val_map = []
  
      # 提前加载验证集数据，以免每次验证时都要重新加载一次数据，节省时间
      val_data = get_coco_api_from_dataset(val_data_loader.dataset)
      for epoch in range(parser_data.start_epoch, parser_data.epochs):
          mean_loss, lr = utils.train_one_epoch(
              model=model,
              optimizer=optimizer,
              data_loader=train_data_loader,
              device=device,
              epoch=epoch,
              print_freq=50,
          )
          train_loss.append(mean_loss.item())
          learning_rate.append(lr)
  
          # update learning rate
          lr_scheduler.step()
  
          coco_info = utils.evaluate(
              model=model,
              data_loader=val_data_loader,
              device=device,
              data_set=val_data,
          )
  
          # write into txt
          with open(results_file, "a") as f:
              # 写入的数据包括coco指标还有loss和learning rate
              result_info = [
                  str(round(i, 4))
                  for i in coco_info + [mean_loss.item()]
              ] + [str(round(lr, 6))]
              txt = "epoch:{} {}".format(epoch, "  ".join(result_info))
              f.write(txt + "\n")
  
          val_map.append(coco_info[1])  # pascal mAP
  
          # save weights
          save_files = {
              "model": model.state_dict(),
              "optimizer": optimizer.state_dict(),
              "lr_scheduler": lr_scheduler.state_dict(),
              "epoch": epoch,
          }
          torch.save(
              save_files, "./save_weights/ssd300-{}.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/ssd/train_ssd300.py
imports: import os
  import datetime
  import torch
  import transforms
  from my_dataset import VOCDataSet
  from src import SSD300, Backbone
  import train_utils.train_eval_utils
  from train_utils import get_coco_api_from_dataset
label_desc: Use image node_id 2 for calling a global function with example usage: main(parser_data) without return types