Plantnet源码解析

from utils import set_seed, load_model, save, get_model, update_optimizer, get_data
from epoch import train_epoch, val_epoch, test_epoch
from cli import add_all_parsers


def train(args):
    set_seed(args, use_gpu=torch.cuda.is_available())
    train_loader, val_loader, test_loader, dataset_attributes = get_data(args.root, args.image_size, args.crop_size,
                                                                         args.batch_size, args.num_workers, args.pretrained)

    model = get_model(args, n_classes=dataset_attributes['n_classes'])
    criteria = CrossEntropyLoss()

    if args.use_gpu:
        print('USING GPU')
        torch.cuda.set_device(0)
        model.cuda()
        criteria.cuda()

    optimizer = SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=args.mu, nesterov=True)

    # Containers for storing metrics over epochs
    loss_train, acc_train, topk_acc_train = [], [], []
    loss_val, acc_val, topk_acc_val, avgk_acc_val, class_acc_val = [], [], [], [], []

    save_name = args.save_name_xp.strip()
    save_dir = os.path.join(os.getcwd(), 'results', save_name)
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)

    print('args.k : ', args.k)

    lmbda_best_acc = None
    best_val_acc = float('-inf')

GPT-4-Turbo

Poe

以下是对代码中每一行的逐行解释：

1. from utils import set_seed, load_model, save, get_model, update_optimizer, get_data

   • 从 utils 模块导入 set_seed, load_model, save, get_model, update_optimizer, get_data 几个函数。

2. from epoch import train_epoch, val_epoch, test_epoch

   • 从 epoch 模块导入 train_epoch, val_epoch, test_epoch 几个函数。

3. from cli import add_all_parsers

   • 从 cli 模块导入 add_all_parsers 函数。

4. def train(args):

   • 定义一个名为 train 的函数，它接受一个参数 args。

5. set_seed(args, use_gpu=torch.cuda.is_available())

   • 调用 set_seed 函数，设置随机数种子以确保结果可复制。此函数接受一个参数 use_gpu，其值由 torch.cuda.is_available() 决定，这个表达式检测是否有可用的 GPU。

6. train_loader, val_loader, test_loader, dataset_attributes = get_data(args.root, args.image_size, args.crop_size, args.batch_size, args.num_workers, args.pretrained)

   • 调用 get_data 函数获取数据加载器和数据集属性。此函数接受多个参数，包括数据根目录、图像大小、裁剪大小、批次大小、工作线程数量和是否使用预训练模型。

7. model = get_model(args, n_classes=dataset_attributes['n_classes'])

   • 调用 get_model 函数创建模型，模型类型由 args 决定，类别数 (n_classes) 从 dataset_attributes 字典中获取。

8. criteria = CrossEntropyLoss()

   • 创建一个 CrossEntropyLoss 对象，用于计算模型的损失函数。

9. if args.use_gpu:

   • 检查 args 中的 use_gpu 参数是否为真，判断是否使用 GPU。

10. print('USING GPU')

    • 如果使用 GPU，打印 "USING GPU"。

11. torch.cuda.set_device(0)

    • 设置当前 CUDA 设备为第一个 GPU（编号为0）。

12. model.cuda()

    • 将模型移至 GPU。

13. criteria.cuda()

    • 将损失函数对象移至 GPU。

14. optimizer = SGD(model.parameters(), lr=args.lr, momentum=0.9, weight_decay=args.mu, nesterov=True)

    • 创建 SGD 优化器，配置学习率、动量、权重衰减和 Nesterov 加速。

15. loss_train, acc_train, topk_acc_train = [], [], []

    • 初始化训练过程中损失、准确率和 top-k 准确率的列表。

16. loss_val, acc_val, topk_acc_val, avgk_acc_val, class_acc_val = [], [], [], [], []

    • 初始化验证过程中损失、准确率、top-k 准确率、平均 k 准确率和类别准确率的列表。

17. save_name = args.save_name_xp.strip()

    • 从 args 中获取实验保存名称，并去除可能的首尾空白字符。

18. save_dir = os.path.join(os.getcwd(), 'results', save_name)

    • 创建一个路径，该路径是当前工作目录、"results" 文件夹和实验名称的组合。

19. if not os.path.exists(save_dir):

    • 检查该路径是否存在。

20. os.makedirs(save_dir)

    • 如果路径不存在，则创建该目录。

21. print('args.k : ', args.k)

    • 打印 args 中的 k 值。

22. lmbda_best_acc = None

    • 初始化一个变量 lmbda_best_acc，用于可能的最佳准确率存储。

23. best_val_acc = float('-inf')

    • 初始化 best_val_acc 为负无穷大，用于比较和更新验证过程中的最佳准确率。

这段代码是一个训练模型的函数，涵盖了从数据加载、模型构建、设备设置到优化器配置等多个方面。

