We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN ISSUE ---
[BUG] FedOpt algorithm not working as expected in cifar10 example
**Describe the bug**
The FedOpt algorithm is not working as expected in cifar10 example when I change the model from the pre-existing ModerateCNN to another model like MobileNetv2 or Resnet18 and others. The problem is that the accuracy of the global model is not increasing or increasing too slow with the FedOpt algorithm while the other algorithms works just fine even changing the model.

**To Reproduce**
1. Add in 'cifar10_nets.py' the new model : 
   class MyModel(nn.Module):
     def __init__(self):
        super(MyModel, self).__init__()
        model = models.mobilenet_v2(weights='DEFAULT')
        model.classifier = nn.Sequential(
            nn.Dropout(0.4),
            nn.Linear(1280, 10),
        )
        self.model = model
    
    def forward(self, x):
        return self.model(x)
2. Import and change the model in file 'cifar10_learner.py'
3. Launch the example with `./run_simulator.sh cifar10_fedopt 0.1 8 8`
4. See the results in tensorboard with `tensorboard --logdir=/tmp/nvflare/sim_cifar10` under the section 'val_acc_global_model'

**Expected behavior**
I expect reading the algorithm proposed in Reddi, Sashank, et al. "Adaptive federated optimization." arXiv preprint arXiv:2003.00295 (2020), to obtain the same performance of FedAvg using SGD optimizer with lr = 1.0 and no scheduler. Also obtain better results changing optimizer and adding a scheduler.

