returnn 1.20240712.22135__tar.gz → 1.20240713.225734__tar.gz

{returnn-1.20240712.22135 → returnn-1.20240713.225734}/PKG-INFO

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

returnn-1.20240713.225734/_setup_info_generated.py

@@ -0,0 +1,2 @@
+version = '1.20240713.225734'
+long_version = '1.20240713.225734+git.1f00da6'

{returnn-1.20240712.22135 → returnn-1.20240713.225734}/returnn/frontend/label_smoothing.py

@@ -63,36 +63,55 @@ def label_smoothed_log_prob_gradient(
     """
     :param log_prob: shape [...,D] (not necessarily the same as loss)
     :param smoothing: smoothing factor, for :func:`label_smoothing`
-    :param axis: label axis. uses feature_dim by default
+    :param axis: label axis in ``log_prob`` (D). uses feature_dim by default
     :param exclude_labels: list of labels to exclude from smoothing (e.g. blank)
+    :return: ``log_prob``, but the gradient is smoothed
 
-    Assume some cross-entropy-like loss:
+    Assume some cross-entropy-like loss::
 
-      loss = - sum_i target_prob[i] * log_prob[i] .
+        loss = - sum_i target_prob[i] * log_prob[i] .
 
     The sum is over the label indices i (corresponding to the ``axis`` argument).
-    Then the gradient of loss w.r.t. log_prob[i] is:
+    Then the gradient of loss w.r.t. log_prob[i] is::
 
-      grad_logprob[i] loss = -target_prob[i] .
+        grad_logprob[i] loss = -target_prob[i] .
 
-    We assume that the negative gradient is a probability distribution, and apply :func:`label_smoothing` on it.
+    We assume that the negative gradient is a probability distribution
+    (potentially scaled by some factor, e.g. when you scale the loss by some factor)
+    and apply :func:`label_smoothing` on it.
     More specifically, we apply the same scale and shift as in the :func:`label_smoothing` function
-    via :func:`scaled_gradient`.
+    via :func:`scaled_gradient_ext`.
 
-    Just as a side remark: assume
+    Note, this is also the case for CTC or RNNT loss,
+    that the negative gradient of the loss w.r.t. the log-probabilities
+    is a probability distribution.
 
-      log_prob = log_softmax(z) .
+    Common usage example::
 
-    The gradient of log_softmax is:
+        # E.g. there was some log_softmax, or anything to get log probs.
+        log_probs = model.get_log_probs(...)
 
-      grad_z[j] log_prob[i] = delta(i==j) - softmax(z)[j] .
+        # Now apply label smoothing on the log prob gradients.
+        log_probs = rf.label_smoothed_log_prob_gradient(log_probs, 0.1)
 
-    Then the gradient w.r.t. z[j] is:
+        # E.g. CE, CTC, or similar, any kind of NLL should work.
+        loss = loss_func(log_probs, targets)
+        loss.sum().backward()
 
-      grad_z[j] loss = sum_i (grad_logprob[i] loss) (grad_z[j] logprob[i])
-                     = sum_i -target_prob[i] delta(i==j) + target_prob[i] softmax(z)[j]
-                     = -target_prob[j] + (sum_i target_prob[i]) softmax(z)[j]
-                     = softmax(z)[j] - target_prob[j]    # assuming (sum_i target_prob[i]) == 1
+    Just as a side remark: assume::
+
+        log_prob = log_softmax(z) .
+
+    The gradient of log_softmax is::
+
+        grad_z[j] log_prob[i] = delta(i==j) - softmax(z)[j] .
+
+    Then the gradient w.r.t. z[j] is::
+
+        grad_z[j] loss = sum_i (grad_logprob[i] loss) (grad_z[j] logprob[i])
+                       = sum_i -target_prob[i] delta(i==j) + target_prob[i] softmax(z)[j]
+                       = -target_prob[j] + (sum_i target_prob[i]) softmax(z)[j]
+                       = softmax(z)[j] - target_prob[j]    # assuming (sum_i target_prob[i]) == 1
 
     """
     if not axis:
@@ -100,9 +119,15 @@ def label_smoothed_log_prob_gradient(
         axis = log_prob.feature_dim
     # See formula above for label_smoothing.
     dim = axis.dimension
-    floor_prob = smoothing / (dim - 1)
-    factor = 1.0 - dim * floor_prob
-    if exclude_labels:
+    if not exclude_labels:
+        # See formula in code comments in label_smoothing above.
+        floor_prob = smoothing / (dim - 1)
+        factor = 1.0 - dim * floor_prob
+    else:  # have exclude_labels
+        # Sum of prob over included labels does not change with this factor.
+        # For prob[i] == 1, we still have (1 - smoothing) for the smoothed prob.
+        floor_prob = smoothing / (dim - len(exclude_labels) - 1)
+        factor = 1.0 - (dim - len(exclude_labels)) * floor_prob
         indices = rf.range_over_dim(axis)
         mask = True
         for label in exclude_labels:

{returnn-1.20240712.22135 → returnn-1.20240713.225734}/returnn/torch/util/gradient_checkpoint.py

@@ -530,8 +530,14 @@ def _unregister_custom_saved_tensors_hooks():
         assert thread in _custom_saved_tensors_hooks_registered_threads
         assert (
             not _custom_saved_tensors_hooks_tls_ctx.stack
-            and not _custom_saved_tensors_hooks_tls_ctx.callbacks
+            and (_custom_saved_tensors_hooks_tls_ctx.in_callback or not _custom_saved_tensors_hooks_tls_ctx.callbacks)
             and not _custom_saved_tensors_hooks_tls_ctx.queued_exits
+        ), (
+            f"_unregister_custom_saved_tensors_hooks:"
+            f" stack {_custom_saved_tensors_hooks_tls_ctx.stack},"
+            f" in_callback {_custom_saved_tensors_hooks_tls_ctx.in_callback},"
+            f" callbacks {_custom_saved_tensors_hooks_tls_ctx.callbacks},"
+            f" queued_exits {_custom_saved_tensors_hooks_tls_ctx.queued_exits}"
         )
         _custom_saved_tensors_hooks_tls_ctx.active = False
         _custom_saved_tensors_hooks_registered_threads.remove(thread)
@@ -579,6 +585,7 @@ def _register_custom_saved_tensors_hooks_thread_local_callback(cb: Callable[[],
     Register some thread-local callback function which is called on saved_tensors_hooks __enter__ and __exit__.
     If it returns True, it is kept alive, otherwise removed.
     """
+    assert not _custom_saved_tensors_hooks_tls_ctx.in_callback
     _custom_saved_tensors_hooks_tls_ctx.callbacks.append(cb)
 
 
@@ -590,5 +597,7 @@ def _custom_saved_tensors_hooks_call_callbacks():
         _custom_saved_tensors_hooks_tls_ctx.callbacks = [
             cb for cb in _custom_saved_tensors_hooks_tls_ctx.callbacks if cb()
         ]
+        if not _custom_saved_tensors_hooks_tls_ctx.active:  # was cleaned up by some of the callbacks
+            assert not _custom_saved_tensors_hooks_tls_ctx.callbacks and not _custom_saved_tensors_hooks_tls_ctx.stack
     finally:
         _custom_saved_tensors_hooks_tls_ctx.in_callback = False

{returnn-1.20240712.22135 → returnn-1.20240713.225734}/returnn.egg-info/PKG-INFO

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

{returnn-1.20240712.22135 → returnn-1.20240713.225734}/tests/test_torch_util.py

@@ -151,6 +151,51 @@ def test_gradient_checkpoint_scope():
         torch.testing.assert_allclose(param_post_state[k], param_post_state_[k])
 
 
+def test_gradient_checkpoint_scope_twice():
+    # https://github.com/rwth-i6/returnn/issues/1579
+    shape = (101, 103)
+
+    class _Model(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.var = torch.nn.Parameter(torch.randn(shape))
+            self.input_var = torch.nn.Parameter(torch.randn(shape))
+            self.opt = torch.optim.SGD(self.parameters(), lr=0.1)  # not common to have this here but ok for the test
+
+        def forward(self, x: torch.Tensor) -> torch.Tensor:
+            return self.get_var() * x
+
+        def get_var(self) -> torch.Tensor:
+            with gradient_checkpoint_scope():
+                return self.var + torch.randn(shape)
+
+        def get_input(self) -> torch.Tensor:
+            x = self.input_var
+            with gradient_checkpoint_scope():
+                return x + torch.randn(shape)
+
+        def demo_run(self):
+            self.opt.zero_grad()
+            y = self(self.get_input())
+            loss = y.sum()  # dummy loss
+            del y  # not needed anymore
+            loss.backward()
+            del loss  # not needed anymore
+            self.opt.step()
+
+    orig_gradient_checkpoint_scope_tensor_del_hook = gradient_checkpoint_scope._tensor_del_hook
+    try:
+        # Overwrite this here to trigger the case where the tensor del hook will not do the cleanup.
+        gradient_checkpoint_scope._tensor_del_hook = lambda self: None
+
+        model = _Model()
+        model.demo_run()
+        model.demo_run()
+
+    finally:
+        gradient_checkpoint_scope._tensor_del_hook = orig_gradient_checkpoint_scope_tensor_del_hook
+
+
 def _report_profile(prof: torch.profiler.profiler, check_events=(), *, _size_threshold=100):
     # Note: I tried prof.events(), prof.profiler.kineto_results.events(), prof._memory_profile().timeline,
     # but they all are not really giving me the information I want.

returnn-1.20240712.22135/_setup_info_generated.py

@@ -1,2 +0,0 @@
-version = '1.20240712.022135'
-long_version = '1.20240712.022135+git.3da9c4b'