From dbe5de00dbed8d656c1683c7a7e4a188f54f134c Mon Sep 17 00:00:00 2001
From: CoinCheung <867153576@qq.com>
Date: Wed, 5 Jun 2019 21:20:56 +0800
Subject: [PATCH] cosine

---
 maskrcnn_benchmark/config/defaults.py     |  2 ++
 maskrcnn_benchmark/solver/build.py        | 28 ++++++++++-----
 maskrcnn_benchmark/solver/lr_scheduler.py | 44 +++++++++++++++++++++++
 3 files changed, 65 insertions(+), 9 deletions(-)

diff --git a/maskrcnn_benchmark/config/defaults.py b/maskrcnn_benchmark/config/defaults.py
index 65fbdaddd..81eea0e41 100644
--- a/maskrcnn_benchmark/config/defaults.py
+++ b/maskrcnn_benchmark/config/defaults.py
@@ -399,8 +399,10 @@
 _C.SOLVER.WEIGHT_DECAY = 0.0005
 _C.SOLVER.WEIGHT_DECAY_BIAS = 0
 
+_C.SOLVER.SCHEDULER = "multi_step" # or "cosine"
 _C.SOLVER.GAMMA = 0.1
 _C.SOLVER.STEPS = (30000,)
+_C.SOLVER.ETA_MIN = 0
 
 _C.SOLVER.WARMUP_FACTOR = 1.0 / 3
 _C.SOLVER.WARMUP_ITERS = 500
diff --git a/maskrcnn_benchmark/solver/build.py b/maskrcnn_benchmark/solver/build.py
index 865a4ec8d..2d4128a75 100644
--- a/maskrcnn_benchmark/solver/build.py
+++ b/maskrcnn_benchmark/solver/build.py
@@ -1,7 +1,7 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
 import torch
 
-from .lr_scheduler import WarmupMultiStepLR
+from .lr_scheduler import WarmupMultiStepLR, WarmupCosineAnnealingLR
 
 
 def make_optimizer(cfg, model):
@@ -21,11 +21,21 @@ def make_optimizer(cfg, model):
 
 
 def make_lr_scheduler(cfg, optimizer):
-    return WarmupMultiStepLR(
-        optimizer,
-        cfg.SOLVER.STEPS,
-        cfg.SOLVER.GAMMA,
-        warmup_factor=cfg.SOLVER.WARMUP_FACTOR,
-        warmup_iters=cfg.SOLVER.WARMUP_ITERS,
-        warmup_method=cfg.SOLVER.WARMUP_METHOD,
-    )
+    if cfg.SOLVER.SCHEDULER == "multi_step":
+        return WarmupMultiStepLR(
+            optimizer,
+            cfg.SOLVER.STEPS,
+            cfg.SOLVER.GAMMA,
+            warmup_factor=cfg.SOLVER.WARMUP_FACTOR,
+            warmup_iters=cfg.SOLVER.WARMUP_ITERS,
+            warmup_method=cfg.SOLVER.WARMUP_METHOD,
+        )
+    elif cfg.SOLVER.SCHEDULER == "cosine":
+        return WarmupCosineAnnealingLR(
+            optimizer,
+            cfg.SOLVER.MAX_ITER,
+            cfg.SOLVER.ETA_MIN,
+            warmup_factor=cfg.SOLVER.WARMUP_FACTOR,
+            warmup_iters=cfg.SOLVER.WARMUP_ITERS,
+            warmup_method=cfg.SOLVER.WARMUP_METHOD,
+        )
diff --git a/maskrcnn_benchmark/solver/lr_scheduler.py b/maskrcnn_benchmark/solver/lr_scheduler.py
index d7d45b6c6..c98ad5e91 100644
--- a/maskrcnn_benchmark/solver/lr_scheduler.py
+++ b/maskrcnn_benchmark/solver/lr_scheduler.py
@@ -1,4 +1,5 @@
 # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+import math
 from bisect import bisect_right
 
 import torch
@@ -50,3 +51,46 @@ def get_lr(self):
             * self.gamma ** bisect_right(self.milestones, self.last_epoch)
             for base_lr in self.base_lrs
         ]
+
+
+class WarmupCosineAnnealingLR(torch.optim.lr_scheduler._LRScheduler):
+    def __init__(
+        self,
+        optimizer,
+        max_iters,
+        eta_min=0,
+        warmup_factor=1.0 / 3,
+        warmup_iters=500,
+        warmup_method="linear",
+        last_epoch=-1,
+    ):
+        if warmup_method not in ("constant", "linear"):
+            raise ValueError(
+                "Only 'constant' or 'linear' warmup_method accepted"
+                "got {}".format(warmup_method)
+            )
+        self.max_iters = max_iters
+        self.eta_min = eta_min
+        self.warmup_factor = warmup_factor
+        self.warmup_iters = warmup_iters
+        self.warmup_method = warmup_method
+        super(WarmupCosineAnnealingLR, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        warmup_factor = 1
+        if self.last_epoch < self.warmup_iters:
+            if self.warmup_method == "constant":
+                warmup_factor = self.warmup_factor
+            elif self.warmup_method == "linear":
+                alpha = float(self.last_epoch) / self.warmup_iters
+                warmup_factor = self.warmup_factor * (1 - alpha) + alpha
+            lr_group = [base_lr * warmup_factor for base_lr in self.base_lrs]
+        else:
+            cos_iters = self.max_iters - self.warmup_iters
+            cos_last_epoch = self.last_epoch - self.warmup_iters
+            cos_factor = (1 + math.cos(math.pi * cos_last_epoch / cos_iters)) / 2.
+            lr_group = [
+                (base_lr - self.eta_min) * cos_factor
+                for base_lr in self.base_lrs
+            ]
+        return lr_group