【CenterFusion】训练执行过程CenterFusion/src/main.py

HIT_Vanni

已于 2024-01-19 17:39:57 修改

阅读量684

点赞数 11

分类专栏： CenterFusion 文章标签：深度学习人工智能自动驾驶机器学习

于 2024-01-19 17:39:44 首次发布

本文链接：https://blog.csdn.net/qq_34972053/article/details/135703690

版权

CenterFusion 专栏收录该内容

16 篇文章 3 订阅

订阅专栏

文章详细描述了CenterFusion项目的训练过程，涉及代码片段，如检查文件、设置随机数种子、选择优化级别（cudnn.benchmark）、数据加载、模型创建、优化器配置以及训练和验证步骤。着重讲解了cudnn.benchmark在性能优化中的作用和适用场景。

摘要由CSDN通过智能技术生成

文件作用：CenterFusion 项目训练的执行过程

首先查看文件的最末尾，最开始执行这段代码

if __name__ == '__main__':
  opt = opts().parse()
  '''
  调用 opts 类中的 parse() 函数
  这个函数在 CenterFusion/src/lib/opts.py 第 305 行
  '''
  
  main(opt)
  '''
  执行 main() 函数
  '''

然后执行 main() 函数中的代码

def main(opt):

  torch.manual_seed(opt.seed)
  '''
  这里引用了 torch 库内的函数 manual_seed() ：设置 CPU 生成随机数种子
  seed 在 opts.py 文件第 50 行定义，默认值为 317
  '''

  torch.backends.cudnn.benchmark = not opt.not_cuda_benchmark and not opt.eval
  '''
  参数 not_cuda_benchmark 在 opts.py 第 48 行，含义：不是cuda基准
    train.sh 脚本中没有添加 --not_cuda_benchmark 参数，所以 opt.not_cuda_benchmark 值为 False
  参数 eval 在 opts.py 第 22 行，含义：只评估测试集 mini_val 并退出
    train.sh 脚本中没有添加 --eval 参数，所以 opt.eval 值为 False
  综上得出，torch.backends.cudnn.benchmark = True
  PS:
  设置为 True 会让程序在开始时花费一点额外时间，使用 cudnn 为整个网络的每个卷积层搜索最适合它的卷积实现算法，进而实现网络的加速
  torch.backends.cudnn.benchmark是用于控制在使用CUDNN进行卷积操作时的优化级别的flag，该flag可以加速训练。如果该flag设置为True，则每次卷积操作都会针对输入大小进行优化，以便找到最快的卷积算法。但是，由于该操作会使用一些额外的内存，因此在使用较小的输入数据集时，并不总是最优的选择。
  因此，如果你在使用较大的数据集进行训练，可以使用以下代码来启用优化级别：
  torch.backends.cudnn.benchmark = True
  但是，如果你在使用较小的数据集，并且想减少内存占用，可以使用以下代码来禁用优化级别：
  torch.backends.cudnn.benchmark = False
  '''

  Dataset = get_dataset(opt.dataset)
  '''
  参数 dataset 在 opts.py 第 16 行，含义：设置默认数据集 nuscenes
  get_dataset() 函数在 CenterFusion/src/lib/dataset/dataset_factory.py 第 32 行
  返回一个 nuScenes 对象
  这个 nuScenes 对象定义在 CenterFusion/src/lib/dataset/datasets/nuscenes.py 第 24 行
  '''

  opt = opts().update_dataset_info_and_set_heads(opt, Dataset)
  '''
  更新一些配置信息  
  update_dataset_info_and_set_heads() 函数在 opts.py 第 458 行
  '''

  if not opt.not_set_cuda_env:
    os.environ['CUDA_VISIBLE_DEVICES'] = opt.gpus_str
  '''
  opt.not_set_cuda_env 在脚本中没有添加这个参数，所以值为 False，再用 not 取反，则为 True
  gpus_str 是 opt 中的一个 GPU 索引号字符串，如：'0,1'
  这里是为了给系统添加 cuda 索引号
  '''

  opt.device = torch.device('cuda' if opt.gpus[0] >= 0 else 'cpu')
  '''
  设置模型要分配的位置
  '''

  logger = Logger(opt)
  '''
  新建一个 Logger 对象，保存 opt 的配置信息
  Logger 对象在 CenterFusion/src/lib/logger.py 
  '''

  print('Creating model...')
  model = create_model(opt.arch, opt.heads, opt.head_conv, opt=opt)
  '''
  创建 DLA 模型（CenterNet 中的一种）
  参数：
      arch ：网络结构名称
      heads ：网络的头部
      head_conv ：头部输出的通道个数
  create_model 函数在 CenterFusion/src/lib/model/model.py 第 24 行
  '''
    optimizer = get_optimizer(opt, model)
  '''
  根据超参数和模型定义优化器
  '''

  start_epoch = 0
  '''
  设定初始轮次
  '''

  lr = opt.lr
  '''
  获取设定的学习率：0.00025
  '''

  if opt.load_model != '':
    model, optimizer, start_epoch = load_model(
      model, opt.load_model, opt, optimizer)
  '''
  如果加载的预训练模型不为空，则加载模型
  返回加载后的模型、优化器、初始轮次
  load_model() 函数在 CenterFusion/src/lib/model/model.py 第 31 行
  '''

  trainer = Trainer(opt, model, optimizer)
  '''
  根据参数、模型、优化器得到对应任务的训练器（其中包括损失统计、模型损失等）
  Trainer 类在 CenterFusion/src/lib/trainer.py 第 129 行
  '''

  trainer.set_device(opt.gpus, opt.chunk_sizes, opt.device)
  '''
  将训练器加载到 GPU 上
  路径 CenterFusion/src/lib/trainer.py 第 137 行
  '''

  if opt.val_intervals < opt.num_epochs or opt.eval:
    '''
    val_intervals=1 < num_epochs=60 执行该 if 语句
    '''

    print('Setting up validation data...')
    val_loader = torch.utils.data.DataLoader(
      Dataset(opt, opt.val_split), batch_size=1, shuffle=False, 
              num_workers=1, pin_memory=True)
    '''
    加载测试集数据
    torch.utils.data.DataLoader 是一个数据读取的一个接口，参数：
      dataset (Dataset)：加载数据的数据集
      batch_size (int, optional)：每个 batch 加载多少个样本(默认: 1)
      shuffle (bool, optional)：设置为 True 时会在每个 epoch 重新打乱数据(默认: False)
      num_workers (int, optional)：用多少个子进程加载数据。0 表示数据将在主进程中加载(默认: 0)
      pin_memory (bool, optional)：设置 pin_memory=True，则意味着生成的 Tensor 数据最开始是属于内存中的锁页内存，
                                 这样将内存的 Tensor 转义到 GPU 的显存就会更快一些
    '''

    if opt.eval:
      '''
      如果是测试，则执行，训练时，该 if 语句没有执行
      '''

      _, preds = trainer.val(0, val_loader)
      '''
      路径 CenterFusion/src/lib/trainer.py 第 402 行
      '''

      val_loader.dataset.run_eval(preds, opt.save_dir, n_plots=opt.eval_n_plots, 
                                  render_curves=opt.eval_render_curves)
      '''
      进行结果评判
      '''

      return

  print('Setting up train data...')

  train_loader = torch.utils.data.DataLoader(
      Dataset(opt, opt.train_split), batch_size=opt.batch_size, 
        shuffle=opt.shuffle_train, num_workers=opt.num_workers, 
        pin_memory=True, drop_last=True
  )
  '''
  加载训练集
  drop_last (bool, optional) – 如果数据集大小不能被 batch size 整除，则设置为 True 后可删除最后一个不完整的 batch。
  如果设为 False 并且数据集的大小不能被 batch size 整除，则最后一个 batch 将更小。(默认: False)
  '''
  print('Starting training...')
  for epoch in range(start_epoch + 1, opt.num_epochs + 1):
    '''
    循环 60 次，训练 60 轮
    '''

    mark = epoch if opt.save_all else 'last'
    '''
    是否每 5 个 epoch 保存模型到磁盘
    由于 train.sh 中没有添加 save_all 参数，所以为 False
    最后 mark = 'last'，意为最后再保存模型到磁盘中
    '''

    for param_group in optimizer.param_groups:
      lr = param_group['lr']
      logger.scalar_summary('LR', lr, epoch)
      break
    '''
    记录学习率
    scalar_summary() 函数在 CenterFusion/src/lib/logger.py 第 73 行
    '''
    
    log_dict_train, _ = trainer.train(epoch, train_loader)
    logger.write('epoch: {} |'.format(epoch))
    '''
    训练一轮模型，返回 ret 和 result
    这里 log_dict_train = ret
    train() 函数在 CenterFusion/src/lib/trainer.py 第 405 行
    '''

    for k, v in log_dict_train.items():
      logger.scalar_summary('train_{}'.format(k), v, epoch)
      logger.write('{} {:8f} | '.format(k, v))
    '''
    items() 函数：将一个字典以列表的形式返回，因为字典是无序的，所以返回的列表也是无序的
    记录训练一次返回的结果
    '''
    
    if opt.val_intervals > 0 and epoch % opt.val_intervals == 0:
      '''
      val_intervals 的值为 1
      epoch % opt.val_intervals 的值始终为 0
      所以每训练一个 epoch 都要使用测试集进行一次测试
      '''

      save_model(os.path.join(opt.save_dir, 'model_{}.pth'.format(mark)), 
                 epoch, model, optimizer)
      '''
      保存模型，下一轮训练好后又会覆盖上一轮的模型
      save_model() 函数在 CenterFusion/src/lib/model/model.py 第 117 行
      其中参数分别为保存路径、训练轮次、训练模型、优化器
      最后保存的模型放在 ~/CenterFusion/src/lib/../../exp/ddd/centerfusion 文件夹下，后缀名为 .pth
      '''

      with torch.no_grad():
        '''
        with 语句适用于对资源进行访问的场合，确保不管使用过程中是否发生异常都会执行必要的“清理”操作，释放资源
        比如文件使用后自动关闭／线程中锁的自动获取和释放等
        with torch.no_grad(): 强制后面不进行计算图（计算过程的构建，以便梯度反向传播等操作）的构建
        '''

        log_dict_val, preds = trainer.val(epoch, val_loader)
        '''
        对当前轮次训练后的模型进行测试
        '''
        
        if opt.run_dataset_eval:
          '''
          run_dataset_eval 在 train.sh 中添加了该参数，所以值为 True
          执行该 if 语句
          '''

          out_dir = val_loader.dataset.run_eval(preds, opt.save_dir, 
                                                n_plots=opt.eval_n_plots, 
                                                render_curves=opt.eval_render_curves)
          '''
          对测试结果评估
          '''
          
          with open('{}/metrics_summary.json'.format(out_dir), 'r') as f:
            metrics = json.load(f)
          logger.scalar_summary('AP/overall', metrics['mean_ap']*100.0, epoch)
          for k,v in metrics['mean_dist_aps'].items():
            logger.scalar_summary('AP/{}'.format(k), v*100.0, epoch)
          for k,v in metrics['tp_errors'].items():
            logger.scalar_summary('Scores/{}'.format(k), v, epoch)
          logger.scalar_summary('Scores/NDS', metrics['nd_score'], epoch)
          '''
          记录测试数据集的评估指标
          '''
      
      for k, v in log_dict_val.items():
        logger.scalar_summary('val_{}'.format(k), v, epoch)
        logger.write('{} {:8f} | '.format(k, v))
      '''
      记录测试结果
      '''
    
    #保存这个检查点
    else:
      save_model(os.path.join(opt.save_dir, 'model_last.pth'), 
                 epoch, model, optimizer)

    logger.write('\n')
    if epoch in opt.save_point:
      save_model(os.path.join(opt.save_dir, 'model_{}.pth'.format(epoch)), 
                 epoch, model, optimizer)
    '''
    save_point 在 train.sh 中的值为 20,40,50
    所以当 epoch 为 20,40,50 时保存模型
    '''
    
    if epoch in opt.lr_step:
      '''
      lr_step 在 train.sh 中的值为 50
      那么 epoch = 50 轮次时，更新学习率 lr
      '''
      lr = opt.lr * (0.1 ** (opt.lr_step.index(epoch) + 1))
      print('Drop LR to', lr)
      for param_group in optimizer.param_groups:
          param_group['lr'] = lr

  logger.close()

其中有get_optimizer()函数，具体如下：

def get_optimizer(opt, model):

  if opt.optim == 'adam':
    optimizer = torch.optim.Adam(model.parameters(), opt.lr)
    '''
    执行该 if 语句，其中 torch.optim.Adam 是实现 Adam 算法的函数
    '''

  elif opt.optim == 'sgd':
    print('Using SGD')
    optimizer = torch.optim.SGD(
      model.parameters(), opt.lr, momentum=0.9, weight_decay=0.0001)
  else:
    assert 0, opt.optim

  return optimizer
  '''
  返回优化器
  '''