implement variational weight noise

Fix #240

Author: albertz

nn/utils/__init__.py

@@ -9,4 +9,5 @@
 from .label_smoothing import *
 from .stochastic_depth import *
 from .targets import *
+from .variational_weight_noise import *
 from .weight_norm import weight_norm, remove_weight_norm

nn/utils/variational_weight_noise.py

@@ -0,0 +1,33 @@
+"""
+Variational weight noise
+
+https://github.com/rwth-i6/returnn_common/issues/240
+"""
+
+from __future__ import annotations
+from typing import TypeVar
+from ... import nn
+
+
+T_module = TypeVar('T_module', bound=nn.Module)
+
+
+def variational_weight_noise(module: T_module, name: str, weight_noise_std: float) -> T_module:
+  """
+  :param module: module
+  :param name: name of the weight parameter
+  :param weight_noise_std: standard deviation of the weight noise
+  """
+  assert weight_noise_std > 0
+  assert hasattr(module, name)
+  weight = getattr(module, name)
+  assert isinstance(weight, nn.Parameter)
+
+  setattr(module, f"{name}_raw", weight)
+
+  with nn.Cond(nn.train_flag()) as cond:
+    weight_noise = nn.random_normal(weight.shape_ordered, weight.dtype, stddev=weight_noise_std)
+    cond.true = weight + weight_noise
+    cond.false = weight
+  setattr(module, name, cond.result)
+  return module

tests/test_nn_utils.py

@@ -40,3 +40,19 @@ def test_weight_norm():
   y.mark_as_default_output()
   config_str = nn.get_returnn_config().get_complete_py_code_str(net)
   dummy_run_net_single_custom(config_str, eval_flag=True)
+
+
+def test_variational_weight_noise():
+  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.variational_weight_noise(net, "weight", 0.075)
+  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)