Procházet zdrojové kódy

modify fcos config

yjh0410 před 1 rokem
rodič
revize
5a39db43c2
1 změnil soubory, kde provedl 5 přidání a 1 odebrání
  1. 5 1
      yolo/utils/solver/lr_scheduler.py

+ 5 - 1
yolo/utils/solver/lr_scheduler.py

@@ -33,7 +33,11 @@ def build_lr_scheduler(cfg, optimizer, resume=None):
         lr_scheduler = MultiStepLR(optimizer, milestones=lr_step, gamma=0.1)
 
     elif cfg.lr_scheduler == "cosine":
-        lr_scheduler = CosineAnnealingLR(optimizer, T_max=cfg.max_epoch - cfg.warmup_epoch - 1, eta_min=cfg.min_lr)
+        if hasattr(cfg, "warmup_epoch"):
+            total_epochs = cfg.max_epoch - cfg.warmup_epoch - 1
+        else:
+            total_epochs = cfg.max_epoch - 1
+        lr_scheduler = CosineAnnealingLR(optimizer, T_max=total_epochs, eta_min=cfg.min_lr)
     
     else:
         raise NotImplementedError("Unknown lr scheduler: {}".format(cfg.lr_scheduler))