update

odlab/main.py → odlab/train.py

@@ -137,7 +137,7 @@ def main():
 
     # ---------------------------- Build LR Scheduler ----------------------------
     cfg.warmup_iters = cfg.warmup_iters * cfg.grad_accumulate
-    wp_lr_scheduler = build_wp_lr_scheduler(cfg, cfg.base_lr)
+    wp_lr_scheduler = build_wp_lr_scheduler(cfg)
     lr_scheduler    = build_lr_scheduler(cfg, optimizer, args.resume)
 
     # ---------------------------- Build DDP model ----------------------------

odlab/train.sh

@@ -12,23 +12,11 @@ elif [[ $MODEL == *"fcos"* ]]; then
     # Epoch setting
     BATCH_SIZE=16
     EVAL_EPOCH=2
-elif [[ $MODEL == *"retinanet"* ]]; then
-    # Epoch setting
-    BATCH_SIZE=16
-    EVAL_EPOCH=2
-elif [[ $MODEL == *"plain_detr"* ]]; then
-    # Epoch setting
-    BATCH_SIZE=16
-    EVAL_EPOCH=2
-elif [[ $MODEL == *"rtdetr"* ]]; then
-    # Epoch setting
-    BATCH_SIZE=16
-    EVAL_EPOCH=1
 fi
 
 # -------------------------- Train Pipeline --------------------------
 if [ $WORLD_SIZE == 1 ]; then
-    python main.py \
+    python train.py \
         --cuda \
         --dataset ${DATASET}  \
         --root ${DATA_ROOT} \
@@ -37,7 +25,7 @@ if [ $WORLD_SIZE == 1 ]; then
         --eval_epoch ${EVAL_EPOCH}
 elif [[ $WORLD_SIZE -gt 1 && $WORLD_SIZE -le 8 ]]; then
     python -m torch.distributed.run --nproc_per_node=$WORLD_SIZE --master_port ${MASTER_PORT}  \
-        main.py \
+        train.py \
         --cuda \
         --distributed \
         --dataset ${DATASET}  \

odlab/utils/lr_scheduler.py

@@ -23,15 +23,15 @@ class LinearWarmUpScheduler(object):
         self.set_lr(optimizer, tmp_lr)
         
 ## Build WP LR Scheduler
-def build_wp_lr_scheduler(cfg, base_lr=0.01):
+def build_wp_lr_scheduler(cfg):
     print('==============================')
     print('WarmUpScheduler: {}'.format(cfg.warmup))
-    print('--base_lr: {}'.format(base_lr))
+    print('--base_lr: {}'.format(cfg.base_lr))
     print('--warmup_iters: {}'.format(cfg.warmup_iters))
     print('--warmup_factor: {}'.format(cfg.warmup_factor))
 
     if cfg.warmup == 'linear':
-        wp_lr_scheduler = LinearWarmUpScheduler(base_lr, cfg.warmup_iters, cfg.warmup_factor)
+        wp_lr_scheduler = LinearWarmUpScheduler(cfg.base_lr, cfg.warmup_iters, cfg.warmup_factor)
     
     return wp_lr_scheduler