returnn 1.20241018.170937__tar.gz → 1.20241018.213651__tar.gz

{returnn-1.20241018.170937 → returnn-1.20241018.213651}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20241018.170937
+Version: 1.20241018.213651
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

returnn-1.20241018.213651/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20241018.213651'
+long_version = '1.20241018.213651+git.2ac447f'

{returnn-1.20241018.170937 → returnn-1.20241018.213651}/returnn/torch/engine.py

@@ -252,7 +252,9 @@ class Engine(EngineBase):
 
         # Update learning rate
         self._updater.set_learning_rate(self.learning_rate)
-        self._updater.set_current_train_step(global_train_step=self.global_train_step, epoch=self.epoch)
+        self._updater.set_current_train_step(
+            global_train_step=self.global_train_step, epoch=self.epoch, epoch_continuous=self.epoch - 1
+        )
 
         self.learning_rate_control.epoch_data[self.epoch].meta.update(
             {
@@ -453,7 +455,11 @@ class Engine(EngineBase):
 
                 step_idx += 1
                 self.global_train_step += 1
-                self._updater.set_current_train_step(global_train_step=self.global_train_step, epoch=self.epoch)
+                self._updater.set_current_train_step(
+                    global_train_step=self.global_train_step,
+                    epoch=self.epoch,
+                    epoch_continuous=(self.epoch - 1 + (last_seq_idx + 1) / num_seqs) if num_seqs is not None else None,
+                )
         except Exception as exc:
             help_on_torch_exception(exc, step_idx=step_idx, model=self._orig_model, extern_data=extern_data)
             raise

{returnn-1.20241018.170937 → returnn-1.20241018.213651}/returnn/torch/updater.py

@@ -13,7 +13,7 @@ import typing
 
 import returnn
 from returnn.log import log
-from returnn.util.basic import RefIdEq
+from returnn.util.basic import RefIdEq, get_fwd_compat_kwargs
 import returnn.frontend as rf
 from returnn.torch.frontend.bridge import wrapped_pt_module_to_rf_module
 
@@ -96,8 +96,11 @@ class Updater:
         self._effective_learning_rate = self.learning_rate
         self.network = network
         self._device = device
+        # Just set the very first step as initial values here.
+        # They will be overwritten via set_current_train_step() below.
         self._current_train_step = 0
-        self._current_epoch = 0
+        self._current_epoch = 1
+        self._current_epoch_continuous = 0.0
 
         self.learning_rate_function = self.config.typed_value("dynamic_learning_rate", None)
         if self.learning_rate_function is not None:
@@ -163,19 +166,38 @@ class Updater:
         self._effective_learning_rate = self.learning_rate
         if self.learning_rate_function is not None:
             lr = self.learning_rate_function(
-                global_train_step=self._current_train_step, epoch=self._current_epoch, learning_rate=self.learning_rate
+                global_train_step=self._current_train_step,
+                epoch=self._current_epoch,
+                epoch_continuous=self._current_epoch_continuous,
+                learning_rate=self.learning_rate,
+                **get_fwd_compat_kwargs(),
             )
             self._effective_learning_rate = float(lr)
         if self.optimizer:
             for param_group in self.optimizer.param_groups:
                 param_group["lr"] = self._effective_learning_rate
 
-    def set_current_train_step(self, *, global_train_step: int, epoch: int):
+    def set_current_train_step(self, *, global_train_step: int, epoch: int, epoch_continuous: Optional[float] = None):
         """
         Obtains an updated learning rate for the current training step inside a (sub)epoch.
+
+        :param global_train_step: Current global training step over the whole training process.
+            In the first epoch, this starts at 0.
+        :param epoch: Current epoch. (First epoch is 1 by RETURNN convention.)
+        :param epoch_continuous: How much of the epoch is finished.
+            In the first step of the first epoch, this starts at 0.0,
+            and when the fist epoch is finished, this reaches 1.0,
+            and the values in between are the fraction of the epoch that is finished.
+            The second epoch (epoch=2) starts at 1.0,
+            and when the second epoch is finished, this reaches 2.0, and so on.
+            We usually calculate this based on ``epoch-1+(last_seq_idx+1)/num_seqs``,
+            if the dataset can provide ``num_seqs``.
+            Other schemes based on the step_idx might be used as well to calculate this,
+            if the number of steps per epoch is known in advance.
         """
         self._current_train_step = global_train_step
         self._current_epoch = epoch
+        self._current_epoch_continuous = epoch_continuous
         self._update_effective_learning_rate()
 
     def step(self, *, grad_scaler: Optional[torch.cuda.amp.GradScaler] = None):

{returnn-1.20241018.170937 → returnn-1.20241018.213651}/returnn.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: returnn
-Version: 1.20241018.170937
+Version: 1.20241018.213651
 Summary: The RWTH extensible training framework for universal recurrent neural networks
 Home-page: https://github.com/rwth-i6/returnn/
 Author: Albert Zeyer

{returnn-1.20241018.170937 → returnn-1.20241018.213651}/tests/test_torch_engine.py

@@ -4,6 +4,7 @@ Tests for PyTorch engine.
 
 from __future__ import annotations
 import _setup_test_env  # noqa
+from typing import Optional
 import sys
 import unittest
 import tempfile
@@ -294,7 +295,6 @@ def test_forward_beam_seq_lens():
 
 
 def test_min_seq_len():
-
     from returnn.datasets.generating import DummyDataset
 
     config = Config({"min_seq_length": 2, "batch_size": 3})
@@ -318,7 +318,6 @@ def test_min_seq_len():
 
 
 def test_max_seq_len():
-
     from returnn.datasets.generating import DummyDataset
 
     config = Config({"max_seq_length": 4, "batch_size": 3})
@@ -520,6 +519,71 @@ def test_torch_engine_train_exception():
             raise Exception("did not get expected exception")
 
 
+def test_dynamic_learning_rate():
+    num_epochs = 3
+    last_global_train_step: Optional[float] = None
+    last_epoch_continuous: Optional[float] = None
+    epoch_continuous_diffs = []
+
+    def _dynamic_learning_rate(
+        *, global_train_step: int, epoch: int, epoch_continuous: float, learning_rate: float, **_kwargs
+    ) -> float:
+        nonlocal last_global_train_step, last_epoch_continuous
+        assert isinstance(global_train_step, int)
+        assert isinstance(epoch, int)
+        assert isinstance(epoch_continuous, (int, float))
+        assert isinstance(learning_rate, (int, float))
+        print(f"global_train_step: {global_train_step}, epoch: {epoch}, epoch_continuous: {epoch_continuous}")
+        if last_global_train_step is None:
+            assert global_train_step == 0 and epoch == 1
+        else:
+            # The call to this function could be repeated.
+            assert global_train_step in (last_global_train_step, last_global_train_step + 1)
+        if last_epoch_continuous is None:
+            assert epoch_continuous == 0
+        elif global_train_step == last_global_train_step:  # repeated call
+            assert epoch_continuous == last_epoch_continuous
+        else:
+            assert epoch_continuous > last_epoch_continuous
+            assert epoch >= epoch_continuous >= epoch - 1
+            epoch_continuous_diffs.append(epoch_continuous - last_epoch_continuous)
+        last_global_train_step = global_train_step
+        last_epoch_continuous = epoch_continuous
+        return learning_rate * epoch_continuous / num_epochs
+
+    config = Config(
+        dict(
+            task="train",
+            device="cpu",
+            extern_data={"data": {"dim": 9}, "classes": {"dim": 2, "sparse": True}},
+            get_model=TrainTestModel,
+            train_step=TrainTestModel.train_step,
+            batch_size=500,
+            optimizer={"class": "adam"},
+            dynamic_learning_rate=_dynamic_learning_rate,
+            num_epochs=num_epochs,
+        )
+    )
+    num_seqs_per_epoch = 100
+    dataset = init_dataset({"class": "Task12AXDataset", "num_seqs": num_seqs_per_epoch, "name": "train"})
+    dataset.init_seq_order(epoch=1)
+
+    with global_config_ctx(config):
+        engine = Engine(config=config)
+        engine.init_train_from_config(train_data=dataset)
+        engine.train()
+
+    assert last_epoch_continuous == num_epochs
+    assert epoch_continuous_diffs
+    print("epoch continuous diffs:", epoch_continuous_diffs)
+    # Just some sanity check. The exact number here depends on num_seqs_per_epoch, batch_size, etc.
+    assert numpy.min(epoch_continuous_diffs) >= 0.01
+    assert numpy.max(epoch_continuous_diffs) <= 0.1
+    # It's one more (non-repeated) call than num steps (first + very last),
+    # and the diffs is one less, so the length should match final global train step.
+    assert len(epoch_continuous_diffs) == engine.global_train_step
+
+
 if __name__ == "__main__":
     better_exchook.install()
     if len(sys.argv) <= 1:

returnn-1.20241018.170937/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20241018.170937'
-long_version = '1.20241018.170937+git.a8aeea1'