for epoch in tqdm(range(args.n_epochs), desc='epoch', position=0):
        t = time.time()
        optimizer = update_optimizer(optimizer, lr_schedule=args.epoch_decay, epoch=epoch)

        loss_epoch_train, acc_epoch_train, topk_acc_epoch_train = train_epoch(model, optimizer, train_loader,
                                                                              criteria, loss_train, acc_train,
                                                                              topk_acc_train, args.k,
                                                                              dataset_attributes['n_train'],
                                                                              args.use_gpu)

        loss_epoch_val, acc_epoch_val, topk_acc_epoch_val, \
        avgk_acc_epoch_val, lmbda_val = val_epoch(model, val_loader, criteria,
                                                  loss_val, acc_val, topk_acc_val, avgk_acc_val,
                                                  class_acc_val, args.k, dataset_attributes, args.use_gpu)

        # save model at every epoch
        save(model, optimizer, epoch, os.path.join(save_dir, save_name + '_weights.tar'))

        # save model with best val accuracy
        if acc_epoch_val > best_val_acc:
            best_val_acc = acc_epoch_val
            lmbda_best_acc = lmbda_val
            save(model, optimizer, epoch, os.path.join(save_dir, save_name + '_weights_best_acc.tar'))

        print()
        print(f'epoch {epoch} took {time.time()-t:.2f}')
        print(f'loss_train : {loss_epoch_train}')
        print(f'loss_val : {loss_epoch_val}')
        print(f'acc_train : {acc_epoch_train} / topk_acc_train : {topk_acc_epoch_train}')
        print(f'acc_val : {acc_epoch_val} / topk_acc_val : {topk_acc_epoch_val} / '
              f'avgk_acc_val : {avgk_acc_epoch_val}')

    # load weights corresponding to best val accuracy and evaluate on test
    load_model(model, os.path.join(save_dir, save_name + '_weights_best_acc.tar'), args.use_gpu)
    loss_test_ba, acc_test_ba, topk_acc_test_ba, \
    avgk_acc_test_ba, class_acc_test = test_epoch(model, test_loader, criteria, args.k,
                                                  lmbda_best_acc, args.use_gpu,
                                                  dataset_attributes)

    # Save the results as a dictionary and save it as a pickle file in desired location

这段代码为训练模型的主循环部分以及测试模型的代码块，下面是对每行代码的解释：

1. for epoch in tqdm(range(args.n_epochs), desc='epoch', position=0):

   • 使用 tqdm 库创建一个进度条，迭代从 0 到 args.n_epochs（总的训练周期数）。

2. t = time.time()

   • 记录当前时间，用于计算本轮训练的耗时。

3. optimizer = update_optimizer(optimizer, lr_schedule=args.epoch_decay, epoch=epoch)

   • 调用 update_optimizer 函数更新优化器的参数，可能会根据 args.epoch_decay 调整学习率。

4. loss_epoch_train, acc_epoch_train, topk_acc_epoch_train = train_epoch(model, optimizer, train_loader, criteria, loss_train, acc_train, topk_acc_train, args.k, dataset_attributes['n_train'], args.use_gpu)

   • 调用 train_epoch 函数进行一轮训练，返回训练的损失、准确率和 top-k 准确率。

5. loss_epoch_val, acc_epoch_val, topk_acc_epoch_val, \ avgk_acc_epoch_val, lmbda_val = val_epoch(model, val_loader, criteria, loss_val, acc_val, topk_acc_val, avgk_acc_val, class_acc_val, args.k, dataset_attributes, args.use_gpu)

   • 调用 val_epoch 函数进行一轮验证，返回验证的损失、准确率、top-k 准确率、平均 k 准确率和某个验证参数 lmbda_val。

6. save(model, optimizer, epoch, os.path.join(save_dir, save_name + '_weights.tar'))

   • 调用 save 函数将当前模型和优化器状态保存到文件。

7. if acc_epoch_val > best_val_acc:

   • 检查本轮验证准确率是否超过之前的最佳准确率。

8. best_val_acc = acc_epoch_val

   • 更新最佳验证准确率。

9. lmbda_best_acc = lmbda_val

   • 更新与最佳验证准确率相对应的 lmbda_val。

10. save(model, optimizer, epoch, os.path.join(save_dir, save_name + '_weights_best_acc.tar'))

    • 将具有最佳验证准确率的模型和优化器状态保存到文件。

11. print()

    • 打印一个空行。

12. print(f'epoch {epoch} took {time.time()-t:.2f}')

    • 打印本轮训练耗时。

13. print(f'loss_train : {loss_epoch_train}')

    • 打印本轮训练的损失。

14. print(f'loss_val : {loss_epoch_val}')

    • 打印本轮验证的损失。

15. print(f'acc_train : {acc_epoch_train} / topk_acc_train : {topk_acc_epoch_train}')

    打印本轮训练的准确率和 top-k 准确率。

16. print(f'acc_val : {acc_epoch_val} / topk_acc_val : {topk_acc_epoch_val} / ' f'avgk_acc_val : {avgk_acc_epoch_val}')

    打印本轮验证的准确率、top-k 准确率和平均 k 准确率。   
17.  load_model(model, os.path.join(save_dir, save_name + '_weights_best_acc.tar'), args.use_gpu)
18. 加载具有最佳验证准确率的模型权重
19. loss_test_ba, acc_test_ba, topk_acc_test_ba, \ avgk_acc_test_ba, class_acc_test = test_epoch(model, test_loader, criteria, args.k, lmbda_best_acc, args.use_gpu, dataset_attributes)
    • • 使用最佳模型权重进行测试，返回测试的损失、准确率、top-k 准确率、平均 k 准确率和类别准确率。

    • # Save the results as a dictionary and save it as a pickle file in desired location

      • 此行注释说明应该将测试结果作为字典保存，并将其作为 pickle 文件保存在指定位置，尽管实际的保存代码在这里没有显示。