implement weight dropout

Fix #100

Author: albertz

nn/utils/__init__.py

@@ -10,4 +10,5 @@
 from .stochastic_depth import *
 from .targets import *
 from .variational_weight_noise import *
+from .weight_dropout import *
 from .weight_norm import weight_norm, remove_weight_norm

nn/utils/weight_dropout.py

@@ -0,0 +1,40 @@
+"""
+Weight dropout.
+
+Also known as "variational dropout" or "Bayesian dropout",
+sometimes applied for LSTM weights,
+but this can also be applied to any other weights.
+
+https://github.com/rwth-i6/returnn_common/issues/100
+"""
+
+from __future__ import annotations
+from typing import Optional, Union, Sequence, TypeVar
+from ... import nn
+
+
+T_module = TypeVar('T_module', bound=nn.Module)
+
+
+def weight_dropout(
+  module: T_module, name: str, dropout: float,
+  *,
+  axis: Optional[Union[nn.Dim, Sequence[nn.Dim]]] = None,
+) -> T_module:
+  """
+  :param module: module
+  :param name: name of the weight parameter
+  :param dropout: dropout probability
+  :param axis: axis to apply dropout on. see :func:`nn.dropout`
+  """
+  assert hasattr(module, name)
+  weight = getattr(module, name)
+  assert isinstance(weight, nn.Parameter)
+  if not axis:
+    axis = weight.shape_ordered
+
+  assert not hasattr(module, f"{name}_raw")
+  setattr(module, f"{name}_raw", weight)
+  weight = nn.dropout(weight, dropout, axis=axis)
+  setattr(module, name, weight)
+  return module

tests/test_nn_utils.py

@@ -56,3 +56,19 @@ def test_variational_weight_noise():
   y.mark_as_default_output()
   config_str = nn.get_returnn_config().get_complete_py_code_str(net)
   dummy_run_net_single_custom(config_str, train_flag=True)
+
+
+def test_weight_dropout():
+  nn.reset_default_root_name_ctx()
+  time_dim = nn.SpatialDim("time")
+  in_dim = nn.FeatureDim("in", 3)
+  x = nn.Data("data", dim_tags=[nn.batch_dim, time_dim, in_dim])
+  x = nn.get_extern_data(x)
+  net = nn.Linear(in_dim, nn.FeatureDim("out", 5))
+  assert isinstance(net.weight, nn.Parameter)
+  nn.weight_dropout(net, "weight", 0.3)
+  assert not isinstance(net.weight, nn.Parameter) and isinstance(net.weight, nn.Tensor)
+  y = net(x)
+  y.mark_as_default_output()
+  config_str = nn.get_returnn_config().get_complete_py_code_str(net)
+  dummy_run_net_single_custom(config_str, train_flag=True)