CenterFusion/src/main.py 項目訓練執行檔案詳解

目錄

• 一、train.sh 腳本
• 二、main.py 檔案
• 三、參考資料

檔案位置：

CenterFusion/experiments/train.sh

和

CenterFusion/src/main.py

檔案作用：CenterFusion 項目訓練的執行過程

注意：本文中的代碼都是 CenterFusion 原始代碼，一些參數沒有修改

一、train.sh 腳本

• 在 README.md 中訓練模型的指令是：

  bash experiments/train.sh

• 首先執行的就是 train.sh 腳本
• 在腳本中

  --參數 值

  表示可選參數

cd src
python main.py \
    ddd \
    '''
    可以了解為工作空間名，也就是檔案夾名
    '''
    --exp_id centerfusion \
    '''
    項目名稱
    '''
    --shuffle_train \
    '''
    将訓練模型的資料集進行打亂的操作
    '''
    --train_split train \
    '''
    訓練集
    '''
    --val_split mini_val \
    '''
    測試集
    '''
    --val_intervals 1 \
    '''
    運作測試集的 epoch 數為 1
    比如 1000 個樣本用來訓練一個神經網絡，訓練就算完成 1 epoch（期）
    '''
    --run_dataset_eval \
    '''
    在 eval 中使用資料集特定的計算函數
    '''
    --nuscenes_att \
    --velocity \
    --batch_size 32 \
    '''
    一次訓練抓取 32 個資料樣本
    '''
    --lr 2.5e-4 \
    '''
    學習率：0.00025
    '''
    --num_epochs 60 \
    '''
    訓練 60 輪
    '''
    --lr_step 50 \
    '''
    學習率步長
    '''
    --save_point 20,40,50 \
    '''
    模型儲存時間點
    '''
    --gpus 0,1 \
    '''
    這裡使用了兩塊 GPU ：0 号 GPU 和 1 号 GPU
    '''
    --not_rand_crop \
    '''
    不使用來自 CenterNet 的随機裁剪資料增強
    '''
    --flip 0.5 \
    '''
    使用翻轉資料增強的機率 50%
    '''
    --shift 0.1 \
    '''
    當不使用随機裁剪時，10% 機率使用移位增強
    '''
    --pointcloud \
    '''
    雷達點雲
    '''
    --radar_sweeps 3 \
    '''
    點雲圖中雷達掃瞄 3 次
    '''
    --pc_z_offset 0.0 \
    '''
    向 z 軸方面提高雷達點
    '''
    --pillar_dims 1.0,0.2,0.2 \
    '''
    雷達柱尺寸(h、w、l)
    '''
    --max_pc_dist 60.0 \
    '''
    移除最大點雲距離 60 以外的點
    '''
    --load_model ../models/centernet_baseline_e170.pth \
    '''
    導入的模型
    '''
cd ..
           

二、main.py 檔案

• 在 train.sh 腳本中執行了 src 下的 main.py 檔案，并傳遞了一些參數

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import _init_paths
import os

import torch
import torch.utils.data
from opts import opts
from model.model import create_model, load_model, save_model
from model.data_parallel import DataParallel
from logger import Logger
from dataset.dataset_factory import get_dataset
from trainer import Trainer
from test import prefetch_test
import json

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
  '''
  傳回優化器
  '''
           

• 首先檢視檔案的最末尾，最開始執行這段代碼

if __name__ == '__main__':
  opt = opts().parse()
  '''
  調用 opts 類中的 parse() 函數
  這個函數在 CenterFusion/src/lib/opts.py 第 305 行
  '''
  
  main(opt)
  '''
  執行 main() 函數
  '''
           

• 注意：這裡我将 opts.py 檔案中一些必要代碼進行注釋，可以參考部落格：CenterFusion/src/lib/opts.py 檔案代碼詳解
• 然後執行 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
  設定為 True 會讓程式在開始時花費一點額外時間，使用 cudnn 為整個網絡的每個卷積層搜尋最适合它的卷積實作算法，進而實作網絡的加速
  '''

  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() 函數在 opys.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 行
  '''
           

• 注意：這裡我将 model.py 檔案中一些必要代碼進行注釋，可以參考部落格：CenterFusion/src/lib/model/model.py 檔案代碼詳解

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 行
    '''
           

• 訓練過程請參考：CenterFusion/src/lib/trainer.py 訓練一個 epoch 過程

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()
           

三、參考資料

[1] pytorch中tensorboard的add_scalar與add_image