Merge branch 'SymbioticLab:master' into master

.gitignore

@@ -36,4 +36,10 @@ landmark
 coqa
 .idea
 out/
-build/
+build/
+.vscode
+fedscale/cache
+fedscale/tmp*
+fedscale/*.tflite
+fedscale/*.ckpt
+fedscale/cloud/aggregation/cache

benchmark/configs/android/mnn.yml

@@ -0,0 +1,32 @@
+# Configuration file of running tensorflow backend
+
+# ========== Cluster configuration ==========
+# ip address of the parameter server (need 1 GPU process)
+ps_ip: localhost
+
+exp_path: $FEDSCALE_HOME/fedscale/cloud
+
+aggregator_entry: aggregation/aggregator_mnn.py
+
+auth:
+    ssh_user: ""
+    ssh_private_key: ~/.ssh/id_rsa
+
+# cmd to run before we can indeed run FAR (in order)
+setup_commands:
+    - source $HOME/anaconda3/bin/activate fedscale
+
+# ========== Additional job configuration ==========
+# Default parameters are specified in config_parser.py, wherein more description of the parameter can be found
+
+job_conf:
+    - job_name: android-mnn              # Generate logs under this folder: log_path/job_name/time_stamp
+    - log_path: $FEDSCALE_HOME/benchmark # Path of log files
+    - experiment_mode: mobile
+    - num_participants: 1                # Number of participants per round, we use K=100 in our paper, large K will be much slower
+    - model: linear                      # Need to define the model in aggregator_mnn.py
+    - learning_rate: 0.01
+    - batch_size: 32
+    - input_shape: 32 32 3
+    - num_classes: 10
+    - test_bsz: 16

benchmark/configs/android/tflite.yml

@@ -0,0 +1,33 @@
+# Configuration file of running tensorflow backend
+
+# ========== Cluster configuration ==========
+# ip address of the parameter server (need 1 GPU process)
+ps_ip: localhost
+
+exp_path: $FEDSCALE_HOME/fedscale/cloud
+
+aggregator_entry: aggregation/aggregator_tflite.py
+
+auth:
+    ssh_user: ""
+    ssh_private_key: ~/.ssh/id_rsa
+
+# cmd to run before we can indeed run FAR (in order)
+setup_commands:
+    - source $HOME/anaconda3/bin/activate fedscale
+
+# ========== Additional job configuration ==========
+# Default parameters are specified in config_parser.py, wherein more description of the parameter can be found
+
+job_conf:
+    - job_name: android-tflite           # Generate logs under this folder: log_path/job_name/time_stamp
+    - log_path: $FEDSCALE_HOME/benchmark # Path of log files
+    - experiment_mode: mobile
+    - num_participants: 1                # Number of participants per round, we use K=100 in our paper, large K will be much slower
+    - model: linear                      # Need to define the model in tf_aggregator.py
+    - learning_rate: 0.01
+    - batch_size: 32
+    - input_shape: 32 32 3
+    - num_classes: 10
+    - test_bsz: 16
+    - engine: 'tensorflow'

fedscale/cloud/aggregation/aggregator_mnn.py

@@ -1,9 +1,6 @@
-import json
-
 import fedscale.cloud.config_parser as parser
 from fedscale.cloud.aggregation.aggregator import Aggregator
-from fedscale.utils.models.simple.linear_model import LinearModel
-from fedscale.utils.models.mnn_convert import *
+from fedscale.utils.models.mnn_model_provider import *
 from fedscale.cloud.internal.torch_model_adapter import TorchModelAdapter
 
 
@@ -14,55 +11,58 @@ class MNNAggregator(Aggregator):
         args (dictionary): Variable arguments for fedscale runtime config. 
                            Defaults to the setup in arg_parser.py.
     """
+
     def __init__(self, args):
         super().__init__(args)
-        
+
         # == mnn model and keymap ==
-        self.mnn_json = None
+        self.mnn_model = None
         self.keymap_mnn2torch = {}
         self.input_shape = args.input_shape
-    
+
     def init_model(self):
         """
         Load the model architecture and convert to mnn.
         NOTE: MNN does not support dropout.
         """
-        if self.args.model == 'linear':
-            self.model_wrapper = TorchModelAdapter(LinearModel())
-            self.model_weights = self.model_wrapper.get_weights()
-        else:
-            super().init_model()
-        self.mnn_json = torch_to_mnn(self.model_wrapper.get_model(), self.input_shape, True)
-        self.keymap_mnn2torch = init_keymap(self.model_wrapper.get_model().state_dict(), self.mnn_json)
+        self.model_wrapper = TorchModelAdapter(
+            get_mnn_model(self.args.model, self.args))
+        self.model_weights = self.model_wrapper.get_weights()
+        self.mnn_model = torch_to_mnn(
+            self.model_wrapper.get_model(), self.input_shape)
+        self.keymap_mnn2torch = init_keymap(
+            self.model_wrapper.get_model().state_dict())
 
     def update_weight_aggregation(self, update_weights):
         """
         Update model when the round completes.
         Then convert new model to mnn json.
-        
+
         Args:
             last_model (list): A list of global model weight in last round.
         """
         super().update_weight_aggregation(update_weights)
         if self.model_in_update == self.tasks_round:
-            self.mnn_json = torch_to_mnn(self.model_wrapper.get_model(), self.input_shape)
+            self.mnn_model = torch_to_mnn(
+                self.model_wrapper.get_model(), self.input_shape)
 
     def deserialize_response(self, responses):
         """
         Deserialize the response from executor.
         If the response contains mnn json model, convert to pytorch state_dict.
-        
+
         Args:
             responses (byte stream): Serialized response from executor.
 
         Returns:
             string, bool, or bytes: The deserialized response object from executor.
         """
-        data = json.loads(responses.decode('utf-8'))
+        data = super().deserialize_response(responses)
         if "update_weight" in data:
             data["update_weight"] = mnn_to_torch(
                 self.keymap_mnn2torch,
-                json.loads(data["update_weight"]))
+                data["update_weight"],
+                data["client_id"])
         return data
 
     def serialize_response(self, responses):
@@ -76,10 +76,9 @@ def serialize_response(self, responses):
         Returns:
             bytes: The serialized response object to server.
         """
-        if type(responses) is list and all([np.array_equal(a, b) for a, b in zip(responses, self.model_wrapper.get_weights())]):
-            responses = self.mnn_json
-        data = json.dumps(responses)
-        return data.encode('utf-8')
+        if type(responses) is list:
+            responses = self.mnn_model
+        return super().serialize_response(responses)
 
 
 if __name__ == "__main__":