**Screenshots**
![Screenshot from 2023-04-28 10-09-15](https://user-images.githubusercontent.com/102248925/235092305-891cc9b1-e0d6-48dd-90f1-48898d6b918b.png)
Purple = FedAvg
Pink = FedOpt

**Desktop (please complete the following information):**
 - OS: ubuntu 22.04
 - Python Version 3.10
 - NVFlare Version 2.3.0

Ty in advance!


--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `nvflare/app_opt/pt/fedopt.py`
Content:
```
1 # Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.
2 #
3 # Licensed under the Apache License, Version 2.0 (the "License");
4 # you may not use this file except in compliance with the License.
5 # You may obtain a copy of the License at
6 #
7 #     http://www.apache.org/licenses/LICENSE-2.0
8 #
9 # Unless required by applicable law or agreed to in writing, software
10 # distributed under the License is distributed on an "AS IS" BASIS,
11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 # See the License for the specific language governing permissions and
13 # limitations under the License.
14 
15 import time
16 
17 import torch
18 
19 from nvflare.apis.dxo import DataKind, MetaKey, from_shareable
20 from nvflare.apis.event_type import EventType
21 from nvflare.apis.fl_context import FLContext
22 from nvflare.apis.shareable import Shareable
23 from nvflare.app_common.abstract.learnable import Learnable
24 from nvflare.app_common.abstract.model import make_model_learnable
25 from nvflare.app_common.app_constant import AppConstants
26 from nvflare.app_common.shareablegenerators.full_model_shareable_generator import FullModelShareableGenerator
27 from nvflare.security.logging import secure_format_exception
28 
29 
30 class PTFedOptModelShareableGenerator(FullModelShareableGenerator):
31     def __init__(
32         self,
33         optimizer_args: dict = None,
34         lr_scheduler_args: dict = None,
35         source_model="model",
36         device=None,
37     ):
38         """Implement the FedOpt algorithm.
39 
40         The algorithm is proposed in Reddi, Sashank, et al. "Adaptive federated optimization." arXiv preprint arXiv:2003.00295 (2020).
41         This SharableGenerator will update the global model using the specified
42         PyTorch optimizer and learning rate scheduler.
43 
44         Args:
45             optimizer_args: dictionary of optimizer arguments, e.g.
46                 {'path': 'torch.optim.SGD', 'args': {'lr': 1.0}} (default).
47             lr_scheduler_args: dictionary of server-side learning rate scheduler arguments, e.g.
48                 {'path': 'torch.optim.lr_scheduler.CosineAnnealingLR', 'args': {'T_max': 100}} (default: None).
49             source_model: either a valid torch model object or a component ID of a torch model object
50             device: specify the device to run server-side optimization, e.g. "cpu" or "cuda:0"
51                 (will default to cuda if available and no device is specified).
52 
53         Raises:
54             TypeError: when any of input arguments does not have correct type
55         """
56         super().__init__()
57         if not optimizer_args:
58             self.logger("No optimizer_args provided. Using FedOpt with SGD and lr 1.0")
59             optimizer_args = {"name": "SGD", "args": {"lr": 1.0}}
60 
61         if not isinstance(optimizer_args, dict):
62             raise TypeError(
63                 "optimizer_args must be a dict of format, e.g. {'path': 'torch.optim.SGD', 'args': {'lr': 1.0}}."
64             )
65         if lr_scheduler_args is not None:
66             if not isinstance(lr_scheduler_args, dict):
67                 raise TypeError(
68                     "optimizer_args must be a dict of format, e.g. "
69                     "{'path': 'torch.optim.lr_scheduler.CosineAnnealingLR', 'args': {'T_max': 100}}."
70                 )
71         self.source_model = source_model
72         self.optimizer_args = optimizer_args
73         self.lr_scheduler_args = lr_scheduler_args
74         self.model = None
75         self.optimizer = None
76         self.lr_scheduler = None
77         if device is None:
78             self.device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
79         else:
80             self.device = torch.device(device)
81         self.optimizer_name = None
82         self.lr_scheduler_name = None
83 
84     def _get_component_name(self, component_args):
85         if component_args is not None:
86             name = component_args.get("path", None)
87             if name is None:
88                 name = component_args.get("name", None)
89             return name
90         else:
91             return None
92 
93     def handle_event(self, event_type: str, fl_ctx: FLContext):
94         if event_type == EventType.START_RUN:
95             # Initialize the optimizer with current global model params
96             engine = fl_ctx.get_engine()
97 
98             if isinstance(self.source_model, str):
99                 self.model = engine.get_component(self.source_model)
100             else:
101                 self.model = self.source_model
102 
103             if self.model is None:
104                 self.system_panic(
105                     "Model is not available",
106                     fl_ctx,
107                 )
108                 return
109             elif not isinstance(self.model, torch.nn.Module):
110                 self.system_panic(
111                     f"Expected model to be a torch.nn.Module but got {type(self.model)}",
112                     fl_ctx,
113                 )
114                 return
115             else:
116                 print("server model", self.model)
117 
118             self.model.to(self.device)
119 
120             # set up optimizer
121             try:
122                 # use provided or default optimizer arguments and add the model parameters
123                 if "args" not in self.optimizer_args:
124                     self.optimizer_args["args"] = {}
125                 self.optimizer_args["args"]["params"] = self.model.parameters()
126                 self.optimizer = engine.build_component(self.optimizer_args)
127                 # get optimizer name for log
128                 self.optimizer_name = self._get_component_name(self.optimizer_args)
129             except Exception as e:
130                 self.system_panic(
131                     f"Exception while parsing `optimizer_args`({self.optimizer_args}): {secure_format_exception(e)}",
132                     fl_ctx,
133                 )
134                 return
135 
136             # set up lr scheduler
137             if self.lr_scheduler_args is not None:
138                 try:
139                     self.lr_scheduler_name = self._get_component_name(self.lr_scheduler_args)
140                     # use provided or default lr scheduler argument and add the optimizer
141                     if "args" not in self.lr_scheduler_args:
142                         self.lr_scheduler_args["args"] = {}
143                     self.lr_scheduler_args["args"]["optimizer"] = self.optimizer
144                     self.lr_scheduler = engine.build_component(self.lr_scheduler_args)
145                 except Exception as e:
146                     self.system_panic(
147                         f"Exception while parsing `lr_scheduler_args`({self.lr_scheduler_args}): {secure_format_exception(e)}",
148                         fl_ctx,
149                     )
150                     return
151 
152     def server_update(self, model_diff):
153         """Updates the global model using the specified optimizer.
154 
155         Args:
156             model_diff: the aggregated model differences from clients.
157 
158         Returns:
159             The updated PyTorch model state dictionary.
160 
161         """
162         self.model.train()
163         self.optimizer.zero_grad()
164 
165         # Apply the update to the model. We must multiply weights_delta by -1.0 to
166         # view it as a gradient that should be applied to the server_optimizer.
167         for name, param in self.model.named_parameters():
168             param.grad = torch.tensor(-1.0 * model_diff[name]).to(self.device)
169 
170         self.optimizer.step()
171         if self.lr_scheduler is not None:
172             self.lr_scheduler.step()
173 
174         return self.model.state_dict()
175 
176     def shareable_to_learnable(self, shareable: Shareable, fl_ctx: FLContext) -> Learnable:
177         """Convert Shareable to Learnable while doing a FedOpt update step.
178 
179         Supporting data_kind == DataKind.WEIGHT_DIFF
180 
181         Args:
182             shareable (Shareable): Shareable to be converted
183             fl_ctx (FLContext): FL context
184 
185         Returns:
186             Model: Updated global ModelLearnable.
187         """
188         # check types
189         dxo = from_shareable(shareable)
190 
191         if dxo.data_kind != DataKind.WEIGHT_DIFF:
192             self.system_panic(
193                 "FedOpt is only implemented for " "data_kind == DataKind.WEIGHT_DIFF",
194                 fl_ctx,
195             )
196             return Learnable()
197 
198         processed_algorithm = dxo.get_meta_prop(MetaKey.PROCESSED_ALGORITHM)
199         if processed_algorithm is not None:
200             self.system_panic(
201                 f"FedOpt is not implemented for shareable processed by {processed_algorithm}",
202                 fl_ctx,
203             )
204             return Learnable()
205 
206         model_diff = dxo.data
207 
208         start = time.time()
209         weights = self.server_update(model_diff)
210         secs = time.time() - start
211 
212         # convert to numpy dict of weights
213         start = time.time()
214         for key in weights:
215             weights[key] = weights[key].detach().cpu().numpy()
216         secs_detach = time.time() - start
217 
218         self.log_info(
219             fl_ctx,
220             f"FedOpt ({self.optimizer_name}, {self.device}) server model update "
221             f"round {fl_ctx.get_prop(AppConstants.CURRENT_ROUND)}, "
222             f"{self.lr_scheduler_name if self.lr_scheduler_name else ''} "
223             f"lr: {self.optimizer.param_groups[-1]['lr']}, "
224             f"update: {secs} secs., detach: {secs_detach} secs.",
225         )
226         # TODO: write server-side lr to tensorboard
227 
228         return make_model_learnable(weights, dxo.get_meta_props())
229
```
--- 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/nvflare/app_opt/pt/fedopt.py b/nvflare/app_opt/pt/fedopt.py
--- a/nvflare/app_opt/pt/fedopt.py
+++ b/nvflare/app_opt/pt/fedopt.py
@@ -21,7 +21,7 @@
 from nvflare.apis.fl_context import FLContext
 from nvflare.apis.shareable import Shareable
 from nvflare.app_common.abstract.learnable import Learnable
-from nvflare.app_common.abstract.model import make_model_learnable
+from nvflare.app_common.abstract.model import ModelLearnableKey, make_model_learnable
 from nvflare.app_common.app_constant import AppConstants
 from nvflare.app_common.shareablegenerators.full_model_shareable_generator import FullModelShareableGenerator
 from nvflare.security.logging import secure_format_exception
@@ -40,6 +40,8 @@
         The algorithm is proposed in Reddi, Sashank, et al. "Adaptive federated optimization." arXiv preprint arXiv:2003.00295 (2020).
         This SharableGenerator will update the global model using the specified
         PyTorch optimizer and learning rate scheduler.
+        Note: This class will use FedOpt to optimize the global trainable parameters (i.e. `self.model.named_parameters()`)
+        but use FedAvg to update any other layers such as batch norm statistics.
 
         Args:
             optimizer_args: dictionary of optimizer arguments, e.g.
@@ -164,14 +166,17 @@
 
         # Apply the update to the model. We must multiply weights_delta by -1.0 to
         # view it as a gradient that should be applied to the server_optimizer.
+        updated_params = []
         for name, param in self.model.named_parameters():
-            param.grad = torch.tensor(-1.0 * model_diff[name]).to(self.device)
+            if name in model_diff:
+                param.grad = torch.tensor(-1.0 * model_diff[name]).to(self.device)
+                updated_params.append(name)
 
         self.optimizer.step()
         if self.lr_scheduler is not None:
             self.lr_scheduler.step()
 
-        return self.model.state_dict()
+        return self.model.state_dict(), updated_params
 
     def shareable_to_learnable(self, shareable: Shareable, fl_ctx: FLContext) -> Learnable:
         """Convert Shareable to Learnable while doing a FedOpt update step.
@@ -206,7 +211,7 @@
         model_diff = dxo.data
 
         start = time.time()
-        weights = self.server_update(model_diff)
+        weights, updated_params = self.server_update(model_diff)
         secs = time.time() - start
 
         # convert to numpy dict of weights
@@ -215,12 +220,28 @@
             weights[key] = weights[key].detach().cpu().numpy()
         secs_detach = time.time() - start
 
+        # update unnamed parameters such as batch norm layers if there are any using the averaged update
+        base_model = fl_ctx.get_prop(AppConstants.GLOBAL_MODEL)
+        if not base_model:
+            self.system_panic(reason="No global base model!", fl_ctx=fl_ctx)
+            return base_model
+
+        base_model_weights = base_model[ModelLearnableKey.WEIGHTS]
+
+        n_fedavg = 0
+        for key, value in model_diff.items():
+            if key not in updated_params:
+                weights[key] = base_model_weights[key] + value
+                n_fedavg += 1
+
         self.log_info(
             fl_ctx,
             f"FedOpt ({self.optimizer_name}, {self.device}) server model update "
             f"round {fl_ctx.get_prop(AppConstants.CURRENT_ROUND)}, "
             f"{self.lr_scheduler_name if self.lr_scheduler_name else ''} "
             f"lr: {self.optimizer.param_groups[-1]['lr']}, "
+            f"fedopt layers: {len(updated_params)}, "
+            f"fedavg layers: {n_fedavg}, "
             f"update: {secs} secs., detach: {secs_detach} secs.",
         )
         # TODO: write server-side lr to tensorboard

{"golden_diff": "diff --git a/nvflare/app_opt/pt/fedopt.py b/nvflare/app_opt/pt/fedopt.py\n--- a/nvflare/app_opt/pt/fedopt.py\n+++ b/nvflare/app_opt/pt/fedopt.py\n@@ -21,7 +21,7 @@\n from nvflare.apis.fl_context import FLContext\n from nvflare.apis.shareable import Shareable\n from nvflare.app_common.abstract.learnable import Learnable\n-from nvflare.app_common.abstract.model import make_model_learnable\n+from nvflare.app_common.abstract.model import ModelLearnableKey, make_model_learnable\n from nvflare.app_common.app_constant import AppConstants\n from nvflare.app_common.shareablegenerators.full_model_shareable_generator import FullModelShareableGenerator\n from nvflare.security.logging import secure_format_exception\n@@ -40,6 +40,8 @@\n         The algorithm is proposed in Reddi, Sashank, et al. \"Adaptive federated optimization.\" arXiv preprint arXiv:2003.00295 (2020).\n         This SharableGenerator will update the global model using the specified\n         PyTorch optimizer and learning rate scheduler.\n+        Note: This class will use FedOpt to optimize the global trainable parameters (i.e. `self.model.named_parameters()`)\n+        but use FedAvg to update any other layers such as batch norm statistics.\n \n         Args:\n             optimizer_args: dictionary of optimizer arguments, e.g.\n@@ -164,14 +166,17 @@\n \n         # Apply the update to the model. We must multiply weights_delta by -1.0 to\n         # view it as a gradient that should be applied to the server_optimizer.\n+        updated_params = []\n         for name, param in self.model.named_parameters():\n-            param.grad = torch.tensor(-1.0 * model_diff[name]).to(self.device)\n+            if name in model_diff:\n+                param.grad = torch.tensor(-1.0 * model_diff[name]).to(self.device)\n+                updated_params.append(name)\n \n         self.optimizer.step()\n         if self.lr_scheduler is not None:\n             self.lr_scheduler.step()\n \n-        return self.model.state_dict()\n+        return self.model.state_dict(), updated_params\n \n     def shareable_to_learnable(self, shareable: Shareable, fl_ctx: FLContext) -> Learnable:\n         \"\"\"Convert Shareable to Learnable while doing a FedOpt update step.\n@@ -206,7 +211,7 @@\n         model_diff = dxo.data\n \n         start = time.time()\n-        weights = self.server_update(model_diff)\n+        weights, updated_params = self.server_update(model_diff)\n         secs = time.time() - start\n \n         # convert to numpy dict of weights\n@@ -215,12 +220,28 @@\n             weights[key] = weights[key].detach().cpu().numpy()\n         secs_detach = time.time() - start\n \n+        # update unnamed parameters such as batch norm layers if there are any using the averaged update\n+        base_model = fl_ctx.get_prop(AppConstants.GLOBAL_MODEL)\n+        if not base_model:\n+            self.system_panic(reason=\"No global base model!\", fl_ctx=fl_ctx)\n+            return base_model\n+\n+        base_model_weights = base_model[ModelLearnableKey.WEIGHTS]\n+\n+        n_fedavg = 0\n+        for key, value in model_diff.items():\n+            if key not in updated_params:\n+                weights[key] = base_model_weights[key] + value\n+                n_fedavg += 1\n+\n         self.log_info(\n             fl_ctx,\n             f\"FedOpt ({self.optimizer_name}, {self.device}) server model update \"\n             f\"round {fl_ctx.get_prop(AppConstants.CURRENT_ROUND)}, \"\n             f\"{self.lr_scheduler_name if self.lr_scheduler_name else ''} \"\n             f\"lr: {self.optimizer.param_groups[-1]['lr']}, \"\n+            f\"fedopt layers: {len(updated_params)}, \"\n+            f\"fedavg layers: {n_fedavg}, \"\n             f\"update: {secs} secs., detach: {secs_detach} secs.\",\n         )\n         # TODO: write server-side lr to tensorboard\n", "issue": "[BUG] FedOpt algorithm not working as expected in cifar10 example\n**Describe the bug**\r\nThe FedOpt algorithm is not working as expected in cifar10 example when I change the model from the pre-existing ModerateCNN to another model like MobileNetv2 or Resnet18 and others. The problem is that the accuracy of the global model is not increasing or increasing too slow with the FedOpt algorithm while the other algorithms works just fine even changing the model.\r\n\r\n**To Reproduce**\r\n1. Add in 'cifar10_nets.py' the new model : \r\n   class MyModel(nn.Module):\r\n     def __init__(self):\r\n        super(MyModel, self).__init__()\r\n        model = models.mobilenet_v2(weights='DEFAULT')\r\n        model.classifier = nn.Sequential(\r\n            nn.Dropout(0.4),\r\n            nn.Linear(1280, 10),\r\n        )\r\n        self.model = model\r\n    \r\n    def forward(self, x):\r\n        return self.model(x)\r\n2. Import and change the model in file 'cifar10_learner.py'\r\n3. Launch the example with `./run_simulator.sh cifar10_fedopt 0.1 8 8`\r\n4. See the results in tensorboard with `tensorboard --logdir=/tmp/nvflare/sim_cifar10` under the section 'val_acc_global_model'\r\n\r\n**Expected behavior**\r\nI expect reading the algorithm proposed in Reddi, Sashank, et al. \"Adaptive federated optimization.\" arXiv preprint arXiv:2003.00295 (2020), to obtain the same performance of FedAvg using SGD optimizer with lr = 1.0 and no scheduler. Also obtain better results changing optimizer and adding a scheduler.\r\n\r\n**Screenshots**\r\n![Screenshot from 2023-04-28 10-09-15](https://user-images.githubusercontent.com/102248925/235092305-891cc9b1-e0d6-48dd-90f1-48898d6b918b.png)\r\nPurple = FedAvg\r\nPink = FedOpt\r\n\r\n**Desktop (please complete the following information):**\r\n - OS: ubuntu 22.04\r\n - Python Version 3.10\r\n - NVFlare Version 2.3.0\r\n\r\nTy in advance!\r\n\n", "before_files": [{"content": "# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport time\n\nimport torch\n\nfrom nvflare.apis.dxo import DataKind, MetaKey, from_shareable\nfrom nvflare.apis.event_type import EventType\nfrom nvflare.apis.fl_context import FLContext\nfrom nvflare.apis.shareable import Shareable\nfrom nvflare.app_common.abstract.learnable import Learnable\nfrom nvflare.app_common.abstract.model import make_model_learnable\nfrom nvflare.app_common.app_constant import AppConstants\nfrom nvflare.app_common.shareablegenerators.full_model_shareable_generator import FullModelShareableGenerator\nfrom nvflare.security.logging import secure_format_exception\n\n\nclass PTFedOptModelShareableGenerator(FullModelShareableGenerator):\n    def __init__(\n        self,\n        optimizer_args: dict = None,\n        lr_scheduler_args: dict = None,\n        source_model=\"model\",\n        device=None,\n    ):\n        \"\"\"Implement the FedOpt algorithm.\n\n        The algorithm is proposed in Reddi, Sashank, et al. \"Adaptive federated optimization.\" arXiv preprint arXiv:2003.00295 (2020).\n        This SharableGenerator will update the global model using the specified\n        PyTorch optimizer and learning rate scheduler.\n\n        Args:\n            optimizer_args: dictionary of optimizer arguments, e.g.\n                {'path': 'torch.optim.SGD', 'args': {'lr': 1.0}} (default).\n            lr_scheduler_args: dictionary of server-side learning rate scheduler arguments, e.g.\n                {'path': 'torch.optim.lr_scheduler.CosineAnnealingLR', 'args': {'T_max': 100}} (default: None).\n            source_model: either a valid torch model object or a component ID of a torch model object\n            device: specify the device to run server-side optimization, e.g. \"cpu\" or \"cuda:0\"\n                (will default to cuda if available and no device is specified).\n\n        Raises:\n            TypeError: when any of input arguments does not have correct type\n        \"\"\"\n        super().__init__()\n        if not optimizer_args:\n            self.logger(\"No optimizer_args provided. Using FedOpt with SGD and lr 1.0\")\n            optimizer_args = {\"name\": \"SGD\", \"args\": {\"lr\": 1.0}}\n\n        if not isinstance(optimizer_args, dict):\n            raise TypeError(\n                \"optimizer_args must be a dict of format, e.g. {'path': 'torch.optim.SGD', 'args': {'lr': 1.0}}.\"\n            )\n        if lr_scheduler_args is not None:\n            if not isinstance(lr_scheduler_args, dict):\n                raise TypeError(\n                    \"optimizer_args must be a dict of format, e.g. \"\n                    \"{'path': 'torch.optim.lr_scheduler.CosineAnnealingLR', 'args': {'T_max': 100}}.\"\n                )\n        self.source_model = source_model\n        self.optimizer_args = optimizer_args\n        self.lr_scheduler_args = lr_scheduler_args\n        self.model = None\n        self.optimizer = None\n        self.lr_scheduler = None\n        if device is None:\n            self.device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n        else:\n            self.device = torch.device(device)\n        self.optimizer_name = None\n        self.lr_scheduler_name = None\n\n    def _get_component_name(self, component_args):\n        if component_args is not None:\n            name = component_args.get(\"path\", None)\n            if name is None:\n                name = component_args.get(\"name\", None)\n            return name\n        else:\n            return None\n\n    def handle_event(self, event_type: str, fl_ctx: FLContext):\n        if event_type == EventType.START_RUN:\n            # Initialize the optimizer with current global model params\n            engine = fl_ctx.get_engine()\n\n            if isinstance(self.source_model, str):\n                self.model = engine.get_component(self.source_model)\n            else:\n                self.model = self.source_model\n\n            if self.model is None:\n                self.system_panic(\n                    \"Model is not available\",\n                    fl_ctx,\n                )\n                return\n            elif not isinstance(self.model, torch.nn.Module):\n                self.system_panic(\n                    f\"Expected model to be a torch.nn.Module but got {type(self.model)}\",\n                    fl_ctx,\n                )\n                return\n            else:\n                print(\"server model\", self.model)\n\n            self.model.to(self.device)\n\n            # set up optimizer\n            try:\n                # use provided or default optimizer arguments and add the model parameters\n                if \"args\" not in self.optimizer_args:\n                    self.optimizer_args[\"args\"] = {}\n                self.optimizer_args[\"args\"][\"params\"] = self.model.parameters()\n                self.optimizer = engine.build_component(self.optimizer_args)\n                # get optimizer name for log\n                self.optimizer_name = self._get_component_name(self.optimizer_args)\n            except Exception as e:\n                self.system_panic(\n                    f\"Exception while parsing `optimizer_args`({self.optimizer_args}): {secure_format_exception(e)}\",\n                    fl_ctx,\n                )\n                return\n\n            # set up lr scheduler\n            if self.lr_scheduler_args is not None:\n                try:\n                    self.lr_scheduler_name = self._get_component_name(self.lr_scheduler_args)\n                    # use provided or default lr scheduler argument and add the optimizer\n                    if \"args\" not in self.lr_scheduler_args:\n                        self.lr_scheduler_args[\"args\"] = {}\n                    self.lr_scheduler_args[\"args\"][\"optimizer\"] = self.optimizer\n                    self.lr_scheduler = engine.build_component(self.lr_scheduler_args)\n                except Exception as e:\n                    self.system_panic(\n                        f\"Exception while parsing `lr_scheduler_args`({self.lr_scheduler_args}): {secure_format_exception(e)}\",\n                        fl_ctx,\n                    )\n                    return\n\n    def server_update(self, model_diff):\n        \"\"\"Updates the global model using the specified optimizer.\n\n        Args:\n            model_diff: the aggregated model differences from clients.\n\n        Returns:\n            The updated PyTorch model state dictionary.\n\n        \"\"\"\n        self.model.train()\n        self.optimizer.zero_grad()\n\n        # Apply the update to the model. We must multiply weights_delta by -1.0 to\n        # view it as a gradient that should be applied to the server_optimizer.\n        for name, param in self.model.named_parameters():\n            param.grad = torch.tensor(-1.0 * model_diff[name]).to(self.device)\n\n        self.optimizer.step()\n        if self.lr_scheduler is not None:\n            self.lr_scheduler.step()\n\n        return self.model.state_dict()\n\n    def shareable_to_learnable(self, shareable: Shareable, fl_ctx: FLContext) -> Learnable:\n        \"\"\"Convert Shareable to Learnable while doing a FedOpt update step.\n\n        Supporting data_kind == DataKind.WEIGHT_DIFF\n\n        Args:\n            shareable (Shareable): Shareable to be converted\n            fl_ctx (FLContext): FL context\n\n        Returns:\n            Model: Updated global ModelLearnable.\n        \"\"\"\n        # check types\n        dxo = from_shareable(shareable)\n\n        if dxo.data_kind != DataKind.WEIGHT_DIFF:\n            self.system_panic(\n                \"FedOpt is only implemented for \" \"data_kind == DataKind.WEIGHT_DIFF\",\n                fl_ctx,\n            )\n            return Learnable()\n\n        processed_algorithm = dxo.get_meta_prop(MetaKey.PROCESSED_ALGORITHM)\n        if processed_algorithm is not None:\n            self.system_panic(\n                f\"FedOpt is not implemented for shareable processed by {processed_algorithm}\",\n                fl_ctx,\n            )\n            return Learnable()\n\n        model_diff = dxo.data\n\n        start = time.time()\n        weights = self.server_update(model_diff)\n        secs = time.time() - start\n\n        # convert to numpy dict of weights\n        start = time.time()\n        for key in weights:\n            weights[key] = weights[key].detach().cpu().numpy()\n        secs_detach = time.time() - start\n\n        self.log_info(\n            fl_ctx,\n            f\"FedOpt ({self.optimizer_name}, {self.device}) server model update \"\n            f\"round {fl_ctx.get_prop(AppConstants.CURRENT_ROUND)}, \"\n            f\"{self.lr_scheduler_name if self.lr_scheduler_name else ''} \"\n            f\"lr: {self.optimizer.param_groups[-1]['lr']}, \"\n            f\"update: {secs} secs., detach: {secs_detach} secs.\",\n        )\n        # TODO: write server-side lr to tensorboard\n\n        return make_model_learnable(weights, dxo.get_meta_props())\n", "path": "nvflare/app_opt/pt/fedopt.py"}], "after_files": [{"content": "# Copyright (c) 2023, NVIDIA CORPORATION.  All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport time\n\nimport torch\n\nfrom nvflare.apis.dxo import DataKind, MetaKey, from_shareable\nfrom nvflare.apis.event_type import EventType\nfrom nvflare.apis.fl_context import FLContext\nfrom nvflare.apis.shareable import Shareable\nfrom nvflare.app_common.abstract.learnable import Learnable\nfrom nvflare.app_common.abstract.model import ModelLearnableKey, make_model_learnable\nfrom nvflare.app_common.app_constant import AppConstants\nfrom nvflare.app_common.shareablegenerators.full_model_shareable_generator import FullModelShareableGenerator\nfrom nvflare.security.logging import secure_format_exception\n\n\nclass PTFedOptModelShareableGenerator(FullModelShareableGenerator):\n    def __init__(\n        self,\n        optimizer_args: dict = None,\n        lr_scheduler_args: dict = None,\n        source_model=\"model\",\n        device=None,\n    ):\n        \"\"\"Implement the FedOpt algorithm.\n\n        The algorithm is proposed in Reddi, Sashank, et al. \"Adaptive federated optimization.\" arXiv preprint arXiv:2003.00295 (2020).\n        This SharableGenerator will update the global model using the specified\n        PyTorch optimizer and learning rate scheduler.\n        Note: This class will use FedOpt to optimize the global trainable parameters (i.e. `self.model.named_parameters()`)\n        but use FedAvg to update any other layers such as batch norm statistics.\n\n        Args:\n            optimizer_args: dictionary of optimizer arguments, e.g.\n                {'path': 'torch.optim.SGD', 'args': {'lr': 1.0}} (default).\n            lr_scheduler_args: dictionary of server-side learning rate scheduler arguments, e.g.\n                {'path': 'torch.optim.lr_scheduler.CosineAnnealingLR', 'args': {'T_max': 100}} (default: None).\n            source_model: either a valid torch model object or a component ID of a torch model object\n            device: specify the device to run server-side optimization, e.g. \"cpu\" or \"cuda:0\"\n                (will default to cuda if available and no device is specified).\n\n        Raises:\n            TypeError: when any of input arguments does not have correct type\n        \"\"\"\n        super().__init__()\n        if not optimizer_args:\n            self.logger(\"No optimizer_args provided. Using FedOpt with SGD and lr 1.0\")\n            optimizer_args = {\"name\": \"SGD\", \"args\": {\"lr\": 1.0}}\n\n        if not isinstance(optimizer_args, dict):\n            raise TypeError(\n                \"optimizer_args must be a dict of format, e.g. {'path': 'torch.optim.SGD', 'args': {'lr': 1.0}}.\"\n            )\n        if lr_scheduler_args is not None:\n            if not isinstance(lr_scheduler_args, dict):\n                raise TypeError(\n                    \"optimizer_args must be a dict of format, e.g. \"\n                    \"{'path': 'torch.optim.lr_scheduler.CosineAnnealingLR', 'args': {'T_max': 100}}.\"\n                )\n        self.source_model = source_model\n        self.optimizer_args = optimizer_args\n        self.lr_scheduler_args = lr_scheduler_args\n        self.model = None\n        self.optimizer = None\n        self.lr_scheduler = None\n        if device is None:\n            self.device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n        else:\n            self.device = torch.device(device)\n        self.optimizer_name = None\n        self.lr_scheduler_name = None\n\n    def _get_component_name(self, component_args):\n        if component_args is not None:\n            name = component_args.get(\"path\", None)\n            if name is None:\n                name = component_args.get(\"name\", None)\n            return name\n        else:\n            return None\n\n    def handle_event(self, event_type: str, fl_ctx: FLContext):\n        if event_type == EventType.START_RUN:\n            # Initialize the optimizer with current global model params\n            engine = fl_ctx.get_engine()\n\n            if isinstance(self.source_model, str):\n                self.model = engine.get_component(self.source_model)\n            else:\n                self.model = self.source_model\n\n            if self.model is None:\n                self.system_panic(\n                    \"Model is not available\",\n                    fl_ctx,\n                )\n                return\n            elif not isinstance(self.model, torch.nn.Module):\n                self.system_panic(\n                    f\"Expected model to be a torch.nn.Module but got {type(self.model)}\",\n                    fl_ctx,\n                )\n                return\n            else:\n                print(\"server model\", self.model)\n\n            self.model.to(self.device)\n\n            # set up optimizer\n            try:\n                # use provided or default optimizer arguments and add the model parameters\n                if \"args\" not in self.optimizer_args:\n                    self.optimizer_args[\"args\"] = {}\n                self.optimizer_args[\"args\"][\"params\"] = self.model.parameters()\n                self.optimizer = engine.build_component(self.optimizer_args)\n                # get optimizer name for log\n                self.optimizer_name = self._get_component_name(self.optimizer_args)\n            except Exception as e:\n                self.system_panic(\n                    f\"Exception while parsing `optimizer_args`({self.optimizer_args}): {secure_format_exception(e)}\",\n                    fl_ctx,\n                )\n                return\n\n            # set up lr scheduler\n            if self.lr_scheduler_args is not None:\n                try:\n                    self.lr_scheduler_name = self._get_component_name(self.lr_scheduler_args)\n                    # use provided or default lr scheduler argument and add the optimizer\n                    if \"args\" not in self.lr_scheduler_args:\n                        self.lr_scheduler_args[\"args\"] = {}\n                    self.lr_scheduler_args[\"args\"][\"optimizer\"] = self.optimizer\n                    self.lr_scheduler = engine.build_component(self.lr_scheduler_args)\n                except Exception as e:\n                    self.system_panic(\n                        f\"Exception while parsing `lr_scheduler_args`({self.lr_scheduler_args}): {secure_format_exception(e)}\",\n                        fl_ctx,\n                    )\n                    return\n\n    def server_update(self, model_diff):\n        \"\"\"Updates the global model using the specified optimizer.\n\n        Args:\n            model_diff: the aggregated model differences from clients.\n\n        Returns:\n            The updated PyTorch model state dictionary.\n\n        \"\"\"\n        self.model.train()\n        self.optimizer.zero_grad()\n\n        # Apply the update to the model. We must multiply weights_delta by -1.0 to\n        # view it as a gradient that should be applied to the server_optimizer.\n        updated_params = []\n        for name, param in self.model.named_parameters():\n            if name in model_diff:\n                param.grad = torch.tensor(-1.0 * model_diff[name]).to(self.device)\n                updated_params.append(name)\n\n        self.optimizer.step()\n        if self.lr_scheduler is not None:\n            self.lr_scheduler.step()\n\n        return self.model.state_dict(), updated_params\n\n    def shareable_to_learnable(self, shareable: Shareable, fl_ctx: FLContext) -> Learnable:\n        \"\"\"Convert Shareable to Learnable while doing a FedOpt update step.\n\n        Supporting data_kind == DataKind.WEIGHT_DIFF\n\n        Args:\n            shareable (Shareable): Shareable to be converted\n            fl_ctx (FLContext): FL context\n\n        Returns:\n            Model: Updated global ModelLearnable.\n        \"\"\"\n        # check types\n        dxo = from_shareable(shareable)\n\n        if dxo.data_kind != DataKind.WEIGHT_DIFF:\n            self.system_panic(\n                \"FedOpt is only implemented for \" \"data_kind == DataKind.WEIGHT_DIFF\",\n                fl_ctx,\n            )\n            return Learnable()\n\n        processed_algorithm = dxo.get_meta_prop(MetaKey.PROCESSED_ALGORITHM)\n        if processed_algorithm is not None:\n            self.system_panic(\n                f\"FedOpt is not implemented for shareable processed by {processed_algorithm}\",\n                fl_ctx,\n            )\n            return Learnable()\n\n        model_diff = dxo.data\n\n        start = time.time()\n        weights, updated_params = self.server_update(model_diff)\n        secs = time.time() - start\n\n        # convert to numpy dict of weights\n        start = time.time()\n        for key in weights:\n            weights[key] = weights[key].detach().cpu().numpy()\n        secs_detach = time.time() - start\n\n        # update unnamed parameters such as batch norm layers if there are any using the averaged update\n        base_model = fl_ctx.get_prop(AppConstants.GLOBAL_MODEL)\n        if not base_model:\n            self.system_panic(reason=\"No global base model!\", fl_ctx=fl_ctx)\n            return base_model\n\n        base_model_weights = base_model[ModelLearnableKey.WEIGHTS]\n\n        n_fedavg = 0\n        for key, value in model_diff.items():\n            if key not in updated_params:\n                weights[key] = base_model_weights[key] + value\n                n_fedavg += 1\n\n        self.log_info(\n            fl_ctx,\n            f\"FedOpt ({self.optimizer_name}, {self.device}) server model update \"\n            f\"round {fl_ctx.get_prop(AppConstants.CURRENT_ROUND)}, \"\n            f\"{self.lr_scheduler_name if self.lr_scheduler_name else ''} \"\n            f\"lr: {self.optimizer.param_groups[-1]['lr']}, \"\n            f\"fedopt layers: {len(updated_params)}, \"\n            f\"fedavg layers: {n_fedavg}, \"\n            f\"update: {secs} secs., detach: {secs_detach} secs.\",\n        )\n        # TODO: write server-side lr to tensorboard\n\n        return make_model_learnable(weights, dxo.get_meta_props())\n", "path": "nvflare/app_opt/pt/fedopt.py"}]}