We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
PyTorch `weight_decay` ambiguity with shared params
The code logic of `_get_optimizer_param_groups` branch with activated weight decay can not run in its current form. Iterating over all named modules while at the same time iterating recursively over parameters will yield the same parameter multiple times with a different module reference. 

In my case this was:
```python
<class 'i6_experiments.users.rossenbach.experiments.librispeech.tts_architecture_improvement_23.pytorch_networks.ctc_aligner_v1.Model'>                                                                            
<class 'torch.nn.modules.sparse.Embedding'> 
```
for
`speaker_embedding.weight`

This means we need a completely new logic if we want to exclude some modules.


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `returnn/torch/updater.py`
Content:
```
1 """
2 This module covers the optimizer (SGD, Adam, etc) logic,
3 and model param update logic in general.
4 """
5 
6 from __future__ import annotations
7 
8 import torch
9 import typing
10 
11 from returnn.log import log
12 
13 _OptimizerClassesDictInitialized = False
14 _OptimizerClassesDict = {}
15 
16 
17 def _init_optimizer_classes_dict():
18     """
19     Initializes a global dictionary with all optimizers available in PyTorch.
20     """
21     global _OptimizerClassesDictInitialized
22     if _OptimizerClassesDictInitialized:
23         return
24     _OptimizerClassesDictInitialized = True
25     for name, cls in list(vars(torch.optim).items()):
26         assert isinstance(name, str)
27         # Check if cls is a valid subclass of torch.optim.Optimizer
28         if not isinstance(cls, type) or not issubclass(cls, torch.optim.Optimizer):
29             continue
30         assert name not in _OptimizerClassesDict
31         _OptimizerClassesDict[name.lower()] = cls
32 
33 
34 def get_optimizer_class(class_name):
35     """
36     :param str|function|type[torch.optim.Optimizer] class_name: Optimizer data, e.g. "adam", torch.optim.Adam...
37     :return: Optimizer class
38     :rtype: type[torch.optim.Optimizer]
39     """
40     _init_optimizer_classes_dict()
41     if isinstance(class_name, type):
42         assert issubclass(class_name, torch.optim.Optimizer)
43     elif callable(class_name):
44         class_name = class_name()
45     else:
46         assert isinstance(class_name, str)
47         assert (
48             class_name.lower() in _OptimizerClassesDict
49         ), "%s not found in the available torch optimizers list: %s." % (
50             class_name.lower(),
51             ", ".join("'%s'" % key for key in _OptimizerClassesDict),
52         )
53         class_name = _OptimizerClassesDict[class_name.lower()]
54 
55     return class_name
56 
57 
58 def _get_class_init_kwargs(optim_class):
59     """
60     Obtains the keyword arguments of the class provided as parameter that the user can add to their optimizer.
61 
62     :param type[torch.optim.Optimizer] optim_class: Optimizer class.
63     :return: Keyword arguments of the provided class.
64     :rtype: List[str]
65     """
66     from returnn.util.basic import collect_class_init_kwargs
67 
68     optim_class_init_kwargs = collect_class_init_kwargs(optim_class)
69     # We already provide params by default, remove it so that the user doesn't add it to the optimizer dict.
70     optim_class_init_kwargs.remove("params")
71 
72     return optim_class_init_kwargs
73 
74 
75 class Updater(object):
76     """
77     Wraps a torch.optim.Optimizer, and extends it by some further functionality.
78     """
79 
80     def __init__(self, config, network, initial_learning_rate=1.0):
81         """
82         :param returnn.config.Config config: config defining the training conditions.
83         :param torch.nn.Module network: PyTorch Module defining the network.
84         :param float initial_learning_rate:
85         """
86         self.config = config
87         self.learning_rate = initial_learning_rate
88         self.network = network
89         self.optimizer = None  # type: typing.Optional[torch.optim.Optimizer]
90 
91     def set_learning_rate(self, value):
92         """
93         Updates the learning rate of the optimizer at each (sub)epoch.
94 
95         :param float value: New learning rate.
96         """
97         for param_group in self.optimizer.param_groups:
98             param_group["lr"] = value
99 
100     def get_current_step_learning_rate(self):
101         """
102         Obtains an updated learning rate for the current training step inside a (sub)epoch.
103         """
104         pass
105 
106     def create_optimizer(self):
107         """
108         Creates an optimizer and stores it in self.optimizer.
109         """
110         optimizer_opts = self.config.typed_value("optimizer", None)
111         if optimizer_opts is None:
112             raise ValueError("config field 'optimizer' needs to be set explicitely for the Torch backend")
113         self.optimizer = self._create_optimizer(optimizer_opts)
114 
115     def load_optimizer(self, filename):
116         """
117         Loads a torch.optim.Optimizer from disk and stores it in self.optimizer.
118 
119         :param str filename: File from which to load the optimizer state.
120         """
121         print("Load optimizer %s" % filename, file=log.v4)
122         optimizer_state = torch.load(filename)
123         self.optimizer.load_state_dict(optimizer_state)
124 
125     def save_optimizer(self, filename):
126         """
127         Saves the state of self.optimizer to a file.
128 
129         :param str filename: File in which to save the optimizer state.
130         """
131         print("Save optimizer under %s" % filename, file=log.v4)
132         torch.save(self.optimizer.state_dict(), filename)
133 
134     def get_optimizer(self):
135         """
136         :return: Wrapped optimizer object.
137         :rtype: torch.optim.Optimizer
138         """
139         return self.optimizer
140 
141     def _create_optimizer(self, optimizer_opts):
142         """
143         Returns a valid optimizer considering the dictionary given by the user in the config.
144 
145         :param dict[str]|str optimizer_opts: Optimizer configuration specified by the user.
146             If it's a dict, it must contain "class" with the optimizer name or callable.
147             If it's a str, it must be the optimizer name.
148         :return: A valid optimizer.
149         :rtype: torch.optim.Optimizer
150         """
151         lr = self.learning_rate
152 
153         # If the parameter is already a valid optimizer, return it without further processing
154         if isinstance(optimizer_opts, torch.optim.Optimizer):
155             return optimizer_opts
156         elif callable(optimizer_opts):
157             optimizer_opts = {"class": optimizer_opts}
158         else:
159             if not isinstance(optimizer_opts, dict):
160                 raise ValueError("'optimizer' must of type dict, callable or torch.optim.Optimizer instance.")
161             if "class" not in optimizer_opts:
162                 raise ValueError("'class' field of 'optimizer' dict was not set (use e.g. 'SGD', 'Adam', ...)")
163             optimizer_opts = optimizer_opts.copy()
164 
165         # Resolve the optimizer class
166         optim_class_name = optimizer_opts.pop("class")
167         optim_class = get_optimizer_class(optim_class_name)
168 
169         # Resolve the optimizer arguments
170         opt_kwargs = optimizer_opts.copy()
171         optim_class_init_kwargs = _get_class_init_kwargs(optim_class)
172         # epsilon is named eps in torch.
173         # If the user specified it as epsilon, parse it as eps for the optimizer
174         if "eps" in optim_class_init_kwargs and "epsilon" in opt_kwargs:
175             opt_kwargs["eps"] = opt_kwargs.pop("epsilon")
176         if "learning_rate" in optimizer_opts:
177             raise ValueError("'learning_rate' should be set outside of the 'optimizer' dict.")
178         lr = lr * optimizer_opts.get("learning_rate_multiplier", 1.0)
179         opt_kwargs["lr"] = lr
180 
181         param_groups = self._get_optimizer_param_groups(optim_class, opt_kwargs)
182         optimizer = optim_class(param_groups, **opt_kwargs)
183         print("Optimizer: %s" % optimizer, file=log.v1)
184         assert isinstance(optimizer, torch.optim.Optimizer)
185 
186         return optimizer
187 
188     def _create_default_optimizer(self):
189         """
190         :return: SGD optimizer.
191         :rtype: torch.optim.SGD
192         """
193         print("Create SGD optimizer (default).", file=log.v2)
194         optimizer = torch.optim.SGD(self.network.parameters(), lr=self.learning_rate)
195 
196         return optimizer
197 
198     def _get_optimizer_param_groups(self, optim_class, optimizer_opts):
199         """
200         The weight_decay parameter from AdamW affects the weights of layers such as LayerNorm and Embedding.
201         This function creates a blacklist of network modules and splits the optimizer groups in two:
202         those who will receive weight decay, and those who won't receive it.
203         The weight_decay parameter of the rest of the optimizers is L2 regularization.
204 
205         For further reading, see https://github.com/karpathy/minGPT/pull/24#issuecomment-679316025 and
206         https://discuss.pytorch.org/t/weight-decay-in-the-optimizers-is-a-bad-idea-especially-with-batchnorm/16994.
207 
208         This code is based on https://github.com/karpathy/minGPT (MIT license):
209         https://github.com/karpathy/minGPT/blob/3ed14b2cec0dfdad3f4b2831f2b4a86d11aef150/mingpt/model.py#L136.
210 
211         :param type[torch.optim.Optimizer] optim_class: Optimizer class.
212         :param dict[str] optimizer_opts: Optimizer configuration specified by the user.
213         :return: List of configurations for the different sets of parameters.
214         :rtype: List[Dict[str]]
215         """
216         network_params = self.network.parameters()
217 
218         # By default insert the weight_decay constraints in the optimizer, as this is default PyTorch behavior.
219         # If the user doesn't accept this, throw an error message.
220         assert self.config.bool("decouple_constraints", True), (
221             "L2/weight_decay constraints are decoupled in PyTorch, but "
222             "decouple_constraints=False was explicitly specified in the config."
223         )
224 
225         # Split in parameter groups only if decouple_constraints is set and the optimizer accepts weight_decay.
226         cls_init_kwargs = _get_class_init_kwargs(optim_class)
227         if "weight_decay" not in cls_init_kwargs:
228             assert (
229                 "weight_decay" not in optimizer_opts
230             ), "weight_decay not accepted by the chosen optimizer. Accepted values: %s" % ", ".join(
231                 "%s" % optim_name for optim_name in cls_init_kwargs
232             )
233             return [{"params": network_params}]
234 
235         weight_decay = optimizer_opts.get("weight_decay", 0.0)
236         if not weight_decay:
237             return [{"params": network_params}]
238 
239         # Distinguish between parameters with and without weight_decay/L2 regularization.
240         # Parameters without weight decay: biases + LayerNorm/Embedding layers.
241         wd_params = set()
242         no_wd_params = set()
243         blacklist_wd_modules = (torch.nn.LayerNorm, torch.nn.Embedding)
244         for mn, m in self.network.named_modules():
245             for pn, p in m.named_parameters():
246                 fpn = "%s.%s" % (mn, pn) if mn else pn  # Full param name
247                 if pn.endswith("bias"):
248                     no_wd_params.add(fpn)
249                 elif pn.endswith("weight") and isinstance(m, blacklist_wd_modules):
250                     no_wd_params.add(fpn)
251                 else:
252                     wd_params.add(fpn)
253 
254         param_dict = {pn: p for pn, p in self.network.named_parameters()}
255         optim_groups = [
256             {"params": [param_dict[pn] for pn in sorted(list(wd_params))], "weight_decay": weight_decay},
257             {"params": [param_dict[pn] for pn in sorted(list(no_wd_params))], "weight_decay": 0.0},
258         ]
259 
260         return optim_groups
261
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

diff --git a/returnn/torch/updater.py b/returnn/torch/updater.py
--- a/returnn/torch/updater.py
+++ b/returnn/torch/updater.py
@@ -7,8 +7,10 @@
 
 import torch
 import typing
+from typing import Set
 
 from returnn.log import log
+from returnn.util.basic import RefIdEq
 
 _OptimizerClassesDictInitialized = False
 _OptimizerClassesDict = {}
@@ -241,8 +243,14 @@
         wd_params = set()
         no_wd_params = set()
         blacklist_wd_modules = (torch.nn.LayerNorm, torch.nn.Embedding)
+        # Tracker of visited parameters to only add each parameter once, in case two modules share common parameters.
+        # We need the wrapper class RefIdEq because Parameters are compared by value and not by reference.
+        visited_params: Set[RefIdEq[torch.nn.Parameter]] = set()
         for mn, m in self.network.named_modules():
             for pn, p in m.named_parameters():
+                if RefIdEq(p) in visited_params:
+                    continue
+                visited_params.add(RefIdEq(p))
                 fpn = "%s.%s" % (mn, pn) if mn else pn  # Full param name
                 if pn.endswith("bias"):
                     no_wd_params.add(fpn)

{"golden_diff": "diff --git a/returnn/torch/updater.py b/returnn/torch/updater.py\n--- a/returnn/torch/updater.py\n+++ b/returnn/torch/updater.py\n@@ -7,8 +7,10 @@\n \n import torch\n import typing\n+from typing import Set\n \n from returnn.log import log\n+from returnn.util.basic import RefIdEq\n \n _OptimizerClassesDictInitialized = False\n _OptimizerClassesDict = {}\n@@ -241,8 +243,14 @@\n         wd_params = set()\n         no_wd_params = set()\n         blacklist_wd_modules = (torch.nn.LayerNorm, torch.nn.Embedding)\n+        # Tracker of visited parameters to only add each parameter once, in case two modules share common parameters.\n+        # We need the wrapper class RefIdEq because Parameters are compared by value and not by reference.\n+        visited_params: Set[RefIdEq[torch.nn.Parameter]] = set()\n         for mn, m in self.network.named_modules():\n             for pn, p in m.named_parameters():\n+                if RefIdEq(p) in visited_params:\n+                    continue\n+                visited_params.add(RefIdEq(p))\n                 fpn = \"%s.%s\" % (mn, pn) if mn else pn  # Full param name\n                 if pn.endswith(\"bias\"):\n                     no_wd_params.add(fpn)\n", "issue": "PyTorch `weight_decay` ambiguity with shared params\nThe code logic of `_get_optimizer_param_groups` branch with activated weight decay can not run in its current form. Iterating over all named modules while at the same time iterating recursively over parameters will yield the same parameter multiple times with a different module reference. \r\n\r\nIn my case this was:\r\n```python\r\n<class 'i6_experiments.users.rossenbach.experiments.librispeech.tts_architecture_improvement_23.pytorch_networks.ctc_aligner_v1.Model'>                                                                            \r\n<class 'torch.nn.modules.sparse.Embedding'> \r\n```\r\nfor\r\n`speaker_embedding.weight`\r\n\r\nThis means we need a completely new logic if we want to exclude some modules.\r\n\n", "before_files": [{"content": "\"\"\"\nThis module covers the optimizer (SGD, Adam, etc) logic,\nand model param update logic in general.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport torch\nimport typing\n\nfrom returnn.log import log\n\n_OptimizerClassesDictInitialized = False\n_OptimizerClassesDict = {}\n\n\ndef _init_optimizer_classes_dict():\n    \"\"\"\n    Initializes a global dictionary with all optimizers available in PyTorch.\n    \"\"\"\n    global _OptimizerClassesDictInitialized\n    if _OptimizerClassesDictInitialized:\n        return\n    _OptimizerClassesDictInitialized = True\n    for name, cls in list(vars(torch.optim).items()):\n        assert isinstance(name, str)\n        # Check if cls is a valid subclass of torch.optim.Optimizer\n        if not isinstance(cls, type) or not issubclass(cls, torch.optim.Optimizer):\n            continue\n        assert name not in _OptimizerClassesDict\n        _OptimizerClassesDict[name.lower()] = cls\n\n\ndef get_optimizer_class(class_name):\n    \"\"\"\n    :param str|function|type[torch.optim.Optimizer] class_name: Optimizer data, e.g. \"adam\", torch.optim.Adam...\n    :return: Optimizer class\n    :rtype: type[torch.optim.Optimizer]\n    \"\"\"\n    _init_optimizer_classes_dict()\n    if isinstance(class_name, type):\n        assert issubclass(class_name, torch.optim.Optimizer)\n    elif callable(class_name):\n        class_name = class_name()\n    else:\n        assert isinstance(class_name, str)\n        assert (\n            class_name.lower() in _OptimizerClassesDict\n        ), \"%s not found in the available torch optimizers list: %s.\" % (\n            class_name.lower(),\n            \", \".join(\"'%s'\" % key for key in _OptimizerClassesDict),\n        )\n        class_name = _OptimizerClassesDict[class_name.lower()]\n\n    return class_name\n\n\ndef _get_class_init_kwargs(optim_class):\n    \"\"\"\n    Obtains the keyword arguments of the class provided as parameter that the user can add to their optimizer.\n\n    :param type[torch.optim.Optimizer] optim_class: Optimizer class.\n    :return: Keyword arguments of the provided class.\n    :rtype: List[str]\n    \"\"\"\n    from returnn.util.basic import collect_class_init_kwargs\n\n    optim_class_init_kwargs = collect_class_init_kwargs(optim_class)\n    # We already provide params by default, remove it so that the user doesn't add it to the optimizer dict.\n    optim_class_init_kwargs.remove(\"params\")\n\n    return optim_class_init_kwargs\n\n\nclass Updater(object):\n    \"\"\"\n    Wraps a torch.optim.Optimizer, and extends it by some further functionality.\n    \"\"\"\n\n    def __init__(self, config, network, initial_learning_rate=1.0):\n        \"\"\"\n        :param returnn.config.Config config: config defining the training conditions.\n        :param torch.nn.Module network: PyTorch Module defining the network.\n        :param float initial_learning_rate:\n        \"\"\"\n        self.config = config\n        self.learning_rate = initial_learning_rate\n        self.network = network\n        self.optimizer = None  # type: typing.Optional[torch.optim.Optimizer]\n\n    def set_learning_rate(self, value):\n        \"\"\"\n        Updates the learning rate of the optimizer at each (sub)epoch.\n\n        :param float value: New learning rate.\n        \"\"\"\n        for param_group in self.optimizer.param_groups:\n            param_group[\"lr\"] = value\n\n    def get_current_step_learning_rate(self):\n        \"\"\"\n        Obtains an updated learning rate for the current training step inside a (sub)epoch.\n        \"\"\"\n        pass\n\n    def create_optimizer(self):\n        \"\"\"\n        Creates an optimizer and stores it in self.optimizer.\n        \"\"\"\n        optimizer_opts = self.config.typed_value(\"optimizer\", None)\n        if optimizer_opts is None:\n            raise ValueError(\"config field 'optimizer' needs to be set explicitely for the Torch backend\")\n        self.optimizer = self._create_optimizer(optimizer_opts)\n\n    def load_optimizer(self, filename):\n        \"\"\"\n        Loads a torch.optim.Optimizer from disk and stores it in self.optimizer.\n\n        :param str filename: File from which to load the optimizer state.\n        \"\"\"\n        print(\"Load optimizer %s\" % filename, file=log.v4)\n        optimizer_state = torch.load(filename)\n        self.optimizer.load_state_dict(optimizer_state)\n\n    def save_optimizer(self, filename):\n        \"\"\"\n        Saves the state of self.optimizer to a file.\n\n        :param str filename: File in which to save the optimizer state.\n        \"\"\"\n        print(\"Save optimizer under %s\" % filename, file=log.v4)\n        torch.save(self.optimizer.state_dict(), filename)\n\n    def get_optimizer(self):\n        \"\"\"\n        :return: Wrapped optimizer object.\n        :rtype: torch.optim.Optimizer\n        \"\"\"\n        return self.optimizer\n\n    def _create_optimizer(self, optimizer_opts):\n        \"\"\"\n        Returns a valid optimizer considering the dictionary given by the user in the config.\n\n        :param dict[str]|str optimizer_opts: Optimizer configuration specified by the user.\n            If it's a dict, it must contain \"class\" with the optimizer name or callable.\n            If it's a str, it must be the optimizer name.\n        :return: A valid optimizer.\n        :rtype: torch.optim.Optimizer\n        \"\"\"\n        lr = self.learning_rate\n\n        # If the parameter is already a valid optimizer, return it without further processing\n        if isinstance(optimizer_opts, torch.optim.Optimizer):\n            return optimizer_opts\n        elif callable(optimizer_opts):\n            optimizer_opts = {\"class\": optimizer_opts}\n        else:\n            if not isinstance(optimizer_opts, dict):\n                raise ValueError(\"'optimizer' must of type dict, callable or torch.optim.Optimizer instance.\")\n            if \"class\" not in optimizer_opts:\n                raise ValueError(\"'class' field of 'optimizer' dict was not set (use e.g. 'SGD', 'Adam', ...)\")\n            optimizer_opts = optimizer_opts.copy()\n\n        # Resolve the optimizer class\n        optim_class_name = optimizer_opts.pop(\"class\")\n        optim_class = get_optimizer_class(optim_class_name)\n\n        # Resolve the optimizer arguments\n        opt_kwargs = optimizer_opts.copy()\n        optim_class_init_kwargs = _get_class_init_kwargs(optim_class)\n        # epsilon is named eps in torch.\n        # If the user specified it as epsilon, parse it as eps for the optimizer\n        if \"eps\" in optim_class_init_kwargs and \"epsilon\" in opt_kwargs:\n            opt_kwargs[\"eps\"] = opt_kwargs.pop(\"epsilon\")\n        if \"learning_rate\" in optimizer_opts:\n            raise ValueError(\"'learning_rate' should be set outside of the 'optimizer' dict.\")\n        lr = lr * optimizer_opts.get(\"learning_rate_multiplier\", 1.0)\n        opt_kwargs[\"lr\"] = lr\n\n        param_groups = self._get_optimizer_param_groups(optim_class, opt_kwargs)\n        optimizer = optim_class(param_groups, **opt_kwargs)\n        print(\"Optimizer: %s\" % optimizer, file=log.v1)\n        assert isinstance(optimizer, torch.optim.Optimizer)\n\n        return optimizer\n\n    def _create_default_optimizer(self):\n        \"\"\"\n        :return: SGD optimizer.\n        :rtype: torch.optim.SGD\n        \"\"\"\n        print(\"Create SGD optimizer (default).\", file=log.v2)\n        optimizer = torch.optim.SGD(self.network.parameters(), lr=self.learning_rate)\n\n        return optimizer\n\n    def _get_optimizer_param_groups(self, optim_class, optimizer_opts):\n        \"\"\"\n        The weight_decay parameter from AdamW affects the weights of layers such as LayerNorm and Embedding.\n        This function creates a blacklist of network modules and splits the optimizer groups in two:\n        those who will receive weight decay, and those who won't receive it.\n        The weight_decay parameter of the rest of the optimizers is L2 regularization.\n\n        For further reading, see https://github.com/karpathy/minGPT/pull/24#issuecomment-679316025 and\n        https://discuss.pytorch.org/t/weight-decay-in-the-optimizers-is-a-bad-idea-especially-with-batchnorm/16994.\n\n        This code is based on https://github.com/karpathy/minGPT (MIT license):\n        https://github.com/karpathy/minGPT/blob/3ed14b2cec0dfdad3f4b2831f2b4a86d11aef150/mingpt/model.py#L136.\n\n        :param type[torch.optim.Optimizer] optim_class: Optimizer class.\n        :param dict[str] optimizer_opts: Optimizer configuration specified by the user.\n        :return: List of configurations for the different sets of parameters.\n        :rtype: List[Dict[str]]\n        \"\"\"\n        network_params = self.network.parameters()\n\n        # By default insert the weight_decay constraints in the optimizer, as this is default PyTorch behavior.\n        # If the user doesn't accept this, throw an error message.\n        assert self.config.bool(\"decouple_constraints\", True), (\n            \"L2/weight_decay constraints are decoupled in PyTorch, but \"\n            \"decouple_constraints=False was explicitly specified in the config.\"\n        )\n\n        # Split in parameter groups only if decouple_constraints is set and the optimizer accepts weight_decay.\n        cls_init_kwargs = _get_class_init_kwargs(optim_class)\n        if \"weight_decay\" not in cls_init_kwargs:\n            assert (\n                \"weight_decay\" not in optimizer_opts\n            ), \"weight_decay not accepted by the chosen optimizer. Accepted values: %s\" % \", \".join(\n                \"%s\" % optim_name for optim_name in cls_init_kwargs\n            )\n            return [{\"params\": network_params}]\n\n        weight_decay = optimizer_opts.get(\"weight_decay\", 0.0)\n        if not weight_decay:\n            return [{\"params\": network_params}]\n\n        # Distinguish between parameters with and without weight_decay/L2 regularization.\n        # Parameters without weight decay: biases + LayerNorm/Embedding layers.\n        wd_params = set()\n        no_wd_params = set()\n        blacklist_wd_modules = (torch.nn.LayerNorm, torch.nn.Embedding)\n        for mn, m in self.network.named_modules():\n            for pn, p in m.named_parameters():\n                fpn = \"%s.%s\" % (mn, pn) if mn else pn  # Full param name\n                if pn.endswith(\"bias\"):\n                    no_wd_params.add(fpn)\n                elif pn.endswith(\"weight\") and isinstance(m, blacklist_wd_modules):\n                    no_wd_params.add(fpn)\n                else:\n                    wd_params.add(fpn)\n\n        param_dict = {pn: p for pn, p in self.network.named_parameters()}\n        optim_groups = [\n            {\"params\": [param_dict[pn] for pn in sorted(list(wd_params))], \"weight_decay\": weight_decay},\n            {\"params\": [param_dict[pn] for pn in sorted(list(no_wd_params))], \"weight_decay\": 0.0},\n        ]\n\n        return optim_groups\n", "path": "returnn/torch/updater.py"}], "after_files": [{"content": "\"\"\"\nThis module covers the optimizer (SGD, Adam, etc) logic,\nand model param update logic in general.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport torch\nimport typing\nfrom typing import Set\n\nfrom returnn.log import log\nfrom returnn.util.basic import RefIdEq\n\n_OptimizerClassesDictInitialized = False\n_OptimizerClassesDict = {}\n\n\ndef _init_optimizer_classes_dict():\n    \"\"\"\n    Initializes a global dictionary with all optimizers available in PyTorch.\n    \"\"\"\n    global _OptimizerClassesDictInitialized\n    if _OptimizerClassesDictInitialized:\n        return\n    _OptimizerClassesDictInitialized = True\n    for name, cls in list(vars(torch.optim).items()):\n        assert isinstance(name, str)\n        # Check if cls is a valid subclass of torch.optim.Optimizer\n        if not isinstance(cls, type) or not issubclass(cls, torch.optim.Optimizer):\n            continue\n        assert name not in _OptimizerClassesDict\n        _OptimizerClassesDict[name.lower()] = cls\n\n\ndef get_optimizer_class(class_name):\n    \"\"\"\n    :param str|function|type[torch.optim.Optimizer] class_name: Optimizer data, e.g. \"adam\", torch.optim.Adam...\n    :return: Optimizer class\n    :rtype: type[torch.optim.Optimizer]\n    \"\"\"\n    _init_optimizer_classes_dict()\n    if isinstance(class_name, type):\n        assert issubclass(class_name, torch.optim.Optimizer)\n    elif callable(class_name):\n        class_name = class_name()\n    else:\n        assert isinstance(class_name, str)\n        assert (\n            class_name.lower() in _OptimizerClassesDict\n        ), \"%s not found in the available torch optimizers list: %s.\" % (\n            class_name.lower(),\n            \", \".join(\"'%s'\" % key for key in _OptimizerClassesDict),\n        )\n        class_name = _OptimizerClassesDict[class_name.lower()]\n\n    return class_name\n\n\ndef _get_class_init_kwargs(optim_class):\n    \"\"\"\n    Obtains the keyword arguments of the class provided as parameter that the user can add to their optimizer.\n\n    :param type[torch.optim.Optimizer] optim_class: Optimizer class.\n    :return: Keyword arguments of the provided class.\n    :rtype: List[str]\n    \"\"\"\n    from returnn.util.basic import collect_class_init_kwargs\n\n    optim_class_init_kwargs = collect_class_init_kwargs(optim_class)\n    # We already provide params by default, remove it so that the user doesn't add it to the optimizer dict.\n    optim_class_init_kwargs.remove(\"params\")\n\n    return optim_class_init_kwargs\n\n\nclass Updater(object):\n    \"\"\"\n    Wraps a torch.optim.Optimizer, and extends it by some further functionality.\n    \"\"\"\n\n    def __init__(self, config, network, initial_learning_rate=1.0):\n        \"\"\"\n        :param returnn.config.Config config: config defining the training conditions.\n        :param torch.nn.Module network: PyTorch Module defining the network.\n        :param float initial_learning_rate:\n        \"\"\"\n        self.config = config\n        self.learning_rate = initial_learning_rate\n        self.network = network\n        self.optimizer = None  # type: typing.Optional[torch.optim.Optimizer]\n\n    def set_learning_rate(self, value):\n        \"\"\"\n        Updates the learning rate of the optimizer at each (sub)epoch.\n\n        :param float value: New learning rate.\n        \"\"\"\n        for param_group in self.optimizer.param_groups:\n            param_group[\"lr\"] = value\n\n    def get_current_step_learning_rate(self):\n        \"\"\"\n        Obtains an updated learning rate for the current training step inside a (sub)epoch.\n        \"\"\"\n        pass\n\n    def create_optimizer(self):\n        \"\"\"\n        Creates an optimizer and stores it in self.optimizer.\n        \"\"\"\n        optimizer_opts = self.config.typed_value(\"optimizer\", None)\n        if optimizer_opts is None:\n            raise ValueError(\"config field 'optimizer' needs to be set explicitely for the Torch backend\")\n        self.optimizer = self._create_optimizer(optimizer_opts)\n\n    def load_optimizer(self, filename):\n        \"\"\"\n        Loads a torch.optim.Optimizer from disk and stores it in self.optimizer.\n\n        :param str filename: File from which to load the optimizer state.\n        \"\"\"\n        print(\"Load optimizer %s\" % filename, file=log.v4)\n        optimizer_state = torch.load(filename)\n        self.optimizer.load_state_dict(optimizer_state)\n\n    def save_optimizer(self, filename):\n        \"\"\"\n        Saves the state of self.optimizer to a file.\n\n        :param str filename: File in which to save the optimizer state.\n        \"\"\"\n        print(\"Save optimizer under %s\" % filename, file=log.v4)\n        torch.save(self.optimizer.state_dict(), filename)\n\n    def get_optimizer(self):\n        \"\"\"\n        :return: Wrapped optimizer object.\n        :rtype: torch.optim.Optimizer\n        \"\"\"\n        return self.optimizer\n\n    def _create_optimizer(self, optimizer_opts):\n        \"\"\"\n        Returns a valid optimizer considering the dictionary given by the user in the config.\n\n        :param dict[str]|str optimizer_opts: Optimizer configuration specified by the user.\n            If it's a dict, it must contain \"class\" with the optimizer name or callable.\n            If it's a str, it must be the optimizer name.\n        :return: A valid optimizer.\n        :rtype: torch.optim.Optimizer\n        \"\"\"\n        lr = self.learning_rate\n\n        # If the parameter is already a valid optimizer, return it without further processing\n        if isinstance(optimizer_opts, torch.optim.Optimizer):\n            return optimizer_opts\n        elif callable(optimizer_opts):\n            optimizer_opts = {\"class\": optimizer_opts}\n        else:\n            if not isinstance(optimizer_opts, dict):\n                raise ValueError(\"'optimizer' must of type dict, callable or torch.optim.Optimizer instance.\")\n            if \"class\" not in optimizer_opts:\n                raise ValueError(\"'class' field of 'optimizer' dict was not set (use e.g. 'SGD', 'Adam', ...)\")\n            optimizer_opts = optimizer_opts.copy()\n\n        # Resolve the optimizer class\n        optim_class_name = optimizer_opts.pop(\"class\")\n        optim_class = get_optimizer_class(optim_class_name)\n\n        # Resolve the optimizer arguments\n        opt_kwargs = optimizer_opts.copy()\n        optim_class_init_kwargs = _get_class_init_kwargs(optim_class)\n        # epsilon is named eps in torch.\n        # If the user specified it as epsilon, parse it as eps for the optimizer\n        if \"eps\" in optim_class_init_kwargs and \"epsilon\" in opt_kwargs:\n            opt_kwargs[\"eps\"] = opt_kwargs.pop(\"epsilon\")\n        if \"learning_rate\" in optimizer_opts:\n            raise ValueError(\"'learning_rate' should be set outside of the 'optimizer' dict.\")\n        lr = lr * optimizer_opts.get(\"learning_rate_multiplier\", 1.0)\n        opt_kwargs[\"lr\"] = lr\n\n        param_groups = self._get_optimizer_param_groups(optim_class, opt_kwargs)\n        optimizer = optim_class(param_groups, **opt_kwargs)\n        print(\"Optimizer: %s\" % optimizer, file=log.v1)\n        assert isinstance(optimizer, torch.optim.Optimizer)\n\n        return optimizer\n\n    def _create_default_optimizer(self):\n        \"\"\"\n        :return: SGD optimizer.\n        :rtype: torch.optim.SGD\n        \"\"\"\n        print(\"Create SGD optimizer (default).\", file=log.v2)\n        optimizer = torch.optim.SGD(self.network.parameters(), lr=self.learning_rate)\n\n        return optimizer\n\n    def _get_optimizer_param_groups(self, optim_class, optimizer_opts):\n        \"\"\"\n        The weight_decay parameter from AdamW affects the weights of layers such as LayerNorm and Embedding.\n        This function creates a blacklist of network modules and splits the optimizer groups in two:\n        those who will receive weight decay, and those who won't receive it.\n        The weight_decay parameter of the rest of the optimizers is L2 regularization.\n\n        For further reading, see https://github.com/karpathy/minGPT/pull/24#issuecomment-679316025 and\n        https://discuss.pytorch.org/t/weight-decay-in-the-optimizers-is-a-bad-idea-especially-with-batchnorm/16994.\n\n        This code is based on https://github.com/karpathy/minGPT (MIT license):\n        https://github.com/karpathy/minGPT/blob/3ed14b2cec0dfdad3f4b2831f2b4a86d11aef150/mingpt/model.py#L136.\n\n        :param type[torch.optim.Optimizer] optim_class: Optimizer class.\n        :param dict[str] optimizer_opts: Optimizer configuration specified by the user.\n        :return: List of configurations for the different sets of parameters.\n        :rtype: List[Dict[str]]\n        \"\"\"\n        network_params = self.network.parameters()\n\n        # By default insert the weight_decay constraints in the optimizer, as this is default PyTorch behavior.\n        # If the user doesn't accept this, throw an error message.\n        assert self.config.bool(\"decouple_constraints\", True), (\n            \"L2/weight_decay constraints are decoupled in PyTorch, but \"\n            \"decouple_constraints=False was explicitly specified in the config.\"\n        )\n\n        # Split in parameter groups only if decouple_constraints is set and the optimizer accepts weight_decay.\n        cls_init_kwargs = _get_class_init_kwargs(optim_class)\n        if \"weight_decay\" not in cls_init_kwargs:\n            assert (\n                \"weight_decay\" not in optimizer_opts\n            ), \"weight_decay not accepted by the chosen optimizer. Accepted values: %s\" % \", \".join(\n                \"%s\" % optim_name for optim_name in cls_init_kwargs\n            )\n            return [{\"params\": network_params}]\n\n        weight_decay = optimizer_opts.get(\"weight_decay\", 0.0)\n        if not weight_decay:\n            return [{\"params\": network_params}]\n\n        # Distinguish between parameters with and without weight_decay/L2 regularization.\n        # Parameters without weight decay: biases + LayerNorm/Embedding layers.\n        wd_params = set()\n        no_wd_params = set()\n        blacklist_wd_modules = (torch.nn.LayerNorm, torch.nn.Embedding)\n        # Tracker of visited parameters to only add each parameter once, in case two modules share common parameters.\n        # We need the wrapper class RefIdEq because Parameters are compared by value and not by reference.\n        visited_params: Set[RefIdEq[torch.nn.Parameter]] = set()\n        for mn, m in self.network.named_modules():\n            for pn, p in m.named_parameters():\n                if RefIdEq(p) in visited_params:\n                    continue\n                visited_params.add(RefIdEq(p))\n                fpn = \"%s.%s\" % (mn, pn) if mn else pn  # Full param name\n                if pn.endswith(\"bias\"):\n                    no_wd_params.add(fpn)\n                elif pn.endswith(\"weight\") and isinstance(m, blacklist_wd_modules):\n                    no_wd_params.add(fpn)\n                else:\n                    wd_params.add(fpn)\n\n        param_dict = {pn: p for pn, p in self.network.named_parameters()}\n        optim_groups = [\n            {\"params\": [param_dict[pn] for pn in sorted(list(wd_params))], \"weight_decay\": weight_decay},\n            {\"params\": [param_dict[pn] for pn in sorted(list(no_wd_params))], \"weight_decay\": 0.0},\n        ]\n\n        return optim_groups\n", "path": "returnn/torch/updater.py"}]}