Sfoglia il codice sorgente

modify fcos config

yjh0410 1 anno fa
parent
commit
5a39db43c2
1 ha cambiato i file con 5 aggiunte e 1 eliminazioni
  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))