bugfix: Correctly set learning rate from LR Finder when resuming

Author: torzdf

.pylintrc

@@ -36,10 +36,6 @@ persistent=yes
 # Specify a configuration file.
 #rcfile=
 
-# When enabled, pylint would attempt to guess common misconfiguration and emit
-# user-friendly hints instead of false-positive error messages.
-suggestion-mode=yes
-
 # Allow loading of arbitrary C extensions. Extensions are imported into the
 # active Python interpreter and may run arbitrary code.
 unsafe-load-any-extension=no

lib/config.py

@@ -302,7 +302,7 @@ def add_section(self, title: str, info: str) -> None:
         logger.debug("Add section: (title: '%s', info: '%s')", title, info)
         self.defaults[title] = ConfigSection(helptext=info, items=OrderedDict())
 
-    def add_item(self,
+    def add_item(self,  # pylint:disable=too-many-arguments,too-many-positional-arguments
                  section: str | None = None,
                  title: str | None = None,
                  datatype: type = str,
@@ -374,7 +374,7 @@ def add_item(self,
                                                          group=group)
 
     @classmethod
-    def _expand_helptext(cls,
+    def _expand_helptext(cls,  # pylint:disable=too-many-positional-arguments
                          helptext: str,
                          choices: str | list[str],
                          default: ConfigValueType,
@@ -453,7 +453,7 @@ def insert_config_section(self,
         config.set(section, helptext)
         logger.debug("Inserted section: '%s'", section)
 
-    def _insert_config_item(self,
+    def _insert_config_item(self,  # pylint:disable=too-many-positional-arguments
                             section: str,
                             item: str,
                             default: ConfigValueType,
@@ -526,7 +526,7 @@ def _load_config(self) -> None:
 
     def save_config(self) -> None:
         """ Save a config file """
-        logger.info("Updating config at: '%s'", self.configfile)
+        logger.debug("Updating config at: '%s'", self.configfile)
         with open(self.configfile, "w", encoding="utf-8", errors="replace") as f_cfgfile:
             self.config.write(f_cfgfile)
         logger.debug("Updated config at: '%s'", self.configfile)

lib/training/lr_finder.py

@@ -48,7 +48,7 @@ class LearningRateFinder:
     beta: float
         Amount to smooth loss by, for graphing purposes
     """
-    def __init__(self,
+    def __init__(self,  # pylint:disable=too-many-positional-arguments
                  model: ModelBase,
                  config: dict[str, ConfigValueType],
                  feeder: Feeder,
@@ -149,7 +149,7 @@ def _reset_model(self, original_lr: float, new_lr: float) -> None:
         new_lr: float
             The discovered optimal learning rate
         """
-        self._model.state.update_session_config("learning_rate", new_lr)
+        self._model.state.add_lr_finder(new_lr)
         self._model.state.save()
 
         logger.debug("Loading initial weights")

plugins/train/model/_base/model.py

@@ -367,11 +367,13 @@ def _output_summary(self) -> None:
         else:
             # print to logger
             print_fn = lambda x: logger.verbose("%s", x)  #type:ignore[attr-defined]  # noqa[E731]  # pylint:disable=C3001
+        parent = None
         for idx, model in enumerate(get_all_sub_models(self.model)):
             if idx == 0:
                 parent = model
                 continue
             model.summary(line_length=100, print_fn=print_fn)
+        assert parent is not None
         parent.summary(line_length=100, print_fn=print_fn)
 
     def _compile_model(self) -> None:
@@ -453,6 +455,7 @@ def __init__(self,
         self._name = model_name
         self._iterations = 0
         self._mixed_precision_layers: list[str] = []
+        self._lr_finder = -1.0
         self._rebuild_model = False
         self._sessions: dict[int, dict] = {}
         self._lowest_avg_loss: dict[str, float] = {}
@@ -503,6 +506,11 @@ def mixed_precision_layers(self) -> list[str]:
         """list: Layers that can be switched between mixed-float16 and float32. """
         return self._mixed_precision_layers
 
+    @property
+    def lr_finder(self) -> float:
+        """ The value discovered from the learning rate finder. -1 if no value stored """
+        return self._lr_finder
+
     @property
     def model_needs_rebuild(self) -> bool:
         """bool: ``True`` if mixed precision policy has changed so model needs to be rebuilt
@@ -595,6 +603,17 @@ def add_mixed_precision_layers(self, layers: list[str]) -> None:
         logger.debug("Storing mixed precision layers: %s", layers)
         self._mixed_precision_layers = layers
 
+    def add_lr_finder(self, learning_rate: float) -> None:
+        """ Add the optimal discovered learning rate from the learning rate finder
+
+        Parameters
+        ----------
+        learning_rate : float
+            The discovered learning rate
+        """
+        logger.debug("Storing learning rate from LR Finder: %s", learning_rate)
+        self._lr_finder = learning_rate
+
     def _load(self, config_changeable_items: dict) -> None:
         """ Load a state file and set the serialized values to the class instance.
 
@@ -616,6 +635,7 @@ def _load(self, config_changeable_items: dict) -> None:
         self._lowest_avg_loss = state.get("lowest_avg_loss", {})
         self._iterations = state.get("iterations", 0)
         self._mixed_precision_layers = state.get("mixed_precision_layers", [])
+        self._lr_finder = state.get("lr_finder", -1.0)
         self._config = state.get("config", {})
         logger.debug("Loaded state: %s", state)
         self._replace_config(config_changeable_items)
@@ -624,10 +644,12 @@ def save(self) -> None:
         """ Save the state values to the serialized state file. """
         logger.debug("Saving State")
         state = {"name": self._name,
-                 "sessions": self._sessions,
+                 "sessions": {k: v for k, v in self._sessions.items()
+                              if v.get("iterations", 0) > 0},
                  "lowest_avg_loss": self._lowest_avg_loss,
                  "iterations": self._iterations,
                  "mixed_precision_layers": self._mixed_precision_layers,
+                 "lr_finder": self._lr_finder,
                  "config": _CONFIG}
         self._serializer.save(self._filename, state)
         logger.debug("Saved State")

plugins/train/trainer/_base.py

@@ -15,6 +15,7 @@
 import cv2
 import numpy as np
 
+import keras.backend as K
 import tensorflow as tf
 from tensorflow.python.framework import (  # pylint:disable=no-name-in-module
     errors_impl as tf_errors)
@@ -162,9 +163,14 @@ def _handle_lr_finder(self) -> bool:
             success = lrf.find()
             return self._config["lr_finder_mode"] == "graph_and_exit" or not success
 
-        learning_rate = self._model.state.sessions[1]["config"]["learning_rate"]
+        learning_rate = self._model.state.lr_finder
+        if learning_rate < 0.:
+            logger.debug("No learning rate finder rate stored. Not setting")
+            return False
+
         logger.info("Setting learning rate from Learning Rate Finder to %s",
                     f"{learning_rate:.1e}")
+        K.set_value(self._model.model.optimizer.lr, learning_rate)
         return False
 
     def _set_tensorboard(self) -> tf.keras.callbacks.TensorBoard:
@@ -807,7 +813,7 @@ class _Timelapse():  # pylint:disable=too-few-public-methods
     image_paths: dict
         The full paths to the training images for each side of the model
     """
-    def __init__(self,
+    def __init__(self,  # pylint:disable=too-many-positional-arguments
                  model: ModelBase,
                  coverage_ratio: float,
                  image_count: int,