code: def main(args):
      device = torch.device(
          args.device if torch.cuda.is_available() else "cpu"
      )
  
      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")
  
      (
          train_images_path,
          train_images_label,
          val_images_path,
          val_images_label,
      ) = read_split_data(args.data_path)
  
      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]
                  ),
              ]
          ),
      }
  
      # 实例化训练数据集
      train_dataset = MyDataSet(
          images_path=train_images_path,
          images_class=train_images_label,
          transform=data_transform["train"],
      )
  
      # 实例化验证数据集
      val_dataset = MyDataSet(
          images_path=val_images_path,
          images_class=val_images_label,
          transform=data_transform["val"],
      )
  
      batch_size = args.batch_size
      nw = min(
          [os.cpu_count(), batch_size if batch_size > 1 else 0, 8]
      )  # number of workers
      print("Using {} dataloader workers every process".format(nw))
      train_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,
      )
  
      val_loader = torch.utils.data.DataLoader(
          val_dataset,
          batch_size=batch_size,
          shuffle=False,
          pin_memory=True,
          num_workers=nw,
          collate_fn=val_dataset.collate_fn,
      )
  
      # 如果存在预训练权重则载入
      model = create_regnet(
          model_name=args.model_name, num_classes=args.num_classes
      ).to(device)
      # print(model)
  
      if args.weights != "":
          if os.path.exists(args.weights):
              weights_dict = torch.load(
                  args.weights, map_location=device
              )
              load_weights_dict = {
                  k: v
                  for k, v in weights_dict.items()
                  if model.state_dict()[k].numel() == v.numel()
              }
              print(
                  model.load_state_dict(load_weights_dict, strict=False)
              )
          else:
              raise FileNotFoundError(
                  "not found weights file: {}".format(args.weights)
              )
  
      # 是否冻结权重
      if args.freeze_layers:
          for name, para in model.named_parameters():
              # 除最后的全连接层外，其他权重全部冻结
              if "head" not in name:
                  para.requires_grad_(False)
              else:
                  print("train {}".format(name))
  
      pg = [p for p in model.parameters() if p.requires_grad]
      optimizer = optim.SGD(
          pg, lr=args.lr, momentum=0.9, weight_decay=5e-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
          mean_loss = train_one_epoch(
              model=model,
              optimizer=optimizer,
              data_loader=train_loader,
              device=device,
              epoch=epoch,
          )
  
          scheduler.step()
  
          # validate
          acc = evaluate(
              model=model, data_loader=val_loader, device=device
          )
  
          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.state_dict(), "./weights/model-{}.pth".format(epoch)
          )
  
file_path: deep-learning-for-image-processing/pytorch_classification/Test10_regnet/train.py
imports: import os
  import math
  import argparse
  import torch
  import torch.optim
  from torch.utils.tensorboard import SummaryWriter
  from torchvision import transforms
  import torch.optim.lr_scheduler
  from model import create_regnet
  from my_dataset import MyDataSet
  from utils import read_split_data, train_one_epoch, evaluate
label_desc: Use image node_id 1 for calling a global function with example usage: main(args) without return types