Switched to numel method to obtain the number of parameters. (#380)

examples/secure_aggregation/maskcrypt/maskcrypt_CIFAR10_resnet18.yml

@@ -0,0 +1,66 @@
+clients:
+  # Type
+  type: simple_he
+
+  # The total number of clients
+  total_clients: 10
+
+  # encrypt ratio
+  encrypt_ratio: 0.1
+  random_mask: False
+
+  # The number of clients selected in each round
+  per_round: 5
+
+  # Should the clients compute test accuracy locally?
+  do_test: false
+
+server:
+  address: 127.0.0.1
+  port: 8001
+  random_seed: 1
+  simulate_wall_time: true
+
+data:
+  # The training and testing dataset
+  datasource: CIFAR10
+
+  # Number of samples in each partition
+  partition_size: 1000
+
+  # IID or non-IID?
+  sampler: iid
+
+trainer:
+  # The maximum number of training rounds
+  rounds: 25
+
+  # The maximum number of clients running concurrently
+  max_concurrency: 5
+
+  # The target accuracy
+  target_accuracy: 0.99
+
+  # The machine learning model
+  model_name: resnet_18
+
+  # Number of epoches for local training in each communication round
+  epochs: 5
+  batch_size: 32
+  optimizer: SGD
+algorithm:
+  # Aggregation algorithm
+  type: fedavg
+
+results:
+  # Write the following parameter(s) into a CSV
+  types: round, elapsed_time, accuracy, comm_overhead
+
+parameters:
+  model:
+    num_classes: 10
+
+  optimizer:
+      lr: 0.01
+      momentum: 0.9
+      weight_decay: 0.0

plato/servers/fedavg_he.py

@@ -1,6 +1,7 @@
 """
 A federated learning server using federated averaging to aggregate updates after homomorphic encryption.
 """
+
 from functools import reduce
 from plato.servers import fedavg
 from plato.utils import homo_enc
@@ -36,7 +37,7 @@ def configure(self) -> None:
 
         for key in extract_model.keys():
             self.weight_shapes[key] = extract_model[key].size()
-            self.para_nums[key] = reduce(lambda a, b: a * b, self.weight_shapes[key])
+            self.para_nums[key] = extract_model[key].numel()
 
         self.encrypted_model = homo_enc.encrypt_weights(
             extract_model, True, self.context, []