remove calls for legacy optimizer

Signed-off-by: kta-intel <[email protected]>
securefederatedai · Jun 7, 2024 · 05014d2 · 05014d2
1 parent bea8729
commit 05014d2
Showing 1 changed file with 7 additions and 28 deletions.
diff --git a/openfl/federated/task/runner_tf.py b/openfl/federated/task/runner_tf.py
@@ -257,12 +257,7 @@ def _get_weights_names(obj, with_opt_vars):
             The weight name list
         """
         if with_opt_vars:
-            # When acquiring optimizer weights, check optimizer version.
-            # Current optimizer does not use 'weights' attributes
-            if 'legacy' in obj.__class__.__module__:
-                weight_names = [weight.name for weight in obj.weights]
-            else:
-                weight_names = [weight.name for weight in obj.variables]
+            weight_names = [weight.name for weight in obj.variables]
 
         weight_names = [weight.name for weight in obj.weights]
         return weight_names
@@ -287,14 +282,8 @@ def _get_weights_dict(obj, suffix='', with_opt_vars=False):
 
         weights_dict = {}
         if with_opt_vars:
-            # When acquiring optimizer weights, check optimizer version.
-            # Current optimizer does not use 'weights' or '.get_weights()' attributes
-            if 'legacy' in obj.__class__.__module__:
-                weight_names = [weight.name for weight in obj.weights]
-                weight_values = obj.get_weights()
-            else:
-                weight_names = [weight.name for weight in obj.variables]
-                weight_values = [weight.numpy() for weight in obj.variables]
+            weight_names = [weight.name for weight in obj.variables]
+            weight_values = [weight.numpy() for weight in obj.variables]
         else:
             weight_names = [weight.name for weight in obj.weights]
             weight_values = obj.get_weights()
@@ -319,12 +308,7 @@ def _set_weights_dict(obj, weights_dict, with_opt_vars=False):
         """
 
         if with_opt_vars:
-            # When acquiring optimizer weights, check optimizer version.
-            # Current optimizer does not use 'weights' attributes
-            if 'legacy' in obj.__class__.__module__:
-                weight_names = [weight.name for weight in obj.weights]
-            else:
-                weight_names = [weight.name for weight in obj.variables]
+            weight_names = [weight.name for weight in obj.variables]
         else:
             weight_names = [weight.name for weight in obj.weights]
 
@@ -383,15 +367,10 @@ def set_tensor_dict(self, tensor_dict, with_opt_vars):
             model_weights_dict = {
                 name: tensor_dict[name] for name in model_weight_names
             }
-            if 'legacy' in self.model.optimizer.__class__.__module__:
-                opt_weight_names = [
-                    weight.name for weight in self.model.optimizer.weights
-                ]
-            else:
-                opt_weight_names = [
-                    weight.name for weight in self.model.optimizer.variables
-                ]
 
+            opt_weight_names = [
+                weight.name for weight in self.model.optimizer.variables
+            ]
             opt_weights_dict = {
                 name: tensor_dict[name] for name in opt_weight_names
             }