added state_dict/load_state_dict for ReduceLROnPlateau (#7201)

kahne · apaszke · commit a257bd19a2f9 · 2018-05-10T12:02:28.000+02:00
diff --git a/test/test_optim.py b/test/test_optim.py
@@ -647,6 +647,16 @@ def test_cosine_lr_state_dict(self):
             lambda: CosineAnnealingLR(self.opt, T_max=epochs // 2, eta_min=eta_min / 2),
             epochs=epochs)
 
+    def test_reduce_lr_on_plateau_state_dict(self):
+        scheduler = ReduceLROnPlateau(self.opt, mode='min', factor=0.1, patience=2)
+        for score in [1.0, 2.0, 3.0, 4.0, 3.0, 4.0, 5.0, 3.0, 2.0, 1.0]:
+            scheduler.step(score)
+        scheduler_copy = ReduceLROnPlateau(self.opt, mode='max', factor=0.5, patience=10)
+        scheduler_copy.load_state_dict(scheduler.state_dict())
+        for key in scheduler.__dict__.keys():
+            if key not in {'optimizer', 'is_better'}:
+                self.assertEqual(scheduler.__dict__[key], scheduler_copy.__dict__[key], allow_inf=True)
+
     def _check_scheduler_state_dict(self, constr, constr2, epochs=10):
         scheduler = constr()
         for _ in range(epochs):
diff --git a/torch/optim/lr_scheduler.py b/torch/optim/lr_scheduler.py
@@ -378,3 +378,10 @@ def _init_is_better(self, mode, threshold, threshold_mode):
             self.mode_worse = (-float('inf'))
 
         self.is_better = partial(self._cmp, mode, threshold_mode, threshold)
+
+    def state_dict(self):
+        return {key: value for key, value in self.__dict__.items() if key not in {'optimizer', 'is_better'}}
+
+    def load_state_dict(self, state_dict):
+        self.__dict__.update(state_dict)
+        self._init_is_better(mode=self.mode, threshold=self.threshold, threshold_mode=self.threshold_mode)
