training: Add Learning Rate Warmup

Author: torzdf

lib/cli/args_train.py

@@ -164,6 +164,17 @@ def get_argument_list() -> list[dict[str, T.Any]]:
                 "stop training when you are happy with the previews. However, if you want the "
                 "model to stop automatically at a set number of iterations, you can set that "
                 "value here.")})
+        argument_list.append({
+            "opts": ("-a", "--warmup"),
+            "action": Slider,
+            "min_max": (0, 5000),
+            "rounding": 100,
+            "type": int,
+            "default": 0,
+            "group": _("training"),
+            "help": _(
+                "Learning rate warmup. Linearly increase the learning rate from 0 to the chosen "
+                "target rate over the number of iterations given here. 0 to disable.")})
         argument_list.append({
             "opts": ("-D", "--distribution-strategy"),
             "dest": "distribution_strategy",

lib/gui/wrapper.py

@@ -366,7 +366,7 @@ def _read_stdout(self) -> None:
             if output and self._process_progress_stdout(output):
                 continue
 
-            if output:
+            if output.strip():
                 self._process_training_stdout(output)
                 print(output.rstrip())
 

lib/training/__init__.py

@@ -7,6 +7,7 @@
 from .augmentation import ImageAugmentation
 from .generator import Feeder
 from .lr_finder import LearningRateFinder
+from .lr_warmup import LearningRateWarmup
 from .preview_cv import PreviewBuffer, TriggerType
 
 if T.TYPE_CHECKING:

lib/training/lr_warmup.py

@@ -0,0 +1,98 @@
+#! /usr/env/bin/python3
+""" Handles Learning Rate Warmup when training a model """
+from __future__ import annotations
+
+import logging
+import typing as T
+
+import keras.backend as K
+
+logger = logging.getLogger(__name__)
+
+if T.TYPE_CHECKING:
+    from keras.models import Model
+
+
+class LearningRateWarmup():
+    """ Handles the updating of the model's learning rate during Learning Rate Warmup
+
+    Parameters
+    ----------
+    model : :class:`keras.models.Model`
+        The keras model that is to be trained
+    target_learning_rate : float
+        The final learning rate at the end of warmup
+    steps : int
+        The number of iterations to warmup the learning rate for
+    """
+    def __init__(self, model: Model, target_learning_rate: float, steps: int) -> None:
+        self._model = model
+        self._target_lr = target_learning_rate
+        self._steps = steps
+        self._current_lr = 0.0
+        self._current_step = 0
+        self._reporting_points = [int(self._steps * i / 10) for i in range(11)]
+        logger.debug("Initialized %s", self)
+
+    def __repr__(self) -> str:
+        """ Pretty string representation for logging """
+        params = ", ".join(f"{k}={v}" for k, v in self.__dict__.items())
+        return f"{self.__class__.__name__}({params})"
+
+    @classmethod
+    def _format_notation(cls, value: float) -> str:
+        """ Format a float to scientific notation at 1 decimal place
+
+        Parameters
+        ----------
+        value : float
+            The value to format
+
+        Returns
+        -------
+        str
+            The formatted float in scientific notation at 1 decimal place
+        """
+        return f"{value:.1e}"
+
+    def _set_learning_rate(self) -> None:
+        """ Set the learning rate for the current step """
+        self._current_lr = self._current_step / self._steps * self._target_lr
+        K.set_value(self._model.optimizer.lr, self._current_lr)
+        logger.debug("Learning rate set to %s for step %s/%s",
+                     self._current_lr, self._current_step, self._steps)
+
+    def _output_status(self) -> None:
+        """ Output the progress of Learning Rate Warmup at set intervals """
+        if self._current_step == 1:
+            logger.info("[Learning Rate Warmup] Start: %s, Target: %s, Steps: %s",
+                        self._format_notation(self._current_lr),
+                        self._format_notation(self._target_lr), self._steps)
+            return
+
+        if self._current_step == self._steps:
+            print()
+            logger.info("[Learning Rate Warmup] Final Learning Rate: %s",
+                        self._format_notation(self._target_lr))
+            return
+
+        if self._current_step in self._reporting_points:
+            print()
+            progress = int(round(100 / (len(self._reporting_points) - 1) *
+                           self._reporting_points.index(self._current_step), 0))
+            logger.info("[Learning Rate Warmup] Step: %s/%s (%s), Current: %s, Target: %s",
+                        self._current_step,
+                        self._steps,
+                        f"{progress}%",
+                        self._format_notation(self._current_lr),
+                        self._format_notation(self._target_lr))
+
+    def __call__(self) -> None:
+        """ If a learning rate update is required, update the model's learning rate, otherwise
+        do nothing """
+        if self._steps == 0 or self._current_step >= self._steps:
+            return
+
+        self._current_step += 1
+        self._set_learning_rate()
+        self._output_status()

locales/kr/LC_MESSAGES/lib.cli.args_train.mo

@@ -7,16 +7,16 @@ msgid ""
 msgstr ""
 "Project-Id-Version: \n"
 "Report-Msgid-Bugs-To: \n"
-"POT-Creation-Date: 2024-03-28 18:04+0000\n"
-"PO-Revision-Date: 2024-03-28 18:16+0000\n"
+"POT-Creation-Date: 2025-11-21 17:32+0000\n"
+"PO-Revision-Date: 2025-11-21 17:47+0000\n"
 "Last-Translator: \n"
 "Language-Team: \n"
 "Language: ko_KR\n"
 "MIME-Version: 1.0\n"
 "Content-Type: text/plain; charset=UTF-8\n"
 "Content-Transfer-Encoding: 8bit\n"
 "Plural-Forms: nplurals=1; plural=0;\n"
-"X-Generator: Poedit 3.4.2\n"
+"X-Generator: Poedit 3.6\n"
 
 #: lib/cli/args_train.py:30
 msgid ""
@@ -157,8 +157,8 @@ msgstr ""
 "성 옵션이 있을 수 있습니다."
 
 #: lib/cli/args_train.py:147 lib/cli/args_train.py:160
-#: lib/cli/args_train.py:175 lib/cli/args_train.py:191
-#: lib/cli/args_train.py:200
+#: lib/cli/args_train.py:174 lib/cli/args_train.py:186
+#: lib/cli/args_train.py:202 lib/cli/args_train.py:211
 msgid "training"
 msgstr "훈련"
 
@@ -187,7 +187,15 @@ msgstr ""
 "다. 그러나 설정된 반복 횟수에서 모델이 자동으로 중지되도록 하려면 여기에서 해"
 "당 값을 설정할 수 있습니다."
 
-#: lib/cli/args_train.py:177
+#: lib/cli/args_train.py:176
+msgid ""
+"Learning rate warmup. Linearly increase the learning rate from 0 to the "
+"chosen target rate over the number of iterations given here. 0 to disable."
+msgstr ""
+"학습률 워밍업. 여기에 주어진 반복 횟수에 따라 학습률을 0에서 선택한 목표 속도"
+"까지 선형적으로 증가시킵니다. 0으로 설정하면 비활성화됩니다."
+
+#: lib/cli/args_train.py:188
 msgid ""
 "R|Select the distribution stategy to use.\n"
 "L|default: Use Tensorflow's default distribution strategy.\n"
@@ -208,15 +216,15 @@ msgstr ""
 "L|mirrored: 여러 로컬 GPU에서 동기화 분산 훈련을 지원합니다. 모델의 복사본과 "
 "모든 변수는 각 반복에서 각 GPU에 배포된 배치들와 함께 각 GPU에 로드됩니다."
 
-#: lib/cli/args_train.py:193
+#: lib/cli/args_train.py:204
 msgid ""
 "Disables TensorBoard logging. NB: Disabling logs means that you will not be "
 "able to use the graph or analysis for this session in the GUI."
 msgstr ""
 "텐서보드 로깅을 비활성화합니다. 주의: 로그를 비활성화하면 GUI에서 이 세션에 "
 "대한 그래프 또는 분석을 사용할 수 없습니다."
 
-#: lib/cli/args_train.py:202
+#: lib/cli/args_train.py:213
 msgid ""
 "Use the Learning Rate Finder to discover the optimal learning rate for "
 "training. For new models, this will calculate the optimal learning rate for "
@@ -229,28 +237,28 @@ msgstr ""
 "때 발견된 최적의 학습률을 사용합니다. 이 옵션을 설정하면 수동으로 구성된 학습"
 "률(기차 설정에서 구성 가능)이 무시됩니다."
 
-#: lib/cli/args_train.py:215 lib/cli/args_train.py:225
+#: lib/cli/args_train.py:226 lib/cli/args_train.py:236
 msgid "Saving"
 msgstr "저장"
 
-#: lib/cli/args_train.py:216
+#: lib/cli/args_train.py:227
 msgid "Sets the number of iterations between each model save."
 msgstr "각 모델 저장 사이의 반복 횟수를 설정합니다."
 
-#: lib/cli/args_train.py:227
+#: lib/cli/args_train.py:238
 msgid ""
 "Sets the number of iterations before saving a backup snapshot of the model "
 "in it's current state. Set to 0 for off."
 msgstr ""
 "현재 상태에서 모델의 백업 스냅샷을 저장하기 전에 반복할 횟수를 설정합니다. 0"
 "으로 설정하면 꺼집니다."
 
-#: lib/cli/args_train.py:234 lib/cli/args_train.py:246
-#: lib/cli/args_train.py:258
+#: lib/cli/args_train.py:245 lib/cli/args_train.py:257
+#: lib/cli/args_train.py:269
 msgid "timelapse"
 msgstr "타임랩스"
 
-#: lib/cli/args_train.py:236
+#: lib/cli/args_train.py:247
 msgid ""
 "Optional for creating a timelapse. Timelapse will save an image of your "
 "selected faces into the timelapse-output folder at every save iteration. "
@@ -263,7 +271,7 @@ msgstr ""
 "랩스를 만드는 데 사용할 'A' 얼굴의 입력 폴더여야 합니다. 또한 사용자는 --"
 "timelapse-output 및 --timelapse-input-B 매개 변수를 제공해야 합니다."
 
-#: lib/cli/args_train.py:248
+#: lib/cli/args_train.py:259
 msgid ""
 "Optional for creating a timelapse. Timelapse will save an image of your "
 "selected faces into the timelapse-output folder at every save iteration. "
@@ -276,7 +284,7 @@ msgstr ""
 "다. 타임 랩스를 만드는 데 사용할 'B' 얼굴의 입력 폴더여야 합니다. 또한 사용자"
 "는 --timelapse-output 및 --timelapse-input-A 매개 변수를 제공해야 합니다."
 
-#: lib/cli/args_train.py:260
+#: lib/cli/args_train.py:271
 msgid ""
 "Optional for creating a timelapse. Timelapse will save an image of your "
 "selected faces into the timelapse-output folder at every save iteration. If "
@@ -288,27 +296,27 @@ msgstr ""
 "다. 입력 폴더가 제공되었지만 출력 폴더가 없는 경우 모델 폴더에/timelapse/로 "
 "기본 설정됩니다"
 
-#: lib/cli/args_train.py:269 lib/cli/args_train.py:276
+#: lib/cli/args_train.py:280 lib/cli/args_train.py:287
 msgid "preview"
 msgstr "미리보기"
 
-#: lib/cli/args_train.py:270
+#: lib/cli/args_train.py:281
 msgid "Show training preview output. in a separate window."
 msgstr "훈련 미리보기 결과를 각기 다른 창에서 보여줍니다."
 
-#: lib/cli/args_train.py:278
+#: lib/cli/args_train.py:289
 msgid ""
 "Writes the training result to a file. The image will be stored in the root "
 "of your FaceSwap folder."
 msgstr ""
 "훈련 결과를 파일에 씁니다. 이미지는 Faceswap 폴더의 최상위 폴더에 저장됩니다."
 
-#: lib/cli/args_train.py:285 lib/cli/args_train.py:295
-#: lib/cli/args_train.py:305 lib/cli/args_train.py:315
+#: lib/cli/args_train.py:296 lib/cli/args_train.py:306
+#: lib/cli/args_train.py:316 lib/cli/args_train.py:326
 msgid "augmentation"
 msgstr "보정"
 
-#: lib/cli/args_train.py:287
+#: lib/cli/args_train.py:298
 msgid ""
 "Warps training faces to closely matched Landmarks from the opposite face-set "
 "rather than randomly warping the face. This is the 'dfaker' way of doing "
@@ -317,7 +325,7 @@ msgstr ""
 "무작위로 얼굴을 변환하지 않고 반대쪽 얼굴 세트에서 특징점과 밀접하게 일치하도"
 "록 훈련 얼굴을 변환해줍니다. 이것은 변환하는 'dfaker' 방식이다."
 
-#: lib/cli/args_train.py:297
+#: lib/cli/args_train.py:308
 msgid ""
 "To effectively learn, a random set of images are flipped horizontally. "
 "Sometimes it is desirable for this not to occur. Generally this should be "
@@ -327,7 +335,7 @@ msgstr ""
 "런 일이 일어나지 않는 것이 바람직합니다. 일반적으로 'fit training' 중을 제외"
 "하고는 이 작업을 중단해야 합니다."
 
-#: lib/cli/args_train.py:307
+#: lib/cli/args_train.py:318
 msgid ""
 "Color augmentation helps make the model less susceptible to color "
 "differences between the A and B sets, at an increased training time cost. "
@@ -337,7 +345,7 @@ msgstr ""
 "이 되며, 훈련 시간 비용이 증가합니다. 색상 보저를 사용하지 않으려면 이 옵션"
 "을 사용합니다."
 
-#: lib/cli/args_train.py:317
+#: lib/cli/args_train.py:328
 msgid ""
 "Warping is integral to training the Neural Network. This option should only "
 "be enabled towards the very end of training to try to bring out more detail. "