Repository: TL-System/plato (language: python)

Issue:
[BUG] Asynchronous mode does not proceed in the initial round
**Describe the bug**
When using the asynchronous mode, the server will not be able to proceed in the initial round, waiting forever for clients that will never arrive.

**To Reproduce**

Use the following configuration file with the `./run -c` command:

```
clients:
    # Type
    type: simple

    # The total number of clients
    total_clients: 500

    # The number of clients selected in each round
    per_round: 50

    # Should the clients compute test accuracy locally?
    do_test: false

    # Whether client heterogeneity should be simulated
    speed_simulation: true

    # The distribution of client speeds
    simulation_distribution:
        distribution: pareto
        alpha: 1

    # The maximum amount of time for clients to sleep after each epoch
    max_sleep_time: 30

    # Should clients really go to sleep, or should we just simulate the sleep times?
    sleep_simulation: false

    # If we are simulating client training times, what is the average training time?
    avg_training_time: 20

    random_seed: 1

server:
    address: 127.0.0.1
    port: 8000

    ping_timeout: 36000
    ping_interval: 36000

    # Should we operate in sychronous mode?
    synchronous: false

    # Should we simulate the wall-clock time on the server? Useful if max_concurrency is specified
    simulate_wall_time: true

    # What is the minimum number of clients that need to report before aggregation begins?
    minimum_clients_aggregated: 15

    # What is the staleness bound, beyond which the server should wait for stale clients?
    staleness_bound: 10

    # Should we send urgent notifications to stale clients beyond the staleness bound?
    request_update: false

    random_seed: 1

data:
    # The training and testing dataset
    datasource: MNIST 

    # Number of samples in each partition
    partition_size: 600 

    # IID or non-IID?
    sampler: noniid

    # The concentration parameter for the Dirichlet distribution
    concentration: 0.5

    # The random seed for sampling data
    random_seed: 1

trainer:
    # The type of the trainer
    type: basic 

    # The maximum number of training rounds
    rounds: 5

    # The maximum number of clients running concurrently
    max_concurrency: 3

    # Number of epoches for local training in each communication round
    epochs: 1
    batch_size: 32
    optimizer: SGD
    learning_rate: 0.01
    momentum: 0.9
    weight_decay: 0.0

    # The machine learning model
    model_name: lenet5 

algorithm:
    # Aggregation algorithm
    type: fedavg
```



Discussion / hints:

Fix patch:
diff --git a/README.md b/README.md
index 00f96b618..e7c84c982 100644
--- a/README.md
+++ b/README.md
@@ -248,7 +248,7 @@ The Plato federated learning server is designed to use Socket.IO over HTTP and H
 Remove the `conda` environment used to run *Plato* first, and then remove the directory containing *Plato*'s git repository.
 
 ```shell
-conda-env remove -n plato
+conda env remove -n plato
 rm -rf plato/
 ```
 
diff --git a/configs/CIFAR10/fedavg_resnet18_dp.yml b/configs/CIFAR10/fedavg_resnet18_dp.yml
index aec93a2ab..c9e4c43af 100644
--- a/configs/CIFAR10/fedavg_resnet18_dp.yml
+++ b/configs/CIFAR10/fedavg_resnet18_dp.yml
@@ -3,10 +3,10 @@ clients:
     type: simple
 
     # The total number of clients
-    total_clients: 1000
+    total_clients: 1
 
     # The number of clients selected in each round
-    per_round: 200
+    per_round: 1
 
     # Should the clients compute test accuracy locally?
     do_test: false
@@ -23,7 +23,7 @@ data:
     datasource: CIFAR10
 
     # Number of samples in each partition
-    partition_size: 1000
+    partition_size: 50000
 
     # IID or non-IID?
     sampler: iid
@@ -31,12 +31,12 @@ data:
 trainer:
     # Using the differential privacy trainer
     type: diff_privacy 
-    dp_epsilon: 1
+    dp_epsilon: 10
     dp_delta: 0.00001
     dp_max_grad_norm: 1 
 
     # The maximum number of training rounds
-    rounds: 20
+    rounds: 100
 
     # The maximum number of clients running concurrently
     max_concurrency: 1
@@ -45,12 +45,12 @@ trainer:
     target_accuracy: 0.80
 
     # Number of epoches for local training in each communication round
-    epochs: 5
-    batch_size: 10
+    epochs: 1
+    batch_size: 32
     optimizer: SGD
-    learning_rate: 0.1
+    learning_rate: 0.01
     momentum: 0.9
-    weight_decay: 0.0001
+    weight_decay: 0.0
 
     # Learning rate schedule
     lr_schedule: LambdaLR
diff --git a/configs/CelebA/fedavg_gan_dcgan.yml b/configs/CelebA/fedavg_gan_dcgan.yml
new file mode 100644
index 000000000..416604bf0
--- /dev/null
+++ b/configs/CelebA/fedavg_gan_dcgan.yml
@@ -0,0 +1,75 @@
+clients:
+    # Type
+    type: simple
+
+    # The total number of clients
+    total_clients: 2
+
+    # The number of clients selected in each round
+    per_round: 2
+
+    # Should the clients compute test accuracy locally?
+    do_test: false
+
+server:
+    address: 127.0.0.1
+    port: 8000
+
+    # Choose to send Generator or Discriminator to clients at
+    # the end of each round. 
+    # Value here should be one of 'none', 'generator', 'discriminator', or 'both'
+    network_to_sync: generator
+
+data:
+    # The training and testing dataset
+    datasource: CelebA
+
+    # Only add face identity as labels for training
+    celeba_targets:
+        attr: false
+        identity: true
+    
+    # Resize all images to 64x64
+    celeba_img_size: 64
+
+    # Where the dataset is located
+    data_path: data
+
+    # Number of samples in each partition
+    partition_size: 81000
+
+    # IID or non-IID?
+    sampler: iid
+
+    # The concentration parameter for the Dirichlet distribution
+    concentration: 0.5
+
+    # The random seed for sampling data
+    random_seed: 1
+
+trainer:
+    # The type of the trainer
+    type: gan
+
+    # The maximum number of training rounds
+    rounds: 5
+
+    # The maximum number of clients running concurrently
+    max_concurrency: 3
+
+    # The target Frechet Distance
+    target_perplexity: 0
+
+    # Number of epoches for local training in each communication round
+    epochs: 5
+    batch_size: 128
+    optimizer: Adam
+    learning_rate: 0.0002
+    weight_decay: 0.0
+
+    # The machine learning model
+    model_name: dcgan
+
+algorithm:
+    # Aggregation algorithm
+    type: fedavg_gan
diff --git a/configs/CelebA/fedavg_resnet18.yml b/configs/CelebA/fedavg_resnet18.yml
index 912e5553a..4c1231fd5 100644
--- a/configs/CelebA/fedavg_resnet18.yml
+++ b/configs/CelebA/fedavg_resnet18.yml
@@ -24,6 +24,9 @@ data:
         # For ResNet, do not set <attr> to True since it does not match the expected output of ResNet
         attr: false
         identity: true
+    
+    # Resize all images to 32x32; default is 64x64
+    celeba_img_size: 32
 
     # Number of identity in CelebA
     num_classes: 10178
diff --git a/configs/EMNIST/fedavg_lenet5_dp.yml b/configs/EMNIST/fedavg_lenet5_dp.yml
new file mode 100644
index 000000000..96a53b019
--- /dev/null
+++ b/configs/EMNIST/fedavg_lenet5_dp.yml
@@ -0,0 +1,78 @@
+clients:
+    # Type
+    type: simple
+
+    # The total number of clients
+    total_clients: 100
+
+    # The number of clients selected in each round
+    per_round: 5
+
+    # Should the clients compute test accuracy locally?
+    do_test: false
+
+    # Client simulation mode
+    simulation: true
+
+server:
+    address: 127.0.0.1
+    port: 8000
+    simulate_wall_time: true
+
+data:
+    # The training and testing dataset
+    datasource: EMNIST
+
+    # Where the dataset is located
+    data_path: ./data
+
+    # Number of samples in each partition
+    partition_size: 1128
+
+    # IID or non-IID?
+    sampler: iid
+
+    # The random seed for sampling data
+    random_seed: 1
+
+trainer:
+    # Using the differential privacy trainer
+    type: diff_privacy 
+    dp_epsilon: 10
+    dp_delta: 0.0001
+    dp_max_grad_norm: 1
+
+    # The maximum number of training rounds
+    rounds: 100
+
+    # Whether the training should use multiple GPUs if available
+    parallelized: false
+
+    # The maximum number of clients running concurrently
+    max_concurrency: 10
+
+    # The target accuracy
+    target_accuracy: 0.95
+
+    # Number of epoches for local training in each communication round
+    epochs: 5
+    batch_size: 32
+    optimizer: SGD
+    learning_rate: 0.01
+    momentum: 0.9
+    weight_decay: 0.0
+
+    # The machine learning model
+    model_name: lenet5
+    num_classes: 47
+
+algorithm:
+    # Aggregation algorithm
+    type: fedavg
+
+results:
+    # Write the following parameter(s) into a CSV
+    types: round, elapsed_time, accuracy
+
+    # Plot results (x_axis-y_axis)
+    plot: round-accuracy, elapsed_time-accuracy
diff --git a/configs/Models/kinetics_flow_audio_models.yml b/configs/Kinetics/Models/kinetics_flow_audio_models.yml
similarity index 100%
rename from configs/Models/kinetics_flow_audio_models.yml
rename to configs/Kinetics/Models/kinetics_flow_audio_models.yml
diff --git a/configs/Models/kinetics_full_models.yml b/configs/Kinetics/Models/kinetics_full_models.yml
similarity index 100%
rename from configs/Models/kinetics_full_models.yml
rename to configs/Kinetics/Models/kinetics_full_models.yml
diff --git a/configs/Pipeline/gym_audio_feature.yml b/configs/Kinetics/Pipeline/gym_audio_feature.yml
similarity index 100%
rename from configs/Pipeline/gym_audio_feature.yml
rename to configs/Kinetics/Pipeline/gym_audio_feature.yml
diff --git a/configs/Pipeline/gym_csn_rgb.yml b/configs/Kinetics/Pipeline/gym_csn_rgb.yml
similarity index 100%
rename from configs/Pipeline/gym_csn_rgb.yml
rename to configs/Kinetics/Pipeline/gym_csn_rgb.yml
diff --git a/configs/Pipeline/gym_tsn_flow.yml b/configs/Kinetics/Pipeline/gym_tsn_flow.yml
similarity index 100%
rename from configs/Pipeline/gym_tsn_flow.yml
rename to configs/Kinetics/Pipeline/gym_tsn_flow.yml
diff --git a/configs/Pipeline/kinetics_audio_feature.yml b/configs/Kinetics/Pipeline/kinetics_audio_feature.yml
similarity index 100%
rename from configs/Pipeline/kinetics_audio_feature.yml
rename to configs/Kinetics/Pipeline/kinetics_audio_feature.yml
diff --git a/configs/Pipeline/kinetics_csn_rgb.yml b/configs/Kinetics/Pipeline/kinetics_csn_rgb.yml
similarity index 100%
rename from configs/Pipeline/kinetics_csn_rgb.yml
rename to configs/Kinetics/Pipeline/kinetics_csn_rgb.yml
diff --git a/configs/Pipeline/kinetics_tsn_flow.yml b/configs/Kinetics/Pipeline/kinetics_tsn_flow.yml
similarity index 100%
rename from configs/Pipeline/kinetics_tsn_flow.yml
rename to configs/Kinetics/Pipeline/kinetics_tsn_flow.yml
diff --git a/configs/YOLO/mistnet_helmet.yml b/configs/YOLO/mistnet_helmet.yml
index a6b345b0d..cb25297df 100644
--- a/configs/YOLO/mistnet_helmet.yml
+++ b/configs/YOLO/mistnet_helmet.yml
@@ -37,7 +37,7 @@ data:
     test_path: data/helmet/helmet/images/val/
 
     # download command/URL
-    download_urls: ["https://iqua.ece.toronto.edu/~bli/helmet.zip"]
+    download_urls: ["http://iqua.ece.toronto.edu/~bli/helmet.zip"]
 
     # number of training examples
     num_train_examples: 5457
diff --git a/docs/Configuration.md b/docs/Configuration.md
index d94d69a32..c7adac27a 100644
--- a/docs/Configuration.md
+++ b/docs/Configuration.md
@@ -143,9 +143,9 @@ Attributes in **bold** must be included in a configuration file, while attribute
 |**epochs**|Number of epoches for local training in each communication round|Any positive integer||
 |**optimizer**||`SGD`, `Adam` or `FedProx`||
 |**batch_size**||Any positive integer||
-|**learning_rate**||||
+|**learning_rate**|||Decrease value when using `diff_privacy` trainer|
 |**momentum**||||
-|**weight_decay**||||   
+|**weight_decay**|||When using `diff_privacy` trainer, set to 0|   
 |lr_schedule|Learning rate scheduler|`CosineAnnealingLR`, `LambdaLR`, `StepLR`, `ReduceLROnPlateau`|| 
 |**model_name**|The machine learning model|`lenet5`, `resnet_x`, `vgg_x`,`wideresnet`, `feedback_transformer`, `yolov5`, `HuggingFace_CausalLM`, `inceptionv3`, `googlenet`, `unet`, `alexnet`, `squeezenet_x`, `shufflenet_x`|For `resnet_x`, x = 18, 34, 50, 101, or 152; For `vgg_x`, x = 11, 13, 16, or 19; For `squeezenet_x`, x = 0 or 1; For `shufflenet_x`, x = 0.5, 1.0, 1.5, or 2.0|
 |pretrained|Use a model pretrained on ImageNet or not|`true` or `false`. Default is `false`|Can be used for `inceptionv3`, `alexnet`, and `squeezenet_x` models.|
@@ -167,6 +167,6 @@ Attributes in **bold** must be included in a configuration file, while attribute
 
 | Attribute | Meaning | Valid Value | Note |
 |:---------:|:-------:|:-----------:|:----:|
-|types|The set of columns that will be written into a .csv file|`round`, `accuracy`, `elapsed_time`, `comm_time`, `round_time`, `local_epoch_num`, `edge_agg_num` (Use comma `,` to seperate them)|default: `round, accuracy, elapsed_time`|
+|types|The set of columns that will be written into a .csv file|`round`, `accuracy`, `elapsed_time`, `comm_time`, `round_time`, `comm_overhead`, `local_epoch_num`, `edge_agg_num` (Use comma `,` to seperate them)|default: `round, accuracy, elapsed_time`|
 |plot|Plot results |(Format: x\_axis-y\_axis. Use hyphen `-` to seperate axis. Use comma `,` to seperate multiple plots)|default: `round-accuracy, elapsed_time-accuracy`|
 |result_path|The directory of results||default: `<base_path>/results/`|
\ No newline at end of file
diff --git a/examples/FedRep/fedrep_trainer.py b/examples/FedRep/fedrep_trainer.py
index 9d2903b3f..63a237297 100644
--- a/examples/FedRep/fedrep_trainer.py
+++ b/examples/FedRep/fedrep_trainer.py
@@ -84,17 +84,12 @@ def train_model(self, config, trainset, sampler, cut_layer=None):
                                                        batch_size=batch_size,
                                                        sampler=sampler)
 
-        iterations_per_epoch = np.ceil(len(trainset) / batch_size).astype(int)
-        # load the total local update epochs
         epochs = config['epochs']
+
         # load the local update epochs for head optimization
         head_epochs = config[
             'head_epochs'] if 'head_epochs' in config else epochs - 1
 
-        # Sending the model to the device used for training
-        self.model.to(self.device)
-        self.model.train()
-
         # Initializing the loss criterion
         _loss_criterion = getattr(self, "loss_criterion", None)
         if callable(_loss_criterion):
@@ -108,13 +103,16 @@ def train_model(self, config, trainset, sampler, cut_layer=None):
         optimizer = get_optimizer(self.model)
 
         # Initializing the learning rate schedule, if necessary
-        if hasattr(config, 'lr_schedule'):
+        if 'lr_schedule' in config:
             lr_schedule = optimizers.get_lr_schedule(optimizer,
-                                                     iterations_per_epoch,
+                                                     len(train_loader),
                                                      train_loader)
         else:
             lr_schedule = None
 
+        self.model.to(self.device)
+        self.model.train()
+
         for epoch in range(1, epochs + 1):
 
             # As presented in the Section 3 of the FedRep paper,
diff --git a/examples/adaptive_freezing/adaptive_freezing.py b/examples/adaptive_freezing/adaptive_freezing.py
index 97bf45d96..b8aba3bdb 100644
--- a/examples/adaptive_freezing/adaptive_freezing.py
+++ b/examples/adaptive_freezing/adaptive_freezing.py
@@ -9,7 +9,7 @@
 Systems (ICDCS 2021), Online, July 7-10, 2021.
 
 The camera-ready manuscript of this paper is located at:
-https://iqua.ece.toronto.edu/papers/cchen-icdcs21.pdf
+https://ieeexplore.ieee.org/document/9546506
 
 """
 
diff --git a/examples/adaptive_hgb/adaptive_hgb.py b/examples/adaptive_hgb/adaptive_hgb.py
index 9730a7a4f..34dc50754 100644
--- a/examples/adaptive_hgb/adaptive_hgb.py
+++ b/examples/adaptive_hgb/adaptive_hgb.py
@@ -13,16 +13,16 @@
 import adaptive_hgb_server
 import adaptive_hgb_trainer
 
+os.environ['config_file'] = 'configs/Kinetics/kinetics_mm.yml'
+
 from plato.config import Config
 from plato.models.multimodal import multimodal_module
-from plato.datasources.multimodal import kinetics
-
-os.environ['config_file'] = 'configs/Kinetics/kinetics_mm.yml'
+from plato.datasources import kinetics
 
 
 def main():
     """ A Plato federated learning training session using the HGB algorithm. """
-    _ = Config()
+    __ = Config()
     support_modalities = ['rgb', "flow", "audio"]
 
     kinetics_data_source = kinetics.DataSource()
@@ -42,7 +42,7 @@ def main():
         multimodal_nets_configs=Config.multimodal_nets_configs,
         is_fused_head=True)
 
-    trainer = adaptive_hgb_trainer.Trainer(model=multi_model)
+    trainer = adaptive_hgb_trainer.Trainer
 
     client = adaptive_hgb_client.Client(model=multi_model,
                                         datasource=kinetics_data_source,
diff --git a/examples/adaptive_hgb/adaptive_hgb_trainer.py b/examples/adaptive_hgb/adaptive_hgb_trainer.py
index 35f51e730..faede49ae 100644
--- a/examples/adaptive_hgb/adaptive_hgb_trainer.py
+++ b/examples/adaptive_hgb/adaptive_hgb_trainer.py
@@ -235,24 +235,24 @@ def train_process(self, config, trainset, evalset, sampler,
                                                   num_workers=config.data.get(
                                                       'workers_per_gpu', 1))
 
-        iterations_per_epoch = np.ceil(len(trainset) / batch_size).astype(int)
         epochs = config['epochs']
 
-        # Sending the model to the device used for training
-        self.model.to(self.device)
-        self.model.train()
         # Initializing the optimizer
         get_optimizer = getattr(self, "get_optimizer",
                                 optimizers.get_optimizer)
         optimizer = get_optimizer(self.model)
+
         # Initializing the learning rate schedule, if necessary
-        if hasattr(config, 'lr_schedule'):
+        if 'lr_schedule' in config:
             lr_schedule = optimizers.get_lr_schedule(optimizer,
-                                                     iterations_per_epoch,
+                                                     len(train_loader),
                                                      train_loader)
         else:
             lr_schedule = None
 
+        self.model.to(self.device)
+        self.model.train()
+
         # operate the local training
         supported_modalities = trainset.supported_modalities
         # in order to blend the gradients in the server side
diff --git a/examples/afl/afl_trainer.py b/examples/afl/afl_trainer.py
index 494c0e24d..1c2907ca3 100644
--- a/examples/afl/afl_trainer.py
+++ b/examples/afl/afl_trainer.py
@@ -43,13 +43,8 @@ def train_model(self, config, trainset, sampler, cut_layer=None):
                                                        batch_size=batch_size,
                                                        sampler=sampler)
 
-        iterations_per_epoch = np.ceil(len(trainset) / batch_size).astype(int)
         epochs = config['epochs']
 
-        # Sending the model to the device used for training
-        self.model.to(self.device)
-        self.model.train()
-
         # Initializing the loss criterion
         _loss_criterion = getattr(self, "loss_criterion", None)
         if callable(_loss_criterion):
@@ -63,13 +58,16 @@ def train_model(self, config, trainset, sampler, cut_layer=None):
         optimizer = get_optimizer(self.model)
 
         # Initializing the learning rate schedule, if necessary
-        if hasattr(config, 'lr_schedule'):
+        if 'lr_schedule' in config:
             lr_schedule = optimizers.get_lr_schedule(optimizer,
-                                                     iterations_per_epoch,
+                                                     len(train_loader),
                                                      train_loader)
         else:
             lr_schedule = None
 
+        self.model.to(self.device)
+        self.model.train()
+
         for epoch in range(1, epochs + 1):
             # Use a default training loop
             for batch_id, (examples, labels) in enumerate(train_loader):
diff --git a/examples/cs_maml/cs_maml_server.py b/examples/cs_maml/cs_maml_server.py
index 08d9c6fb5..7c3970b05 100644
--- a/examples/cs_maml/cs_maml_server.py
+++ b/examples/cs_maml/cs_maml_server.py
@@ -141,7 +141,7 @@ async def process_clients(self, client_info):
         if self.do_personalization_test:
             client = client_info
         else:
-            client = client_info[1]
+            client = client_info[2]
             client_staleness = self.current_round - client['starting_round']
 
             self.updates.append(
diff --git a/examples/fedsaw/fedsaw_MNIST_lenet5.yml b/examples/fedsaw/fedsaw_MNIST_lenet5.yml
index 77ad20cdb..2c3785717 100644
--- a/examples/fedsaw/fedsaw_MNIST_lenet5.yml
+++ b/examples/fedsaw/fedsaw_MNIST_lenet5.yml
@@ -9,7 +9,7 @@ clients:
     per_round: 1
 
     # Should the clients compute test accuracy locally?
-    do_test: true
+    do_test: false
 
     # Whether simulate clients or not
     simulation: true
@@ -96,6 +96,6 @@ algorithm:
 
 results:
     # Write the following parameter(s) into a CSV
-    types: round, accuracy, average_accuracy, elapsed_time, comm_time, round_time
+    types: round, accuracy, elapsed_time, comm_time, round_time, comm_overhead
 
     result_path: ./results/test
diff --git a/examples/fedsaw/fedsaw_edge.py b/examples/fedsaw/fedsaw_edge.py
index 0f18a90e9..7bb8c2d53 100644
--- a/examples/fedsaw/fedsaw_edge.py
+++ b/examples/fedsaw/fedsaw_edge.py
@@ -18,7 +18,6 @@
 @dataclass
 class Report(edge.Report):
     """ Client report, to be sent to the federated learning server. """
-    comm_overhead: float
 
 
 class Client(edge.Client):
diff --git a/examples/fedsaw/fedsaw_server.py b/examples/fedsaw/fedsaw_server.py
index 78e3e99bb..dc3c4e2d7 100644
--- a/examples/fedsaw/fedsaw_server.py
+++ b/examples/fedsaw/fedsaw_server.py
@@ -3,12 +3,8 @@
 as either central or edge servers.
 """
 
-import logging
 import math
-import os
-import pickle
 import statistics
-import sys
 
 import torch
 
@@ -122,15 +118,6 @@ def get_record_items_values(self):
         record_items_values = super().get_record_items_values()
         record_items_values['pruning_amount'] = Config().clients.pruning_amount
 
-        if Config().is_central_server():
-            edge_comm_overhead = sum([
-                report.comm_overhead for (__, report, __, __) in self.updates
-            ])
-            record_items_values[
-                'comm_overhead'] = edge_comm_overhead + self.comm_overhead
-        else:
-            record_items_values['comm_overhead'] = self.comm_overhead
-
         return record_items_values
 
     async def wrap_up_processing_reports(self):
@@ -139,83 +126,3 @@ async def wrap_up_processing_reports(self):
 
         if Config().is_central_server():
             self.update_pruning_amount_list()
-            self.comm_overhead = 0
-
-    async def send(self, sid, payload, client_id) -> None:
-        """ Sending a new data payload to the client using either S3 or socket.io. """
-        # First apply outbound processors, if any
-        payload = self.outbound_processor.process(payload)
-
-        metadata = {'id': client_id}
-
-        if self.s3_client is not None:
-            s3_key = f'server_payload_{os.getpid()}_{self.current_round}'
-            self.s3_client.send_to_s3(s3_key, payload)
-            data_size = sys.getsizeof(pickle.dumps(payload))
-            metadata['s3_key'] = s3_key
-        else:
-            data_size = 0
-
-            if isinstance(payload, list):
-                for data in payload:
-                    _data = pickle.dumps(data)
-                    await self.send_in_chunks(_data, sid, client_id)
-                    data_size += sys.getsizeof(_data)
-
-            else:
-                _data = pickle.dumps(payload)
-                await self.send_in_chunks(_data, sid, client_id)
-                data_size = sys.getsizeof(_data)
-
-        await self.sio.emit('payload_done', metadata, room=sid)
-
-        logging.info("[%s] Sent %.2f MB of payload data to client #%d.", self,
-                     data_size / 1024**2, client_id)
-        self.comm_overhead += data_size / 1024**2
-
-    async def client_report_arrived(self, sid, client_id, report):
-        """ Upon receiving a report from a client. """
-        self.reports[sid] = pickle.loads(report)
-        self.client_payload[sid] = None
-        self.client_chunks[sid] = []
-
-        if self.comm_simulation:
-            model_name = Config().trainer.model_name if hasattr(
-                Config().trainer, 'model_name') else 'custom'
-            checkpoint_path = Config().params['checkpoint_path']
-            payload_filename = f"{checkpoint_path}/{model_name}_client_{client_id}.pth"
-            with open(payload_filename, 'rb') as payload_file:
-                self.client_payload[sid] = pickle.load(payload_file)
-
-            data_size = sys.getsizeof(pickle.dumps(self.client_payload[sid]))
-            logging.info(
-                "[%s] Received %.2f MB of payload data from client #%d (simulated).",
-                self, data_size / 1024**2, client_id)
-
-            self.comm_overhead += data_size / 1024**2
-
-            await self.process_client_info(client_id, sid)
-
-    async def client_payload_done(self, sid, client_id, s3_key=None):
-        """ Upon receiving all the payload from a client, either via S3 or socket.io. """
-        if s3_key is None:
-            assert self.client_payload[sid] is not None
-
-            payload_size = 0
-            if isinstance(self.client_payload[sid], list):
-                for _data in self.client_payload[sid]:
-                    payload_size += sys.getsizeof(pickle.dumps(_data))
-            else:
-                payload_size = sys.getsizeof(
-                    pickle.dumps(self.client_payload[sid]))
-        else:
-            self.client_payload[sid] = self.s3_client.receive_from_s3(s3_key)
-            payload_size = sys.getsizeof(pickle.dumps(
-                self.client_payload[sid]))
-
-        logging.info("[%s] Received %.2f MB of payload data from client #%d.",
-                     self, payload_size / 1024**2, client_id)
-
-        self.comm_overhead += payload_size / 1024**2
-
-        await self.process_client_info(client_id, sid)
diff --git a/examples/fei/fei_trainer.py b/examples/fei/fei_trainer.py
index e39e03f49..556ad1e2b 100644
--- a/examples/fei/fei_trainer.py
+++ b/examples/fei/fei_trainer.py
@@ -35,13 +35,8 @@ def train_model(self, config, trainset, sampler, cut_layer=None):
                                                        batch_size=batch_size,
                                                        sampler=sampler)
 
-        iterations_per_epoch = np.ceil(len(trainset) / batch_size).astype(int)
         epochs = config['epochs']
 
-        # Sending the model to the device used for training
-        self.model.to(self.device)
-        self.model.train()
-
         # Initializing the loss criterion
         _loss_criterion = getattr(self, "loss_criterion", None)
         if callable(_loss_criterion):
@@ -55,13 +50,16 @@ def train_model(self, config, trainset, sampler, cut_layer=None):
         optimizer = get_optimizer(self.model)
 
         # Initializing the learning rate schedule, if necessary
-        if hasattr(config, 'lr_schedule'):
+        if 'lr_schedule' in config:
             lr_schedule = optimizers.get_lr_schedule(optimizer,
-                                                     iterations_per_epoch,
+                                                     len(train_loader),
                                                      train_loader)
         else:
             lr_schedule = None
 
+        self.model.to(self.device)
+        self.model.train()
+
         total_loss = 0
         for epoch in range(1, epochs + 1):
             # Use a default training loop
diff --git a/examples/fl_maml/fl_maml_server.py b/examples/fl_maml/fl_maml_server.py
index 8441baf31..767338f35 100644
--- a/examples/fl_maml/fl_maml_server.py
+++ b/examples/fl_maml/fl_maml_server.py
@@ -107,7 +107,7 @@ async def process_clients(self, client_info):
         if self.do_personalization_test:
             client = client_info
         else:
-            client = client_info[1]
+            client = client_info[2]
             client_staleness = self.current_round - client['starting_round']
 
             self.updates.append(
diff --git a/examples/fl_maml/fl_maml_trainer.py b/examples/fl_maml/fl_maml_trainer.py
index 1ee97b98e..326fa7ef3 100644
--- a/examples/fl_maml/fl_maml_trainer.py
+++ b/examples/fl_maml/fl_maml_trainer.py
@@ -15,6 +15,7 @@
 
 class Trainer(basic.Trainer):
     """A federated learning trainer for personalized FL using MAML algorithm."""
+
     def __init__(self, model=None):
         """Initializing the trainer with the provided model."""
         super().__init__(model=model)
@@ -37,19 +38,14 @@ def train_model(self, config, trainset, sampler, cut_layer=None):
                                                        batch_size=batch_size,
                                                        sampler=sampler)
 
-        iterations_per_epoch = np.ceil(len(trainset) / batch_size).astype(int)
         epochs = config['epochs']
 
-        # Sending the model to the device used for training
-        self.model.to(self.device)
-        self.model.train()
-
         # Initializing the loss criterion
         _loss_criterion = getattr(self, "loss_criterion", None)
         if callable(_loss_criterion):
             loss_criterion = self.loss_criterion(self.model)
         else:
-            loss_criterion = nn.CrossEntropyLoss()
+            loss_criterion = torch.nn.CrossEntropyLoss()
 
         # Initializing the optimizer for the second stage of MAML
         # The learning rate here is the meta learning rate (beta)
@@ -58,13 +54,17 @@ def train_model(self, config, trainset, sampler, cut_layer=None):
                                     momentum=Config().trainer.momentum,
                                     weight_decay=Config().trainer.weight_decay)
 
+        # Initializing the learning rate schedule, if necessary
         # Initializing the schedule for meta learning rate, if necessary
-        if hasattr(config, 'meta_lr_schedule'):
+        if 'meta_lr_schedule' in config:
             meta_lr_schedule = optimizers.get_lr_schedule(
-                optimizer, iterations_per_epoch, train_loader)
+                optimizer, len(train_loader), train_loader)
         else:
             meta_lr_schedule = None
 
+        self.model.to(self.device)
+        self.model.train()
+
         for epoch in range(1, epochs + 1):
             # Copy the current model due to using MAML
             current_model = copy.deepcopy(self.model)
@@ -74,22 +74,22 @@ def train_model(self, config, trainset, sampler, cut_layer=None):
 
             # Initializing the optimizer for the first stage of MAML
             # The learning rate here is the alpha in the paper
-            temp_optimizer = torch.optim.SGD(
+            local_optimizer = torch.optim.SGD(
                 current_model.parameters(),
                 lr=Config().trainer.learning_rate,
                 momentum=Config().trainer.momentum,
                 weight_decay=Config().trainer.weight_decay)
 
             # Initializing the learning rate schedule, if necessary
-            if hasattr(config, 'lr_schedule'):
+            if 'lr_schedule' in config:
                 lr_schedule = optimizers.get_lr_schedule(
-                    temp_optimizer, iterations_per_epoch, train_loader)
+                    local_optimizer, len(train_loader), train_loader)
             else:
                 lr_schedule = None
 
             # The first stage of MAML
             # Use half of the training dataset
-            self.training_per_stage(1, temp_optimizer, lr_schedule,
+            self.training_per_stage(1, local_optimizer, lr_schedule,
                                     train_loader, cut_layer, current_model,
                                     loss_criterion, log_interval, epoch,
                                     epochs)
@@ -151,7 +151,7 @@ def training_per_stage(self, stage_id, optimizer, lr_schedule,
             if lr_schedule is not None:
                 lr_schedule.step()
 
-    def test_process(self, config, testset, sampler=None):
+    def test_process(self, config, testset, sampler=None, **kwargs):
         """ A customized testing loop for personalized FL. """
         self.model.to(self.device)
         self.model.eval()
diff --git a/examples/mistnetplus/split_learning_trainer.py b/examples/mistnetplus/split_learning_trainer.py
index 36a08b236..647ff7ed9 100644
--- a/examples/mistnetplus/split_learning_trainer.py
+++ b/examples/mistnetplus/split_learning_trainer.py
@@ -36,18 +36,14 @@ def train_model(self, config, trainset, sampler, cut_layer=None):
                                                        batch_size=batch_size,
                                                        sampler=sampler)
 
-        iterations_per_epoch = np.ceil(len(trainset) / batch_size).astype(int)
-
-        # Sending the model to the device used for training
-        self.model.to(self.device)
-        self.model.train()
+        epochs = config['epochs']
 
         # Initializing the loss criterion
         _loss_criterion = getattr(self, "loss_criterion", None)
         if callable(_loss_criterion):
             loss_criterion = self.loss_criterion(self.model)
         else:
-            loss_criterion = nn.CrossEntropyLoss()
+            loss_criterion = torch.nn.CrossEntropyLoss()
 
         # Initializing the optimizer
         get_optimizer = getattr(self, "get_optimizer",
@@ -55,13 +51,16 @@ def train_model(self, config, trainset, sampler, cut_layer=None):
         optimizer = get_optimizer(self.model)
 
         # Initializing the learning rate schedule, if necessary
-        if hasattr(Config().trainer, 'lr_schedule'):
+        if 'lr_schedule' in config:
             lr_schedule = optimizers.get_lr_schedule(optimizer,
-                                                     iterations_per_epoch,
+                                                     len(train_loader),
                                                      train_loader)
         else:
             lr_schedule = None
 
+        self.model.to(self.device)
+        self.model.train()
+
         for __, (examples, labels) in enumerate(train_loader):
             examples, labels = examples.to(self.device), labels.to(self.device)
             optimizer.zero_grad()
diff --git a/examples/sub_fedavg/subfedavg_trainer.py b/examples/sub_fedavg/subfedavg_trainer.py
index 5cab17831..96037ce1a 100644
--- a/examples/sub_fedavg/subfedavg_trainer.py
+++ b/examples/sub_fedavg/subfedavg_trainer.py
@@ -18,6 +18,7 @@
 
 class Trainer(basic.Trainer):
     """A federated learning trainer for Sub-Fedavg algorithm."""
+
     def __init__(self, model=None):
         """Initializing the trainer with the provided model."""
         super().__init__(model=model)
@@ -55,17 +56,8 @@ def train_model(self, config, trainset, sampler, cut_layer):
                                                        batch_size=batch_size,
                                                        sampler=sampler)
 
-        iterations_per_epoch = np.ceil(len(trainset) / batch_size).astype(int)
         epochs = config['epochs']
 
-        if not self.made_init_mask:
-            self.mask = pruning_processor.make_init_mask(self.model)
-            self.made_init_mask = True
-
-        # Sending the model to the device used for training
-        self.model.to(self.device)
-        self.model.train()
-
         # Initializing the loss criterion
         _loss_criterion = getattr(self, "loss_criterion", None)
         if callable(_loss_criterion):
@@ -78,6 +70,17 @@ def train_model(self, config, trainset, sampler, cut_layer):
                                 optimizers.get_optimizer)
         optimizer = get_optimizer(self.model)
 
+        # Initializing the learning rate schedule, if necessary
+        if 'lr_schedule' in config:
+            lr_schedule = optimizers.get_lr_schedule(optimizer,
+                                                     len(train_loader),
+                                                     train_loader)
+        else:
+            lr_schedule = None
+
+        self.model.to(self.device)
+        self.model.train()
+
         # Initializing the learning rate schedule, if necessary
         if hasattr(config, 'lr_schedule'):
             lr_schedule = optimizers.get_lr_schedule(optimizer,
diff --git a/plato/__init__.py b/plato/__init__.py
index 3d26edf77..df1243329 100644
--- a/plato/__init__.py
+++ b/plato/__init__.py
@@ -1,1 +1,1 @@
-__version__ = "0.4.1"
+__version__ = "0.4.2"
diff --git a/plato/algorithms/fedavg_gan.py b/plato/algorithms/fedavg_gan.py
new file mode 100644
index 000000000..74210882b
--- /dev/null
+++ b/plato/algorithms/fedavg_gan.py
@@ -0,0 +1,78 @@
+"""
+The federated averaging algorithm for GAN model.
+"""
+from collections import OrderedDict
+
+from plato.algorithms import fedavg
+from plato.trainers.base import Trainer
+
+
+class Algorithm(fedavg.Algorithm):
+    """ Federated averaging algorithm for GAN models, used by both the client and the server. """
+
+    def __init__(self, trainer: Trainer):
+        super().__init__(trainer=trainer)
+        self.generator = self.model.generator
+        self.discriminator = self.model.discriminator
+
+    def compute_weight_deltas(self, weights_received):
+        """ Extract the weights received from a client and compute the updates. """
+        baseline_weights_gen, baseline_weights_disc = self.extract_weights()
+
+        deltas = []
+        for weight_gen, weight_disc in weights_received:
+            delta_gen = OrderedDict()
+            for name, current_weight in weight_gen.items():
+                baseline = baseline_weights_gen[name]
+
+                delta = current_weight - baseline
+                delta_gen[name] = delta
+
+            delta_disc = OrderedDict()
+            for name, current_weight in weight_disc.items():
+                baseline = baseline_weights_disc[name]
+
+                delta = current_weight - baseline
+                delta_disc[name] = delta
+
+            deltas.append((delta_gen, delta_disc))
+
+        return deltas
+
+    def update_weights(self, deltas):
+        """ Update the existing model weights. """
+        baseline_weights_gen, baseline_weights_disc = self.extract_weights()
+        update_gen, update_disc = deltas
+
+        updated_weights_gen = OrderedDict()
+        for name, weight in baseline_weights_gen.items():
+            updated_weights_gen[name] = weight + update_gen[name]
+
+        updated_weights_disc = OrderedDict()
+        for name, weight in baseline_weights_disc.items():
+            updated_weights_disc[name] = weight + update_disc[name]
+
+        return updated_weights_gen, updated_weights_disc
+
+    def extract_weights(self, model=None):
+        """ Extract weights from the model. """
+        generator = self.generator
+        discriminator = self.discriminator
+        if model is not None:
+            generator = model.generator
+            discriminator = model.discriminator
+
+        gen_weight = generator.cpu().state_dict()
+        disc_weight = discriminator.cpu().state_dict()
+
+        return gen_weight, disc_weight
+
+    def load_weights(self, weights):
+        """ Load the model weights passed in as a parameter. """
+        weights_gen, weights_disc = weights
+        # The client might only receive one or none of the Generator
+        # and Discriminator model weight.
+        if weights_gen is not None:
+            self.generator.load_state_dict(weights_gen, strict=True)
+        if weights_disc is not None:
+            self.discriminator.load_state_dict(weights_disc, strict=True)
diff --git a/plato/algorithms/registry.py b/plato/algorithms/registry.py
index 41071b6f8..28a01c20e 100644
--- a/plato/algorithms/registry.py
+++ b/plato/algorithms/registry.py
@@ -30,11 +30,13 @@
     from plato.algorithms import (
         fedavg,
         mistnet,
+        fedavg_gan,
     )
 
     registered_algorithms = OrderedDict([
         ('fedavg', fedavg.Algorithm),
         ('mistnet', mistnet.Algorithm),
+        ('fedavg_gan', fedavg_gan.Algorithm),
     ])
 
 
diff --git a/plato/clients/edge.py b/plato/clients/edge.py
index 2adf175eb..162fbc446 100644
--- a/plato/clients/edge.py
+++ b/plato/clients/edge.py
@@ -15,10 +15,12 @@ class Report(simple.Report):
     """ Client report, to be sent to the federated learning server. """
     average_accuracy: float
     client_id: str
+    edge_server_comm_overhead: float
 
 
 class Client(base.Client):
     """ A federated learning client at the edge server in a cross-silo training workload. """
+
     def __init__(self, server):
         super().__init__()
         self.server = server
@@ -69,7 +71,9 @@ async def train(self):
         # Generate a report for the central server
         self.report = Report(self.server.total_samples, accuracy,
                              training_time, comm_time, False, average_accuracy,
-                             self.client_id)
+                             self.client_id, self.server.comm_overhead)
+
+        self.server.comm_overhead = 0
 
         return self.report, weights
 
diff --git a/plato/datasources/celeba.py b/plato/datasources/celeba.py
index eb0b14363..474439740 100644
--- a/plato/datasources/celeba.py
+++ b/plato/datasources/celeba.py
@@ -56,26 +56,32 @@ def __init__(self):
         else:
             target_types = ['attr', 'identity']
 
-        image_size = 32
+        image_size = 64
+        if hasattr(Config().data, 'celeba_img_size'):
+            image_size = Config().data.celeba_img_size
+
         _transform = transforms.Compose([
             transforms.Resize(image_size),
             transforms.CenterCrop(image_size),
             transforms.ToTensor(),
             transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
         ])
+        _target_transform = None
+        if target_types:
+            _target_transform = DataSource._target_transform
 
         self.trainset = CelebA(root=_path,
                                split='train',
                                target_type=target_types,
                                download=False,
                                transform=_transform,
-                               target_transform=DataSource._target_transform)
+                               target_transform=_target_transform)
         self.testset = CelebA(root=_path,
                               split='test',
                               target_type=target_types,
                               download=False,
                               transform=_transform,
-                              target_transform=DataSource._target_transform)
+                              target_transform=_target_transform)
 
     @staticmethod
     def _target_transform(label):
@@ -100,7 +106,10 @@ def _target_transform(label):
 
     @staticmethod
     def input_shape():
-        return [162770, 3, 32, 32]
+        image_size = 64
+        if hasattr(Config().data, 'celeba_img_size'):
+            image_size = Config().data.celeba_img_size
+        return [162770, 3, image_size, image_size]
 
     def num_train_examples(self):
         return 162770
diff --git a/plato/datasources/coco.py b/plato/datasources/coco.py
index 8d92216bb..bd6c4d78e 100644
--- a/plato/datasources/coco.py
+++ b/plato/datasources/coco.py
@@ -67,10 +67,10 @@ def __init__(self):
         _path = Config().params['data_path']
         self._data_path_process(data_path=_path, base_data_name=self.data_name)
 
-        base_data_path = self.mm_data_info["base_data_dir_path"]
+        base_data_path = self.mm_data_info["data_path"]
         raw_data_name = self.data_name + "Raw"
         raw_data_path = os.path.join(base_data_path, raw_data_name)
-        if not self._exist_judgement(raw_data_path):
+        if not self._exists(raw_data_path):
             os.makedirs(raw_data_path, exist_ok=True)
 
         download_train_url = Config().data.download_train_url
@@ -90,18 +90,18 @@ def __init__(self):
             split_download_url = splits_downalods[split_name]
             split_file_name = self._download_arrange_data(
                 download_url_address=split_download_url,
-                put_data_dir=raw_data_path,
+                data_path=raw_data_path,
                 extract_to_dir=split_path)
-            # rename of the extracted file to "images"
-            extracted_dir_path = os.path.join(split_path, split_file_name)
-            rename_dir_path = os.path.join(split_path, "images")
-            os.rename(src=extracted_dir_path, dst=rename_dir_path)
+            # renaming the extracted file to "images"
+            extracted_path = os.path.join(split_path, split_file_name)
+            renamed_path = os.path.join(split_path, "images")
+            os.rename(src=extracted_path, dst=renamed_path)
 
         # Download the annotation
         self._download_arrange_data(
             download_url_address=download_annotation_url,
-            put_data_dir=raw_data_path)
-        anno_dir_path = os.path.join(raw_data_path, "annotations")
+            data_path=raw_data_path)
+        annotation_path = os.path.join(raw_data_path, "annotations")
 
         # Move the annotation to each split
         splits_caption_name = {
@@ -112,6 +112,6 @@ def __init__(self):
             split_caption_name = splits_caption_name[split_name]
             to_split_path = os.path.join(self.splits_info[split_name]["path"],
                                          "captions.json")
-            shutil.copyfile(src=os.path.join(anno_dir_path,
+            shutil.copyfile(src=os.path.join(annotation_path,
                                              split_caption_name),
                             dst=to_split_path)
diff --git a/plato/datasources/datalib/flickr30kE_utils.py b/plato/datasources/datalib/flickr30kE_utils.py
index 6b0aef854..0e56f6627 100644
--- a/plato/datasources/datalib/flickr30kE_utils.py
+++ b/plato/datasources/datalib/flickr30kE_utils.py
@@ -13,22 +13,22 @@
 
 def phrase_boxes_alignment(flatten_boxes, ori_phrases_boxes):
     """ align the bounding boxes with corresponding phrases. """
-    phrases_boxes = list()
+    phrases_boxes = []
 
-    ori_pb_boxes_count = list()
+    ori_pb_boxes_count = []
     for ph_boxes in ori_phrases_boxes:
         ori_pb_boxes_count.append(len(ph_boxes))
 
     strat_point = 0
     for pb_boxes_num in ori_pb_boxes_count:
-        sub_boxes = list()
+        sub_boxes = []
         for i in range(strat_point, strat_point + pb_boxes_num):
             sub_boxes.append(flatten_boxes[i])
 
         strat_point += pb_boxes_num
         phrases_boxes.append(sub_boxes)
 
-    pb_boxes_count = list()
+    pb_boxes_count = []
     for ph_boxes in phrases_boxes:
         pb_boxes_count.append(len(ph_boxes))
 
@@ -39,7 +39,7 @@ def phrase_boxes_alignment(flatten_boxes, ori_phrases_boxes):
 
 def filter_bad_boxes(boxes_coor):
     """ Filter the boxes with wrong coordinates """
-    filted_boxes = list()
+    filted_boxes = []
     for box_coor in boxes_coor:
         [xmin, ymin, xmax, ymax] = box_coor
         if xmin < xmax and ymin < ymax:
@@ -117,7 +117,7 @@ def get_sentence_data(parse_file_path):
 
 
 def get_annotations(parse_file_path):
-    """ Parses the xml files in the Flickr30K Entities dataset
+    """ Parses the xml files in the Flickr30K Entities dataset.
     Args:
         parse_file_path - full file path to the annotations file to parse
     Return:
@@ -161,7 +161,7 @@ def get_annotations(parse_file_path):
 
 
 def align_anno_sent(image_sents, image_annos):
-    """Align the items in annotations and sentences
+    """Align the items in annotations and sentences.
 
     Args:
         image_sents ([list]): [each itme is a dict that contains 'sentence', 'phrases']
@@ -173,14 +173,14 @@ def align_anno_sent(image_sents, image_annos):
                 items because for one image, there are 5 sentences. Sometimes,
                 some sentences are useless, making the number of items less than 5]
     """
-    aligned_items = list()  # each item is a dict
+    aligned_items = []  # each item is a dict
     for sent_info in image_sents:
 
         img_sent = sent_info["sentence"]
-        img_sent_phrases = list()
-        img_sent_phrases_type = list()
-        img_sent_phrases_id = list()
-        img_sent_phrases_boxes = list()
+        img_sent_phrases = []
+        img_sent_phrases_type = []
+        img_sent_phrases_id = []
+        img_sent_phrases_boxes = []
         for phrase_info_idx in range(len(sent_info["phrases"])):
             phrase_info = sent_info["phrases"][phrase_info_idx]
 
@@ -226,9 +226,11 @@ def integrate_data_to_json(splits_info,
                            split_wise=True,
                            globally=True):
     """ Integrate the data into one json file that contains aligned
-         annotation-sentence for each image.
-        The integrated data info is presented as the dict type.
-         Each item in dict contains image and one of its annotation.
+        annotation-sentence for each image.
+        
+        The integrated data info is presented as a dict type.
+        
+        Each item in dict contains image and one of its annotation.
 
         For example, one randomly item:
             {
@@ -244,9 +246,10 @@ def integrate_data_to_json(splits_info,
             ....
             }
         """
+
     def operate_integration(images_name, images_annotations_path,
                             images_sentences_path):
-        """ Obtain the integrated for images """
+        """ Obtain the integrated for images. """
         integrated_data = dict()
         for image_name_idx, image_name in enumerate(images_name):
             image_sent_path = images_sentences_path[image_name_idx]
@@ -270,9 +273,8 @@ def operate_integration(images_name, images_annotations_path,
             save_path = os.path.join(path,
                                      split_type + "_integrated_data.json")
             if os.path.exists(save_path):
-                warn_info = ("Integrated {}, but file already existed"
-                             ).format(split_type)
-                logging.info(warn_info)
+                logging.info("Integrating %s: the file already exists.",
+                             split_type)
                 continue
 
             split_data_types_samples_path = []
@@ -297,16 +299,16 @@ def operate_integration(images_name, images_annotations_path,
                 images_name=split_data_types_samples_path[0],
                 images_annotations_path=split_data_types_samples_path[1],
                 images_sentences_path=split_data_types_samples_path[2])
-            with open(save_path, 'w') as outfile:
+            with open(save_path, 'w', encoding='utf-8') as outfile:
                 json.dump(split_integrated_data, outfile)
 
-            logging.info("Integration for %s Done!", split_type)
+            logging.info("The integration process for %s is done.", split_type)
 
     if globally:
-        save_path = os.path.join(mm_data_info["base_data_dir_path"],
+        save_path = os.path.join(mm_data_info["data_path"],
                                  "total_integrated_data.json")
         if os.path.exists(save_path):
-            logging.info("Gloablly Integrated file, Existed!")
+            logging.info("Gloablly integrated file already exists.")
             return
 
         raw_data_types_samples_path = []
@@ -327,7 +329,7 @@ def operate_integration(images_name, images_annotations_path,
             images_name=raw_data_types_samples_path[0],
             images_annotations_path=raw_data_types_samples_path[1],
             images_sentences_path=raw_data_types_samples_path[2])
-        with open(save_path, 'w') as outfile:
+        with open(save_path, 'w', encoding='utf-8') as outfile:
             json.dump(global_integrated_data, outfile)
 
-        logging.info("Integration for the whole dataset, Done!")
+        logging.info("Integration for the whole dataset, Done.")
diff --git a/plato/datasources/datalib/frames_extraction_tools.py b/plato/datasources/datalib/frames_extraction_tools.py
index e866e8c56..dd54ff7e6 100644
--- a/plato/datasources/datalib/frames_extraction_tools.py
+++ b/plato/datasources/datalib/frames_extraction_tools.py
@@ -93,6 +93,7 @@ def extract_optical_flow(videos_items):
 
 class VideoFramesExtractor(modality_extraction_base.VideoExtractorBase):
     """ The class for extracting the frame the video """
+
     def __init__(self,
                  video_src_dir,
                  dir_level=2,
@@ -150,29 +151,29 @@ def build_optical_flow_frames(
                 len(self.videos_path_list) * [self.is_classname_contained]))
 
     def build_frames_gpu(self,
-                         rgb_out_dir_path,
-                         flow_our_dir_path,
+                         rgb_out__path,
+                         flow_our__path,
                          new_short=1,
                          new_width=0,
                          new_height=0):
         """ Get the optical flow frame based on the GPU  """
-        self.build_rgb_frames(rgb_out_dir_path,
+        self.build_rgb_frames(rgb_out__path,
                               new_short=new_short,
                               new_width=new_width,
                               new_height=new_height)
-        self.build_optical_flow_frames(flow_our_dir_path,
+        self.build_optical_flow_frames(flow_our__path,
                                        new_short=new_short,
                                        new_width=new_width,
                                        new_height=new_height)
 
     def build_full_frames_gpu(self,
-                              to_dir_path,
+                              to__path,
                               new_short=1,
                               new_width=0,
                               new_height=0):
         """ The interface for extracting all frames based on the GPU  """
-        self.build_frames_gpu(rgb_out_dir_path=to_dir_path,
-                              flow_our_dir_path=to_dir_path,
+        self.build_frames_gpu(rgb_out__path=to__path,
+                              flow_our__path=to__path,
                               new_short=new_short,
                               new_width=new_width,
                               new_height=new_height)
diff --git a/plato/datasources/femnist.py b/plato/datasources/femnist.py
index 8bf232664..ef6cc4816 100644
--- a/plato/datasources/femnist.py
+++ b/plato/datasources/femnist.py
@@ -74,11 +74,11 @@ def __init__(self, client_id=0):
         if client_id == 0:
             # If we are on the federated learning server
             data_dir = os.path.join(root_path, 'test')
-            data_url = "https://jiangzhifeng.s3.us-east-2.amazonaws.com/FEMNIST/test/" \
+            data_url = "http://iqua.ece.toronto.edu/FEMNIST/test/" \
                        + str(client_id) + ".zip"
         else:
             data_dir = os.path.join(root_path, 'train')
-            data_url = "https://jiangzhifeng.s3.us-east-2.amazonaws.com/FEMNIST/train/" \
+            data_url = "http://iqua.ece.toronto.edu/FEMNIST/train/" \
                        + str(client_id) + ".zip"
 
         if not os.path.exists(os.path.join(data_dir, str(client_id))):
diff --git a/plato/datasources/flickr30k_entities.py b/plato/datasources/flickr30k_entities.py
index 53afc29c0..0932b77b8 100644
--- a/plato/datasources/flickr30k_entities.py
+++ b/plato/datasources/flickr30k_entities.py
@@ -220,15 +220,13 @@ def __init__(self):
         self._data_path_process(data_path=_path, base_data_name=self.data_name)
 
         raw_data_name = self.data_name + "Raw"
-        base_data_path = self.mm_data_info["base_data_dir_path"]
+        base_data_path = self.mm_data_info["data_path"]
 
-        download_file_id = Config().data.download_file_id
+        download_url = Config().data.download_url
 
-        self._download_google_driver_arrange_data(
-            download_file_id=download_file_id,
-            extract_download_file_name=raw_data_name,
-            put_data_dir=base_data_path,
-        )
+        self._download_arrange_data(download_url_address=download_url,
+                                    data_path=base_data_path,
+                                    extract_to_dir=base_data_path)
 
         # define the path of different data source,
         #   the annotation is .xml, the sentence is in .txt
@@ -295,10 +293,10 @@ def create_splits_data(self):
                 split_dt_type_path = self.splits_info[split_type][dt_type][
                     "path"]
 
-                if not self._exist_judgement(split_dt_type_path):
+                if not self._exists(split_dt_type_path):
                     os.makedirs(split_dt_type_path, exist_ok=True)
                 else:
-                    logging.info("The path %s does exist", split_dt_type_path)
+                    logging.info("The path %s exists.", split_dt_type_path)
                     continue
 
                 raw_data_type_path = self.mm_data_info[dt_type]["path"]
@@ -311,7 +309,7 @@ def create_splits_data(self):
                 # 2. saving the splited data into the target file
                 data_utils.copy_files(split_samples_path, split_dt_type_path)
 
-        logging.info(" Done!")
+        logging.info("Done.")
 
     def get_phase_data_info(self, phase):
         """ Obtain the data information for the required phrase """
diff --git a/plato/datasources/gym.py b/plato/datasources/gym.py
index b580fcb04..8a647ecd0 100644
--- a/plato/datasources/gym.py
+++ b/plato/datasources/gym.py
@@ -120,50 +120,51 @@ def __init__(self):
         self._data_path_process(data_path=_path, base_data_name=self.data_name)
         self._create_modalities_path(modality_names=self.modality_names)
 
-        base_data_path = self.mm_data_info["base_data_dir_path"]
+        base_data_path = self.mm_data_info["data_path"]
         # define all the dir here
         gym_anno_dir_name = "annotations"
-        self.data_anno_dir_path = os.path.join(base_data_path,
-                                               gym_anno_dir_name)
+        self.data_annotation_path = os.path.join(base_data_path,
+                                                 gym_anno_dir_name)
 
-        self.data_anno_file_path = os.path.join(self.data_anno_dir_path,
+        self.data_anno_file_path = os.path.join(self.data_annotation_path,
                                                 "annotation.json")
-        self.categoty_anno_file_path = os.path.join(self.data_anno_dir_path,
+        self.categoty_anno_file_path = os.path.join(self.data_annotation_path,
                                                     "gym99_categories.txt")
 
         self.raw_videos_path = os.path.join(base_data_path, "videos")
-        self.event_dir_path = os.path.join(base_data_path, "event")
-        self.event_subsection_dir_path = os.path.join(base_data_path,
-                                                      "subactions")
-        self.data_event_anno_file_path = os.path.join(self.data_anno_dir_path,
-                                                      "event_annotation.json")
-        self.event_subsection_frames_dir_path = os.path.join(
+        self.event__path = os.path.join(base_data_path, "event")
+        self.event_subsection__path = os.path.join(base_data_path,
+                                                   "subactions")
+        self.data_event_anno_file_path = os.path.join(
+            self.data_annotation_path, "event_annotation.json")
+        self.event_subsection_frames__path = os.path.join(
             base_data_path, "subaction_rawframes")
-        self.event_subsection_audios_dir_path = os.path.join(
+        self.event_subsection_audios__path = os.path.join(
             base_data_path, "subaction_audios")
 
-        self.event_subsection_audios_fea_dir_path = os.path.join(
+        self.event_subsection_audios_fea__path = os.path.join(
             base_data_path, "subaction_audios_features")
 
         self.rawframes_splits_list_files_into = {
             "train":
-            os.path.join(self.data_anno_dir_path, "gym99_train_rawframes.txt"),
+            os.path.join(self.data_annotation_path,
+                         "gym99_train_rawframes.txt"),
             "val":
-            os.path.join(self.data_anno_dir_path, "gym99_val_rawframes.txt")
+            os.path.join(self.data_annotation_path, "gym99_val_rawframes.txt")
         }
 
         self.audios_splits_list_files_into = {
             "train":
-            os.path.join(self.data_anno_dir_path, "gym99_train_audios.txt"),
+            os.path.join(self.data_annotation_path, "gym99_train_audios.txt"),
             "val":
-            os.path.join(self.data_anno_dir_path, "gym99_val_audios.txt")
+            os.path.join(self.data_annotation_path, "gym99_val_audios.txt")
         }
         self.audio_features_splits_list_files_into = {
             "train":
-            os.path.join(self.data_anno_dir_path,
+            os.path.join(self.data_annotation_path,
                          "gym99_train_audio_features.txt"),
             "val":
-            os.path.join(self.data_anno_dir_path,
+            os.path.join(self.data_annotation_path,
                          "gym99_val_audio_features.txt")
         }
 
@@ -178,28 +179,28 @@ def __init__(self):
 
         _ = self._download_arrange_data(
             download_url_address=set_level_category_url,
-            put_data_dir=self.data_anno_dir_path,
+            data_path=self.data_annotation_path,
             obtained_file_name="set_categories.txt")
 
         _ = self._download_arrange_data(
             download_url_address=g99_categoty_url,
-            put_data_dir=self.data_anno_dir_path,
+            data_path=self.data_annotation_path,
             obtained_file_name="gym99_categories.txt")
 
         _ = self._download_arrange_data(download_url_address=anno_url,
-                                        put_data_dir=self.data_anno_dir_path,
+                                        data_path=self.data_annotation_path,
                                         obtained_file_name="annotation.json")
 
         _ = self._download_arrange_data(
             download_url_address=train_url,
-            put_data_dir=self.data_anno_dir_path,
+            data_path=self.data_annotation_path,
             obtained_file_name="gym99_train_org.txt")
 
         _ = self._download_arrange_data(download_url_address=eval_url,
-                                        put_data_dir=self.data_anno_dir_path,
+                                        data_path=self.data_annotation_path,
                                         obtained_file_name="gym99_val_org.txt")
 
-        if not self._exist_judgement(self.raw_videos_path):
+        if not self._exists(self.raw_videos_path):
 
             logging.info(
                 "Downloading the raw videos for the Gym dataset. This may take a long time."
@@ -211,28 +212,28 @@ def __init__(self):
             logging.info("Done.")
 
         # Trim Videos into Events
-        if not self._exist_judgement(self.event_dir_path):
+        if not self._exists(self.event__path):
             gym_trim.trim_event(video_root=self.raw_videos_path,
                                 anno_file=self.data_anno_file_path,
                                 event_anno_file=self.data_event_anno_file_path,
-                                event_root=self.event_dir_path)
-        if not self._exist_judgement(self.event_subsection_dir_path):
+                                event_root=self.event__path)
+        if not self._exists(self.event_subsection__path):
             gym_trim.trim_subsection(
                 event_anno_file=self.data_event_anno_file_path,
-                event_root=self.event_dir_path,
-                subaction_root=self.event_subsection_dir_path)
+                event_root=self.event__path,
+                subaction_root=self.event_subsection__path)
 
         logging.info("The Gym dataset has been prepared")
         self.extract_videos_rgb_flow_audio()
 
     def extract_videos_rgb_flow_audio(self):
         """ Extract the rgb optical flow audios from the video """
-        src_videos_dir = self.event_subsection_dir_path
-        frames_out_dir_path = self.event_subsection_frames_dir_path
-        rgb_out_dir_path = self.event_subsection_frames_dir_path
-        flow_our_dir_path = self.event_subsection_frames_dir_path
-        audio_out_dir_path = self.event_subsection_audios_dir_path
-        audio_feature_dir_path = self.event_subsection_audios_fea_dir_path
+        src_videos_dir = self.event_subsection__path
+        frames_out__path = self.event_subsection_frames__path
+        rgb_out__path = self.event_subsection_frames__path
+        flow_our__path = self.event_subsection_frames__path
+        audio_out__path = self.event_subsection_audios__path
+        audio_feature__path = self.event_subsection_audios_fea__path
 
         # define the modalities extractor
         vdf_extractor = frames_extraction_tools.VideoFramesExtractor(
@@ -249,44 +250,43 @@ def extract_videos_rgb_flow_audio(self):
             mixed_ext=False)
 
         if torch.cuda.is_available():
-            if not self._exist_judgement(
-                    rgb_out_dir_path) and not self._exist_judgement(
-                        flow_our_dir_path):
+            if not self._exists(rgb_out__path) and not self._exists(
+                    flow_our__path):
                 logging.info(
                     "Extracting frames by GPU from videos in %s to %s.",
-                    src_videos_dir, rgb_out_dir_path)
-                vdf_extractor.build_full_frames_gpu(
-                    to_dir_path=frames_out_dir_path,
-                    new_short=256,
-                    new_width=0,
-                    new_height=0)
+                    src_videos_dir, rgb_out__path)
+                vdf_extractor.build_full_frames_gpu(to__path=frames_out__path,
+                                                    new_short=256,
+                                                    new_width=0,
+                                                    new_height=0)
         else:
-            if not self._exist_judgement(rgb_out_dir_path):
+            if not self._exists(rgb_out__path):
                 logging.info(
                     "Extracting frames by CPU from videos in %s to %s.",
-                    src_videos_dir, rgb_out_dir_path)
-                vdf_extractor.build_frames_cpu(to_dir=frames_out_dir_path)
+                    src_videos_dir, rgb_out__path)
+                vdf_extractor.build_frames_cpu(to_dir=frames_out__path)
 
-        if not self._exist_judgement(audio_out_dir_path):
+        if not self._exists(audio_out__path):
             logging.info("Extracting audios by CPU from videos in %s to %s.",
-                         src_videos_dir, audio_out_dir_path)
-            vda_extractor.build_audios(to_dir=audio_out_dir_path)
+                         src_videos_dir, audio_out__path)
+            vda_extractor.build_audios(to_dir=audio_out__path)
 
-        if not self._exist_judgement(audio_feature_dir_path):
+        if not self._exists(audio_feature__path):
             logging.info(
                 "Extracting audios feature by CPU from audios in %s to %s.",
-                audio_out_dir_path, audio_feature_dir_path)
+                audio_out__path, audio_feature__path)
             # # window_size:32ms hop_size:16ms
 
             vda_extractor.build_audios_features(
-                audio_src_path=audio_out_dir_path,
-                to_dir=audio_feature_dir_path,
+                audio_src_path=audio_out__path,
+                to_dir=audio_feature__path,
                 fft_size=512,  # fft_size / sample_rate is window size
                 hop_size=256)
         # extract the splits data into list files based on the frames information
-        gym_trim.generate_splits_list(data_root=self.event_subsection_dir_path,
-                                      annotation_root=self.data_anno_dir_path,
-                                      frame_data_root=frames_out_dir_path)
+        gym_trim.generate_splits_list(
+            data_root=self.event_subsection__path,
+            annotation_root=self.data_annotation_path,
+            frame_data_root=frames_out__path)
 
         # generate the audio and audio features splits file
         # just copy the frame files to the audio ones
@@ -315,12 +315,12 @@ def correct_current_config(self, loaded_plato_config, mode, modality_name):
         #   The main reason is that the obtained path here is the full path
         cur_rawframes_anno_file_path = self.rawframes_splits_list_files_into[
             mode]
-        cur_rawframes_data_path = self.event_subsection_frames_dir_path
+        cur_rawframes_data_path = self.event_subsection_frames__path
         cur_videos_anno_file_path = None
-        cur_video_data_path = self.event_subsection_dir_path
+        cur_video_data_path = self.event_subsection__path
         cur_audio_feas_anno_file_path = self.audios_splits_list_files_into[
             mode]
-        cur_audio_feas_data_path = self.event_subsection_audios_dir_path
+        cur_audio_feas_data_path = self.event_subsection_audios__path
 
         if modality_name == "rgb" or modality_name == "flow":
             loaded_config["ann_file"] = cur_rawframes_anno_file_path
diff --git a/plato/datasources/kinetics.py b/plato/datasources/kinetics.py
index e3e9d96cf..16db31663 100644
--- a/plato/datasources/kinetics.py
+++ b/plato/datasources/kinetics.py
@@ -155,7 +155,7 @@ def __init__(self):
         # Generate the basic path for the dataset, it performs:
         #   1.- Assign path to self.mm_data_info
         #   2.- Assign splits path to self.splits_info
-        #       where the root path for splits is the base_data_dir_path
+        #       where the root path for splits is the data_path
         #       in self.mm_data_info
         self._data_path_process(data_path=_path, base_data_name=self.data_name)
         # Generate the modalities path for all splits, it performs:
@@ -166,18 +166,19 @@ def __init__(self):
         self._create_modalities_path(modality_names=self.modality_names)
 
         # Set the annotation file path
-        base_data_path = self.mm_data_info["base_data_dir_path"]
+        base_data_path = self.mm_data_info["data_path"]
 
         # Define all the dir here
         kinetics_anno_dir_name = "annotations"
-        self.data_anno_dir_path = os.path.join(base_data_path,
-                                               kinetics_anno_dir_name)
+        self.data_annotation_path = os.path.join(base_data_path,
+                                                 kinetics_anno_dir_name)
 
         for split in ["train", "test", "validate"]:
             split_anno_path = os.path.join(
-                self.data_anno_dir_path, base_data_name + "_" + split + ".csv")
+                self.data_annotation_path,
+                base_data_name + "_" + split + ".csv")
             split_tiny_anno_path = os.path.join(
-                self.data_anno_dir_path,
+                self.data_annotation_path,
                 base_data_name + "_" + split + "_tiny.csv")
             split_name = split if split != "validate" else "val"
             self.splits_info[split_name]["split_anno_file"] = split_anno_path
@@ -194,7 +195,7 @@ def __init__(self):
         #       'video_path': xxx}
         #  the self.mm_data_info can contain
         #   - source_data_path
-        #   - base_data_dir_path
+        #   - data_path
 
         anno_download_url = (
             "https://storage.googleapis.com/deepmind-media/Datasets/{}.tar.gz"
@@ -202,16 +203,16 @@ def __init__(self):
 
         extracted_anno_file_name = self._download_arrange_data(
             download_url_address=anno_download_url,
-            put_data_dir=self.data_anno_dir_path,
+            data_path=self.data_annotation_path,
             obtained_file_name=None)
-        download_anno_path = os.path.join(self.data_anno_dir_path,
+        download_anno_path = os.path.join(self.data_annotation_path,
                                           extracted_anno_file_name)
 
         downloaded_files = os.listdir(download_anno_path)
         for file_name in downloaded_files:
             new_file_name = base_data_name + "_" + file_name
             shutil.move(os.path.join(download_anno_path, file_name),
-                        os.path.join(self.data_anno_dir_path, new_file_name))
+                        os.path.join(self.data_annotation_path, new_file_name))
 
         # Whether to create the tiny dataset
         for split in ["train", "test", "validate"]:
@@ -238,7 +239,7 @@ def __init__(self):
             video_path_format = self.set_modality_path_key_format(
                 modality_name="video")
             video_dir = self.splits_info[split][video_path_format]
-            if not self._exist_judgement(video_dir):
+            if not self._exists(video_dir):
                 num_workers = Config().data.downloader.num_workers
                 # Set the tmp_dir to save the raw video
                 # Then, the raw video will be clipped to save to
@@ -283,10 +284,10 @@ def rename_classes(self, mode):
         """ Rename classes by replacing whitespace to  'Underscore' """
         video_format_path_key = self.set_modality_path_key_format(
             modality_name="video")
-        videos_root_dir_path = self.splits_info[mode][video_format_path_key]
+        videos_root__path = self.splits_info[mode][video_format_path_key]
         videos_dirs_name = [
-            dir_name for dir_name in os.listdir(videos_root_dir_path)
-            if os.path.isdir(os.path.join(videos_root_dir_path, dir_name))
+            dir_name for dir_name in os.listdir(videos_root__path)
+            if os.path.isdir(os.path.join(videos_root__path, dir_name))
         ]
 
         new_videos_dirs_name = [
@@ -294,11 +295,11 @@ def rename_classes(self, mode):
         ]
 
         videos_dirs_path = [
-            os.path.join(videos_root_dir_path, dir_name)
+            os.path.join(videos_root__path, dir_name)
             for dir_name in videos_dirs_name
         ]
         new_videos_dirs_path = [
-            os.path.join(videos_root_dir_path, dir_name)
+            os.path.join(videos_root__path, dir_name)
             for dir_name in new_videos_dirs_name
         ]
         for i, _ in enumerate(videos_dirs_path):
@@ -318,25 +319,25 @@ def extract_videos_rgb_flow_audio(self, mode="train"):
                                                       modality_name="video")
         src_mode_videos_dir = video_data_path
 
-        rgb_out_dir_path = self.get_modality_data_path(mode=mode,
-                                                       modality_name="rgb")
-        flow_our_dir_path = self.get_modality_data_path(mode=mode,
-                                                        modality_name="flow")
-        audio_out_dir_path = self.get_modality_data_path(mode=mode,
-                                                         modality_name="audio")
-        audio_feature_dir_path = self.get_modality_data_path(
+        rgb_out__path = self.get_modality_data_path(mode=mode,
+                                                    modality_name="rgb")
+        flow_our__path = self.get_modality_data_path(mode=mode,
+                                                     modality_name="flow")
+        audio_out__path = self.get_modality_data_path(mode=mode,
+                                                      modality_name="audio")
+        audio_feature__path = self.get_modality_data_path(
             mode=mode, modality_name="audio_feature")
 
         # define the modalities extractor
-        if not self._exist_judgement(rgb_out_dir_path):
+        if not self._exists(rgb_out__path):
             vdf_extractor = frames_extraction_tools.VideoFramesExtractor(
                 video_src_dir=src_mode_videos_dir,
                 dir_level=2,
                 num_worker=8,
                 video_ext="mp4",
                 mixed_ext=False)
-        if not self._exist_judgement(audio_out_dir_path) \
-            or not self._exist_judgement(audio_feature_dir_path):
+        if not self._exists(audio_out__path) \
+            or not self._exists(audio_feature__path):
             vda_extractor = audio_extraction_tools.VideoAudioExtractor(
                 video_src_dir=src_mode_videos_dir,
                 dir_level=2,
@@ -345,49 +346,49 @@ def extract_videos_rgb_flow_audio(self, mode="train"):
                 mixed_ext=False)
 
         if torch.cuda.is_available():
-            if not self._exist_judgement(rgb_out_dir_path):
+            if not self._exists(rgb_out__path):
                 logging.info(
                     "Extracting frames by GPU from videos in %s to %s.",
-                    src_mode_videos_dir, rgb_out_dir_path)
-                vdf_extractor.build_frames_gpu(rgb_out_dir_path,
-                                               flow_our_dir_path,
+                    src_mode_videos_dir, rgb_out__path)
+                vdf_extractor.build_frames_gpu(rgb_out__path,
+                                               flow_our__path,
                                                new_short=1,
                                                new_width=0,
                                                new_height=0)
         else:
-            if not self._exist_judgement(rgb_out_dir_path):
+            if not self._exists(rgb_out__path):
                 logging.info(
                     "Extracting frames by CPU from videos in %s to %s.",
-                    src_mode_videos_dir, rgb_out_dir_path)
-                vdf_extractor.build_frames_cpu(to_dir=rgb_out_dir_path)
+                    src_mode_videos_dir, rgb_out__path)
+                vdf_extractor.build_frames_cpu(to_dir=rgb_out__path)
 
-        if not self._exist_judgement(audio_out_dir_path):
+        if not self._exists(audio_out__path):
             logging.info("Extracting audios by CPU from videos in %s to %s.",
-                         src_mode_videos_dir, audio_out_dir_path)
-            vda_extractor.build_audios(to_dir=audio_out_dir_path)
+                         src_mode_videos_dir, audio_out__path)
+            vda_extractor.build_audios(to_dir=audio_out__path)
 
-        if not self._exist_judgement(audio_feature_dir_path):
+        if not self._exists(audio_feature__path):
             logging.info(
                 "Extracting audios feature by CPU from audios in %s to %s.",
-                audio_out_dir_path, audio_feature_dir_path)
+                audio_out__path, audio_feature__path)
             # # window_size:32ms hop_size:16ms
             vda_extractor.build_audios_features(
-                audio_src_path=audio_out_dir_path,
-                to_dir=audio_feature_dir_path,
+                audio_src_path=audio_out__path,
+                to_dir=audio_feature__path,
                 fft_size=512,  # fft_size / sample_rate is window size
                 hop_size=256)
 
     def extract_splits_list_files(self, data_format, splits):
         """ Extract and generate the split information of current mode/phase """
         output_format = "json"
-        out_path = self.mm_data_info["base_data_dir_path"]
+        out_path = self.mm_data_info["data_path"]
 
         # obtained a dict that contains the required data splits' file path
         #   it can be full data or tiny data
         required_anno_files = obtain_required_anno_files(self.splits_info)
         data_splits_file_info = required_anno_files
         gen_annots_op = modality_data_anntation_tools.GenerateMDataAnnotation(
-            data_src_dir=self.mm_data_info["base_data_dir_path"],
+            data_src_dir=self.mm_data_info["data_path"],
             data_annos_files_info=data_splits_file_info,
             dataset_name=self.data_name,
             data_format=data_format,  # 'rawframes', 'videos', 'audio_features'
@@ -398,9 +399,9 @@ def extract_splits_list_files(self, data_format, splits):
             output_format=output_format)
 
         target_list_regu = f'_{data_format}.{output_format}'
-        if not self._exist_file_in_dir(tg_file_name=target_list_regu,
-                                       search_dir=out_path,
-                                       is_partial_name=True):
+        if not self._file_exists(tg_file_name=target_list_regu,
+                                 search_path=out_path,
+                                 is_partial_name=True):
             logging.info("Extracting annotation list for %s. ", data_format)
 
             gen_annots_op.read_data_splits_csv_info()
diff --git a/plato/datasources/multimodal_base.py b/plato/datasources/multimodal_base.py
index da5f88884..05b38e9c9 100644
--- a/plato/datasources/multimodal_base.py
+++ b/plato/datasources/multimodal_base.py
@@ -24,6 +24,7 @@ class MultiModalDataSource(base.DataSource):
     """
     The training or testing dataset that accommodates custom augmentation and transforms.
     """
+
     def __init__(self):
         super().__init__()
 
@@ -35,11 +36,11 @@ def __init__(self):
 
         # define the information container for the source data
         #  - source_data_path: the original downloaded data
-        #  - base_data_dir_path: the source data used for the model
-        # For some datasets, we directly utilize the base_data_dir_path as
+        #  - data_path: the source data used for the model
+        # For some datasets, we directly utilize the data_path as
         #  there is no need to process the original downloaded data to put them
-        #  in the base_data_dir_path dir.
-        self.mm_data_info = {"source_data_path": "", "base_data_dir_path": ""}
+        #  in the data_path dir.
+        self.mm_data_info = {"source_data_path": "", "data_path": ""}
 
         # define the paths for the splited root data - train, test, and val
         self.splits_info = {
@@ -106,15 +107,13 @@ def _data_path_process(
         """ Generate the data structure based on the defined data path """
 
         # Create the full path by introducing the project path
-        proj_root_path = os.path.abspath(os.curdir)
-        base_data_path = os.path.join(proj_root_path, data_path,
-                                      base_data_name)
+        base_data_path = os.path.join(data_path, base_data_name)
 
         if not os.path.exists(base_data_path):
             os.makedirs(base_data_path)
 
         #
-        self.mm_data_info["base_data_dir_path"] = base_data_path
+        self.mm_data_info["data_path"] = base_data_path
 
         # create the split dirs for current dataset
         for split_type in list(self.splits_info.keys()):
@@ -129,31 +128,31 @@ def _data_path_process(
     def _download_arrange_data(
         self,
         download_url_address,
-        put_data_dir,
+        data_path,
         extract_to_dir=None,
         obtained_file_name=None,
     ):
         """ Download the raw data and arrange the data """
         # Extract to the same dir as the download dir
         if extract_to_dir is None:
-            extract_to_dir = put_data_dir
+            extract_to_dir = data_path
 
         download_file_name = os.path.basename(download_url_address)
-        download_file_path = os.path.join(put_data_dir, download_file_name)
+        download_file_path = os.path.join(data_path, download_file_name)
 
         download_extracted_file_name = download_file_name.split(".")[0]
-        download_extracted_dir_path = os.path.join(
-            extract_to_dir, download_extracted_file_name)
+        download_extracted_path = os.path.join(extract_to_dir,
+                                               download_extracted_file_name)
         # Download the raw data if necessary
-        if not self._exist_judgement(download_file_path):
+        if not self._exists(download_file_path):
             logging.info("Downloading the %s data.....", download_file_name)
             download_url(url=download_url_address,
-                         root=put_data_dir,
+                         root=data_path,
                          filename=obtained_file_name)
 
         # Extract the data to the specific dir
         if ".zip" in download_file_name or ".tar.gz" in download_file_name:
-            if not self._exist_judgement(download_extracted_dir_path):
+            if not self._exists(download_extracted_path):
                 logging.info("Extracting data to %s dir.....", extract_to_dir)
                 extract_archive(from_path=download_file_path,
                                 to_path=extract_to_dir,
@@ -165,48 +164,42 @@ def _download_google_driver_arrange_data(
         self,
         download_file_id,
         extract_download_file_name,
-        put_data_dir,
+        data_path,
     ):
         download_data_file_name = extract_download_file_name + ".zip"
-        download_data_path = os.path.join(put_data_dir,
-                                          download_data_file_name)
-        extract_data_path = os.path.join(put_data_dir,
-                                         extract_download_file_name)
-        if not self._exist_judgement(download_data_path):
+        download_data_path = os.path.join(data_path, download_data_file_name)
+        extract_data_path = os.path.join(data_path, extract_download_file_name)
+        if not self._exists(download_data_path):
             logging.info("Downloading the data to %s", download_data_path)
             download_file_from_google_drive(file_id=download_file_id,
-                                            root=put_data_dir,
+                                            root=data_path,
                                             filename=download_data_file_name)
-        if not self._exist_judgement(extract_data_path):
+        if not self._exists(extract_data_path):
             extract_archive(from_path=download_data_path,
-                            to_path=put_data_dir,
+                            to_path=data_path,
                             remove_finished=True)
 
-    def _exist_file_in_dir(self,
-                           tg_file_name,
-                           search_dir,
-                           is_partial_name=True):
-        """ Judge whether the input file exists in the search_dir. """
+    def _file_exists(self, tg_file_name, search_path, is_partial_name=True):
+        """ Judge whether the input file exists in the search_path. """
         # the tg_file_name matches one file if it match part of the file name
         if is_partial_name:
-            is_included_fuc = lambda src_f_name: tg_file_name in src_f_name
+            is_included = lambda src_f_name: tg_file_name in src_f_name
         else:
-            is_included_fuc = lambda src_f_name: tg_file_name == src_f_name
-        is_existed = any(
-            [is_included_fuc(f_name) for f_name in os.listdir(search_dir)])
+            is_included = lambda src_f_name: tg_file_name == src_f_name
+        exists = any(is_included(f_name) for f_name in os.listdir(search_path))
 
-        return is_existed
+        return exists
 
-    def _exist_judgement(self, target_path):
-        """ Judeg whether the input file/dir existed and whether it contains useful data """
+    def _exists(self, target_path):
+        """ Does the input path/file exist and does the file contain useful data? """
         if not os.path.exists(target_path):
-            logging.info("The path %s does not exist", target_path)
+            logging.info("The path %s does not exist.", target_path)
             return False
 
         # remove all .DS_Store files
         command = ['find', '.', '-name', '".DS_Store"', '-delete']
         command = ' '.join(command)
-        #cmd = f"find . -name ".DS_Store" -delete"
+        # cmd = f"find . -name ".DS_Store" -delete"
         subprocess.check_output(command, shell=True, stderr=subprocess.STDOUT)
 
         def get_size(folder):
@@ -219,7 +212,7 @@ def get_size(folder):
 
         def is_contain_useful_file(target_dir):
             """ Return True once reaching one useful file """
-            for _, _, files in os.walk(target_dir):
+            for __, __, files in os.walk(target_dir):
                 for file in files:
                     # whether a useful file
                     if not file.startswith('.'):
@@ -229,17 +222,17 @@ def is_contain_useful_file(target_dir):
         if os.path.isdir(target_path):
             if get_size(target_path
                         ) == 0 or not is_contain_useful_file(target_path):
-                logging.info("The path %s does exist but contains no data",
+                logging.info("The path %s exists but contains no data.",
                              target_path)
                 return False
-            else:
-                return True
 
-        logging.info("The file %s does exist", target_path)
+            return True
+
+        logging.info("The file %s exists.", target_path)
         return True
 
     def num_modalities(self) -> int:
-        """ Number of modalities """
+        """ The number of modalities. """
         return len(self.modality_names)
 
     @abstractmethod
@@ -247,25 +240,28 @@ def get_phase_dataset(self, phase, modality_sampler):
         """ Obtain the dataset with the modaltiy_sampler for the
             specific phase (train/test/val) """
         raise NotImplementedError(
-            "Please Implement 'get_phase_dataset' method")
+            "Please implement the 'get_phase_dataset' method.")
 
     @abstractmethod
     def get_train_set(self, modality_sampler):
         """ Obtain the train dataset with the modaltiy_sampler """
-        raise NotImplementedError("Please Implement 'get_train_set' method")
+        raise NotImplementedError(
+            "Please implement the 'get_train_set' method.")
 
     @abstractmethod
     def get_test_set(self, modality_sampler):
-        """ btain the test dataset with the modaltiy_sampler """
-        raise NotImplementedError("Please Implement 'get_test_set' method")
+        """ Obtain the test dataset with the modaltiy_sampler. """
+        raise NotImplementedError(
+            "Please implement the 'get_test_set' method.")
 
 
 class MultiModalDataset(torch.utils.data.Dataset):
-    """ The base interface for the multimodal data """
+    """ The base interface for multimodal data. """
+
     def __init__(self):
         self.phase = None  # the 'train' , 'test', 'val'
 
-        # the recorded samples for current dataset
+        # The recorded samples for current dataset:
         #   In flickr20K entities dataset, this presents as:
         #    this is a dict in which key is the 'sample name/id' ...
         #    the values are the sample's information,
@@ -275,7 +271,7 @@ def __init__(self):
         #    {"rgb": rgb_dataset, "flow": flow_dataset, "audio": audio_dataset}
         self.phase_multimodal_data_record = None
 
-        # the detailed info in selected split
+        # Detailed information in selected split:
         #   i.e., path, path for different modalities, et. al
         self.phase_info = None
         # the data types included,
@@ -326,7 +322,7 @@ def __getitem__(self, sample_idx):
         sampled_multimodal_data = self.get_one_multimodal_sample(sample_idx)
 
         # utilize the modality to mask specific modalities
-        sampled_modality_data = dict()
+        sampled_modality_data = {}
         for item_name, item_data in sampled_multimodal_data.items():
             # maintain the modality data based on the sampler
             # maintain the external data
diff --git a/plato/datasources/referitgame.py b/plato/datasources/referitgame.py
index e3ddb12ad..10ad934ae 100644
--- a/plato/datasources/referitgame.py
+++ b/plato/datasources/referitgame.py
@@ -167,18 +167,17 @@ def __init__(self):
         #  unc, google
         self.split_name = Config().data.split_name
         if self.split_config not in self.split_configs:
-            info_msg = (
-                "{} does not exsit in the official configs {}.....").format(
-                    self.split_config, self.split_configs)
-            logging.info(info_msg)
+            logging.info(
+                "%s does not exist in the official configurations %s.",
+                self.split_config, self.split_configs)
 
         _path = Config().params['data_path']
         self._data_path_process(data_path=_path, base_data_name=self.data_name)
-        base_data_path = self.mm_data_info["base_data_dir_path"]
+        base_data_path = self.mm_data_info["data_path"]
 
         # raw coco images path
         coco_raw_imgs_path = self.base_coco
-        if self._exist_judgement(coco_raw_imgs_path):
+        if self._exists(coco_raw_imgs_path):
             logging.info(
                 "Successfully connecting the source COCO2017 images data from the path %s",
                 coco_raw_imgs_path)
@@ -192,7 +191,7 @@ def __init__(self):
         ).data.download_splits_base_url + self.split_config + ".zip"
         for dd_url in [download_split_url]:
             self._download_arrange_data(download_url_address=dd_url,
-                                        put_data_dir=base_data_path)
+                                        data_path=base_data_path)
 
         self._dataset_refer = referitgame_utils.REFER(
             data_root=base_data_path,
@@ -200,7 +199,7 @@ def __init__(self):
             dataset=self.split_config,
             splitBy=self.split_name)  # default is unc or google
 
-        self._splited_referids_holder = dict()
+        self._splited_referids_holder = {}
         self._connect_to_splits()
 
     def _connect_to_splits(self):
@@ -215,8 +214,8 @@ def get_phase_data(self, phase):
         """ Get phrases from the raw data """
         mode_refer_ids = self._splited_referids_holder[phase]
 
-        mode_elements_holder = dict()
-        mode_flatten_emelemts = list()
+        mode_elements_holder = {}
+        mode_flatten_emelemts = []
 
         for refer_id in mode_refer_ids:
 
@@ -226,11 +225,11 @@ def get_phase_data(self, phase):
             caption_phrases_cate = self._dataset_refer.Cats[ref['category_id']]
             caption_phrases_cate_id = ref['category_id']
 
-            mode_elements_holder[refer_id] = dict()
+            mode_elements_holder[refer_id] = {}
             mode_elements_holder[refer_id]["image_id"] = image_id
             mode_elements_holder[refer_id]["image_file_path"] = image_file_path
 
-            mode_elements_holder[refer_id]["sentences"] = list()
+            mode_elements_holder[refer_id]["sentences"] = []
             for send in ref["sentences"]:
                 caption = send["tokens"]
                 caption_phrase = send["tokens"]
diff --git a/plato/models/dcgan.py b/plato/models/dcgan.py
new file mode 100644
index 000000000..9b4334198
--- /dev/null
+++ b/plato/models/dcgan.py
@@ -0,0 +1,120 @@
+"""
+The DCGAN model.
+
+Reference:
+https://pytorch.org/tutorials/beginner/dcgan_faces_tutorial.html
+"""
+
+from torch import nn
+
+nz = 100
+nc = 3
+ngf = 64
+ndf = 64
+
+
+class Generator(nn.Module):
+    """ Generator network of DCGAN """
+
+    def __init__(self):
+        super().__init__()
+
+        self.main = nn.Sequential(
+            # input is Z, going into a convolution
+            nn.ConvTranspose2d(nz, ngf * 8, 4, 1, 0, bias=False),
+            nn.BatchNorm2d(ngf * 8),
+            nn.ReLU(True),
+            # state size. (ngf*8) x 4 x 4
+            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ngf * 4),
+            nn.ReLU(True),
+            # state size. (ngf*4) x 8 x 8
+            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ngf * 2),
+            nn.ReLU(True),
+            # state size. (ngf*2) x 16 x 16
+            nn.ConvTranspose2d(ngf * 2, ngf, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ngf),
+            nn.ReLU(True),
+            # state size. (ngf) x 32 x 32
+            nn.ConvTranspose2d(ngf, nc, 4, 2, 1, bias=False),
+            nn.Tanh()
+            # state size. (nc) x 64 x 64
+        )
+
+    def forward(self, input_data):
+        """ Forward pass. """
+        return self.main(input_data)
+
+
+class Discriminator(nn.Module):
+    """ Discriminator network of DCGAN. """
+
+    def __init__(self):
+        super().__init__()
+
+        self.main = nn.Sequential(
+            # input is (nc) x 64 x 64
+            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
+            nn.LeakyReLU(0.2, inplace=True),
+            # state size. (ndf) x 32 x 32
+            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ndf * 2),
+            nn.LeakyReLU(0.2, inplace=True),
+            # state size. (ndf*2) x 16 x 16
+            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ndf * 4),
+            nn.LeakyReLU(0.2, inplace=True),
+            # state size. (ndf*4) x 8 x 8
+            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ndf * 8),
+            nn.LeakyReLU(0.2, inplace=True),
+            # state size. (ndf*8) x 4 x 4
+            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),
+            nn.Sigmoid())
+
+    def forward(self, input_data):
+        return self.main(input_data)
+
+
+class Model:
+    """ A wrapper class to hold the Generator and Discriminator models of DCGAN. """
+
+    def __init__(self) -> None:
+        self.generator = Generator()
+        self.discriminator = Discriminator()
+        self.loss_criterion = nn.BCELoss()
+
+        self.nz = nz
+        self.nc = nc
+        self.ngf = ngf
+        self.ndf = ndf
+
+    def weights_init(self, model):
+        classname = model.__class__.__name__
+        if classname.find('Conv') != -1:
+            nn.init.normal_(model.weight.data, 0.0, 0.02)
+        elif classname.find('BatchNorm') != -1:
+            nn.init.normal_(model.weight.data, 1.0, 0.02)
+            nn.init.constant_(model.bias.data, 0)
+
+    def cpu(self):
+        self.generator.cpu()
+        self.discriminator.cpu()
+
+    def to(self, device):
+        self.generator.to(device)
+        self.discriminator.to(device)
+
+    def train(self):
+        self.generator.train()
+        self.discriminator.train()
+
+    def eval(self):
+        self.generator.eval()
+        self.discriminator.eval()
+
+    @staticmethod
+    def get_model(*args):
+        """ Obtaining an instance of this model. """
+        return Model()
diff --git a/plato/models/registry.py b/plato/models/registry.py
index fc58cce99..690c142bd 100644
--- a/plato/models/registry.py
+++ b/plato/models/registry.py
@@ -37,6 +37,7 @@
         hybrid,
         efficientnet,
         regnet,
+        dcgan
     )
     registered_models = OrderedDict([
         ('lenet5', lenet5.Model),
@@ -52,6 +53,7 @@
         ('hybrid', hybrid.Model),
         ('efficientnet', efficientnet.Model),
         ('regnet', regnet.Model),
+        ('dcgan', dcgan.Model),
     ])
 
 
diff --git a/plato/samplers/sampler_utils.py b/plato/samplers/sampler_utils.py
index d5b7fb1c2..a04f35f88 100644
--- a/plato/samplers/sampler_utils.py
+++ b/plato/samplers/sampler_utils.py
@@ -1,14 +1,9 @@
 """
-    Rewrite the Opacus's BatchMemoryManager
-    Reference: https://opacus.ai/api/_modules/opacus/utils/batch_memory_manager.html
+Utility functions that samples local datasets
 """
 
 import numpy as np
 
-from opacus.optimizers import DPOptimizer
-from opacus.utils.batch_memory_manager import BatchSplittingSampler
-from torch.utils.data import DataLoader, BatchSampler
-
 
 def extend_indices(indices, required_total_size):
     """ Extend the indices to obtain the required total size
@@ -192,66 +187,3 @@ def set_min_bound(proportion):
             np.repeat(concentration, number_partitions))
 
     return proportions
-
-
-def wrap_data_loader(*, data_loader: DataLoader,
-                     original_batch_sampler: BatchSampler, max_batch_size: int,
-                     optimizer: DPOptimizer):
-    """
-    Replaces batch_sampler in the input data loader with ``BatchSplittingSampler``
-
-    Args:
-        data_loader: Wrapper DataLoader
-        max_batch_size: max physical batch size we want to emit
-        optimizer: DPOptimizer instance used for training
-
-    Returns:
-        New DataLoader instance with batch_sampler wrapped in ``BatchSplittingSampler``
-    """
-    return DataLoader(
-        dataset=data_loader.dataset,
-        batch_sampler=BatchSplittingSampler(
-            sampler=original_batch_sampler,
-            max_batch_size=max_batch_size,
-            optimizer=optimizer,
-        ),
-        num_workers=data_loader.num_workers,
-        collate_fn=data_loader.collate_fn,
-        pin_memory=data_loader.pin_memory,
-        timeout=data_loader.timeout,
-        worker_init_fn=data_loader.worker_init_fn,
-        multiprocessing_context=data_loader.multiprocessing_context,
-        generator=data_loader.generator,
-        prefetch_factor=data_loader.prefetch_factor,
-        persistent_workers=data_loader.persistent_workers,
-    )
-
-
-class BatchMemoryManager(object):
-    """
-        Enable sampler on top of Opacus BatchMemoryManager
-    """
-
-    def __init__(
-        self,
-        *,
-        data_loader: DataLoader,
-        original_batch_sampler: BatchSampler,
-        max_physical_batch_size: int,
-        optimizer: DPOptimizer,
-    ):
-        self.data_loader = data_loader
-        self.optimizer = optimizer
-        self.original_batch_sampler = original_batch_sampler
-        self.max_physical_batch_size = max_physical_batch_size
-
-    def __enter__(self):
-        return wrap_data_loader(
-            data_loader=self.data_loader,
-            original_batch_sampler=self.original_batch_sampler,
-            max_batch_size=self.max_physical_batch_size,
-            optimizer=self.optimizer,
-        )
-
-    def __exit__(self, type, value, traceback):
-        pass
diff --git a/plato/servers/base.py b/plato/servers/base.py
index cff110c05..53a15cec8 100644
--- a/plato/servers/base.py
+++ b/plato/servers/base.py
@@ -95,8 +95,14 @@ def __init__(self):
         self.comm_simulation = Config().clients.comm_simulation if hasattr(
             Config().clients, 'comm_simulation') else True
 
+        # Accumulated communication overhead (MB) throughout the FL training session
+        self.comm_overhead = 0
+
         # States that need to be maintained for asynchronous FL
 
+        # sids that are currently in use
+        self.training_sids = []
+
         # Clients whose new reports were received but not yet processed
         self.reported_clients = []
 
@@ -381,7 +387,7 @@ async def select_clients(self, for_next_batch=False):
                 # reported may not have been aggregated; they should be excluded from the next
                 # round of client selection
                 reporting_client_ids = [
-                    client[1]['client_id'] for client in self.reported_clients
+                    client[2]['client_id'] for client in self.reported_clients
                 ]
 
                 selectable_clients = [
@@ -443,6 +449,7 @@ async def select_clients(self, for_next_batch=False):
                             max_concurrency, len(self.selected_clients))]
 
                 self.trained_clients += selected_clients
+
             else:
                 selected_clients = self.selected_clients
 
@@ -459,16 +466,14 @@ async def select_clients(self, for_next_batch=False):
                 sid = self.clients[client_id]['sid']
 
                 if self.asynchronous_mode and self.simulate_wall_time:
-                    training_sids = []
-                    for client_info in self.reported_clients:
-                        training_sids.append(client_info[1]['sid'])
 
                     # skip if this sid is currently `training' with reporting clients
                     # or it has already been selected in this round
-                    while sid in training_sids or sid in self.selected_sids:
+                    while sid in self.training_sids or sid in self.selected_sids:
                         client_id = client_id % self.clients_per_round + 1
                         sid = self.clients[client_id]['sid']
 
+                    self.training_sids.append(sid)
                     self.selected_sids.append(sid)
 
                 self.training_clients[self.selected_client_id] = {
@@ -501,10 +506,21 @@ async def select_clients(self, for_next_batch=False):
                     checkpoint_path = Config().params['checkpoint_path']
 
                     payload_filename = f"{checkpoint_path}/{model_name}_{self.selected_client_id}.pth"
+
                     with open(payload_filename, 'wb') as payload_file:
                         pickle.dump(payload, payload_file)
+
                     server_response['payload_filename'] = payload_filename
 
+                    payload_size = sys.getsizeof(
+                        pickle.dumps(payload)) / 1024**2
+
+                    logging.info(
+                        "[%s] Sending %.2f MB of payload data to client #%d (simulated).",
+                        self, payload_size, self.selected_client_id)
+
+                    self.comm_overhead += payload_size
+
                 server_response = await self.customize_server_response(
                     server_response)
 
@@ -621,6 +637,8 @@ async def send(self, sid, payload, client_id) -> None:
         logging.info("[%s] Sent %.2f MB of payload data to client #%d.", self,
                      data_size / 1024**2, client_id)
 
+        self.comm_overhead += data_size / 1024**2
+
     async def client_report_arrived(self, sid, client_id, report):
         """ Upon receiving a report from a client. """
         self.reports[sid] = pickle.loads(report)
@@ -635,11 +653,14 @@ async def client_report_arrived(self, sid, client_id, report):
             with open(payload_filename, 'rb') as payload_file:
                 self.client_payload[sid] = pickle.load(payload_file)
 
+            payload_size = sys.getsizeof(pickle.dumps(
+                self.client_payload[sid])) / 1024**2
+
             logging.info(
                 "[%s] Received %.2f MB of payload data from client #%d (simulated).",
-                self,
-                sys.getsizeof(pickle.dumps(self.client_payload[sid])) /
-                1024**2, client_id)
+                self, payload_size, client_id)
+
+            self.comm_overhead += payload_size
 
             await self.process_client_info(client_id, sid)
 
@@ -684,6 +705,8 @@ async def client_payload_done(self, sid, client_id, s3_key=None):
         logging.info("[%s] Received %.2f MB of payload data from client #%d.",
                      self, payload_size / 1024**2, client_id)
 
+        self.comm_overhead += payload_size / 1024**2
+
         await self.process_client_info(client_id, sid)
 
     async def process_client_info(self, client_id, sid):
@@ -703,8 +726,13 @@ async def process_client_info(self, client_id, sid):
             sid].comm_time + start_time
         starting_round = self.training_clients[client_id]['starting_round']
 
+        if Config().is_central_server():
+            self.comm_overhead += self.reports[sid].edge_server_comm_overhead
+
         client_info = (
             finish_time,  # sorted by the client's finish time
+            random.random(
+            ),  # in case two or more clients have the same finish time
             {
                 'client_id': client_id,
                 'sid': sid,
@@ -715,9 +743,12 @@ async def process_client_info(self, client_id, sid):
             })
 
         heapq.heappush(self.reported_clients, client_info)
-        self.current_reported_clients[client_info[1]['client_id']] = True
+        self.current_reported_clients[client_info[2]['client_id']] = True
         del self.training_clients[client_id]
 
+        if self.asynchronous_mode and self.simulate_wall_time:
+            self.training_sids.remove(client_info[2]['sid'])
+
         await self.process_clients(client_info)
 
     async def process_clients(self, client_info):
@@ -744,7 +775,7 @@ async def process_clients(self, client_info):
 
                 request_sent = False
                 for i, client_info in enumerate(self.reported_clients):
-                    client = client_info[1]
+                    client = client_info[2]
                     client_staleness = self.current_round - client[
                         'starting_round']
 
@@ -790,7 +821,7 @@ async def process_clients(self, client_info):
                         self.minimum_clients)):
                 # Extract a client with the earliest finish time in wall clock time
                 client_info = heapq.heappop(self.reported_clients)
-                client = client_info[1]
+                client = client_info[2]
 
                 # Removing from the list of current reporting clients as well, if needed
                 self.current_processed_clients[client['client_id']] = True
@@ -823,14 +854,14 @@ async def process_clients(self, client_info):
                 client_info = heapq.heappop(self.reported_clients)
                 heapq.heappush(possibly_stale_clients, client_info)
 
-                if client_info[1][
+                if client_info[2][
                         'starting_round'] < self.current_round - self.staleness_bound:
                     for __ in range(0, len(possibly_stale_clients)):
                         stale_client_info = heapq.heappop(
                             possibly_stale_clients)
                         # Update the simulated wall clock time to be the finish time of this client
                         self.wall_time = stale_client_info[0]
-                        client = stale_client_info[1]
+                        client = stale_client_info[2]
 
                         # Add the report and payload of the extracted reporting client into updates
                         logging.info(
@@ -856,7 +887,7 @@ async def process_clients(self, client_info):
             # In both synchronous and asynchronous modes, if we are not simulating the wall clock
             # time, we need to add the client report to the list of updates so far;
             # the same applies when we are running in synchronous mode.
-            client = client_info[1]
+            client = client_info[2]
             client_staleness = self.current_round - client['starting_round']
 
             self.updates.append((client['client_id'], client['report'],
@@ -892,7 +923,9 @@ async def process_clients(self, client_info):
                      'max_concurrency') and not Config().is_central_server():
             # Clients in the current batch finish training
             # The server will select the next batch of clients to train
-            if len(self.updates) >= len(self.trained_clients):
+            if len(self.updates) >= len(self.trained_clients) or len(
+                    self.current_reported_clients) >= len(
+                        self.trained_clients):
                 await self.select_clients(for_next_batch=True)
 
     async def client_disconnected(self, sid):
diff --git a/plato/servers/fedavg.py b/plato/servers/fedavg.py
index 3c4f0d970..829cc9b72 100644
--- a/plato/servers/fedavg.py
+++ b/plato/servers/fedavg.py
@@ -232,6 +232,8 @@ def get_record_items_values(self):
                 report.training_time + report.comm_time
                 for (__, report, __, __) in self.updates
             ]),
+            'comm_overhead':
+            self.comm_overhead,
         }
 
     @staticmethod
diff --git a/plato/servers/fedavg_gan.py b/plato/servers/fedavg_gan.py
new file mode 100644
index 000000000..ec25954a8
--- /dev/null
+++ b/plato/servers/fedavg_gan.py
@@ -0,0 +1,80 @@
+"""
+A federated learning server using federated averaging to train GAN models.
+"""
+import asyncio
+
+from plato.servers import fedavg
+from plato.config import Config
+
+
+class Server(fedavg.Server):
+    """ Federated learning server using federated averaging to train GAN models. """
+
+    async def federated_averaging(self, updates):
+        """Aggregate weight updates from the clients using federated averaging."""
+        weights_received = self.compute_weight_deltas(updates)
+
+        # Total sample is the same for both Generator and Discriminator
+        self.total_samples = sum(
+            [report.num_samples for (__, report, __, __) in updates])
+
+        # Perform weighted averaging for both Generator and Discriminator
+        gen_avg_update = {
+            name: self.trainer.zeros(weights.shape)
+            for name, weights in weights_received[0][0].items()
+        }
+        disc_avg_update = {
+            name: self.trainer.zeros(weights.shape)
+            for name, weights in weights_received[0][1].items()
+        }
+
+        for i, update in enumerate(weights_received):
+            __, report, __, __ = updates[i]
+            num_samples = report.num_samples
+
+            update_from_gen, update_from_disc = update
+
+            for name, delta in update_from_gen.items():
+                gen_avg_update[name] += delta * (num_samples /
+                                                 self.total_samples)
+
+            for name, delta in update_from_disc.items():
+                disc_avg_update[name] += delta * (num_samples /
+                                                  self.total_samples)
+
+            # Yield to other tasks in the server
+            await asyncio.sleep(0)
+
+        return gen_avg_update, disc_avg_update
+
+    def customize_server_payload(self, payload):
+        """
+        Customize the server payload before sending to the client.
+
+        At the end of each round, the server can choose to only send the global Generator
+        or Discriminator (or both or neither) model to the clients next round.
+
+        Reference this paper for more detail:
+        https://deepai.org/publication/federated-generative-adversarial-learning
+
+        By default, both model will be sent to the clients.
+        """
+        if hasattr(Config().server, 'network_to_sync'):
+            network = Config().server.network_to_sync.lower()
+        else:
+            network = 'both'
+
+        weights_gen, weights_disc = payload
+        if network == 'none':
+            server_payload = None, None
+        elif network == 'generator':
+            server_payload = weights_gen, None
+        elif network == 'discriminator':
+            server_payload = None, weights_disc
+        elif network == 'both':
+            server_payload = payload
+        else:
+            raise ValueError(
+                f'Unknown value to attribute network_to_sync: {network}')
+
+        return server_payload
diff --git a/plato/servers/registry.py b/plato/servers/registry.py
index 4141d1da4..d02f883b5 100644
--- a/plato/servers/registry.py
+++ b/plato/servers/registry.py
@@ -13,12 +13,14 @@
     fedavg,
     fedavg_cs,
     mistnet,
+    fedavg_gan,
 )
 
 registered_servers = OrderedDict([
     ('fedavg', fedavg.Server),
     ('fedavg_cross_silo', fedavg_cs.Server),
     ('mistnet', mistnet.Server),
+    ('fedavg_gan', fedavg_gan.Server),
 ])
 
 
diff --git a/plato/trainers/base.py b/plato/trainers/base.py
index 66a2e3038..a9da6eb25 100644
--- a/plato/trainers/base.py
+++ b/plato/trainers/base.py
@@ -43,7 +43,7 @@ def save_accuracy(accuracy, filename=None):
         else:
             accuracy_path = f'{model_path}/{model_name}.acc'
 
-        with open(accuracy_path, 'w', encoding='utf8') as file:
+        with open(accuracy_path, 'w', encoding='utf-8') as file:
             file.write(str(accuracy))
 
     @staticmethod
@@ -57,7 +57,7 @@ def load_accuracy(filename=None):
         else:
             accuracy_path = f'{model_path}/{model_name}.acc'
 
-        with open(accuracy_path, 'r', encoding='utf8') as file:
+        with open(accuracy_path, 'r', encoding='utf-8') as file:
             accuracy = float(file.read())
 
         return accuracy
diff --git a/plato/trainers/basic.py b/plato/trainers/basic.py
index db1948772..f69448477 100644
--- a/plato/trainers/basic.py
+++ b/plato/trainers/basic.py
@@ -30,6 +30,7 @@ def __init__(self, model=None):
 
         self.training_start_time = time.time()
         self.models_per_epoch = {}
+        self.model_state_dict = None
 
         if model is None:
             model = models_registry.get()
@@ -59,7 +60,10 @@ def save_model(self, filename=None, location=None):
         else:
             model_path = f'{model_path}/{model_name}.pth'
 
-        torch.save(self.model.state_dict(), model_path)
+        if self.model_state_dict is None:
+            torch.save(self.model.state_dict(), model_path)
+        else:
+            torch.save(self.model_state_dict, model_path)
 
         if self.client_id == 0:
             logging.info("[Server #%d] Model saved to %s.", os.getpid(),
@@ -86,7 +90,7 @@ def load_model(self, filename=None, location=None):
             logging.info("[Client #%d] Loading a model from %s.",
                          self.client_id, model_path)
 
-        self.model.load_state_dict(torch.load(model_path), strict=False)
+        self.model.load_state_dict(torch.load(model_path), strict=True)
 
     def simulate_sleep_time(self):
         """Simulate client's speed by putting it to sleep."""
@@ -127,13 +131,7 @@ def train_process(self,
                 self.train_model(config, trainset, sampler.get(), cut_layer,
                                  **kwargs)
             else:
-                self.train_loop(
-                    config,
-                    trainset,
-                    sampler.get(),
-                    cut_layer,
-                    **kwargs,
-                )
+                self.train_loop(config, trainset, sampler.get(), cut_layer)
         except Exception as training_exception:
             logging.info("Training on client #%d failed.", self.client_id)
             raise training_exception
@@ -144,7 +142,7 @@ def train_process(self,
             filename = f"{model_type}_{self.client_id}_{config['run_id']}.pth"
             self.save_model(filename)
 
-    def train_loop(self, config, trainset, sampler, cut_layer, **kwargs):
+    def train_loop(self, config, trainset, sampler, cut_layer):
         """ The default training loop when a custom training loop is not supplied. """
         batch_size = config['batch_size']
         log_interval = 10
@@ -162,7 +160,6 @@ def train_loop(self, config, trainset, sampler, cut_layer, **kwargs):
                                                        batch_size=batch_size,
                                                        sampler=sampler)
 
-        iterations_per_epoch = np.ceil(len(trainset) / batch_size).astype(int)
         epochs = config['epochs']
 
         # Initializing the loss criterion
@@ -178,9 +175,9 @@ def train_loop(self, config, trainset, sampler, cut_layer, **kwargs):
         optimizer = get_optimizer(self.model)
 
         # Initializing the learning rate schedule, if necessary
-        if hasattr(config, 'lr_schedule'):
+        if 'lr_schedule' in config:
             lr_schedule = optimizers.get_lr_schedule(optimizer,
-                                                     iterations_per_epoch,
+                                                     len(train_loader),
                                                      train_loader)
         else:
             lr_schedule = None
@@ -269,7 +266,8 @@ def train(self, trainset, sampler, cut_layer=None, **kwargs) -> float:
                 mp.set_start_method('spawn', force=True)
 
             train_proc = mp.Process(target=self.train_process,
-                                    args=(config, trainset, sampler, cut_layer),
+                                    args=(config, trainset, sampler,
+                                          cut_layer),
                                     kwargs=kwargs)
             train_proc.start()
             train_proc.join()
diff --git a/plato/trainers/diff_privacy.py b/plato/trainers/diff_privacy.py
index 936bbd027..23817e469 100644
--- a/plato/trainers/diff_privacy.py
+++ b/plato/trainers/diff_privacy.py
@@ -7,15 +7,16 @@
 
 import numpy as np
 import torch
+from torch.utils.data import Subset
 
 from opacus import GradSampleModule
 from opacus.privacy_engine import PrivacyEngine
 from opacus.validators import ModuleValidator
+from opacus.utils.batch_memory_manager import BatchMemoryManager
 
 from plato.config import Config
 from plato.trainers import basic
 from plato.utils import optimizers
-from plato.samplers.sampler_utils import BatchMemoryManager
 
 
 class Trainer(basic.Trainer):
@@ -25,10 +26,13 @@ def __init__(self, model=None):
         """Initializing the trainer with the provided model.
         """
         super().__init__(model=model)
+
         self.max_physical_batch_size = Config(
         ).trainer.max_physical_batch_size if hasattr(
             Config().trainer, "max_physical_batch_size") else 128
 
+        self.make_model_private()
+
     def make_model_private(self):
         """ Make the model private for use with the differential privacy engine. """
         errors = ModuleValidator.validate(self.model, strict=False)
@@ -37,30 +41,26 @@ def make_model_private(self):
             errors = ModuleValidator.validate(self.model, strict=False)
             assert len(errors) == 0
 
-        self.model = GradSampleModule(self.model)
-
     def train_model(self, config, trainset, sampler, cut_layer):
         """ The training loop that supports differential privacy. """
         batch_size = config['batch_size']
         log_interval = 10
         tic = time.perf_counter()
 
-        self.make_model_private()
+        self.model = GradSampleModule(self.model)
 
-        logging.info("[Client #%d] Loading the dataset.", self.client_id)
+        logging.info("[Client #%d] Loading the dataset with size %d.",
+                     self.client_id, len(list(sampler)))
         _train_loader = getattr(self, "train_loader", None)
+        trainset = Subset(trainset, list(sampler))
 
         if callable(_train_loader):
-            train_loader = self.train_loader(batch_size, trainset, sampler,
-                                             cut_layer)
+            train_loader = self.train_loader(batch_size, trainset, cut_layer)
         else:
             train_loader = torch.utils.data.DataLoader(dataset=trainset,
                                                        shuffle=False,
-                                                       batch_size=batch_size,
-                                                       sampler=sampler)
+                                                       batch_size=batch_size)
 
-        iterations_per_epoch = np.ceil(len(trainset) / batch_size).astype(int)
-        batch_sampler = train_loader.batch_sampler
         epochs = config['epochs']
 
         # Initializing the loss criterion
@@ -76,9 +76,9 @@ def train_model(self, config, trainset, sampler, cut_layer):
         optimizer = get_optimizer(self.model)
 
         # Initializing the learning rate schedule, if necessary
-        if hasattr(config, 'lr_schedule'):
+        if 'lr_schedule' in config:
             lr_schedule = optimizers.get_lr_schedule(optimizer,
-                                                     iterations_per_epoch,
+                                                     len(train_loader),
                                                      train_loader)
         else:
             lr_schedule = None
@@ -108,7 +108,6 @@ def train_model(self, config, trainset, sampler, cut_layer):
         for epoch in range(1, epochs + 1):
             with BatchMemoryManager(
                     data_loader=train_loader,
-                    original_batch_sampler=batch_sampler,
                     max_physical_batch_size=self.max_physical_batch_size,
                     optimizer=optimizer) as memory_safe_train_loader:
                 for batch_id, (examples,
@@ -165,3 +164,10 @@ def train_model(self, config, trainset, sampler, cut_layer):
                 filename = f"{self.client_id}_{epoch}_{training_time}.pth"
                 self.save_model(filename)
                 self.model.to(self.device)
+
+        # After GradSampleModule() conversion, the state_dict names have a `_module` prefix
+        # We will need to save the weights with the original layer names without the prefix
+        self.model_state_dict = {
+            k[8:] if '_module.' in k else k: v
+            for k, v in self.model.state_dict().items()
+        }
diff --git a/plato/trainers/gan.py b/plato/trainers/gan.py
new file mode 100644
index 000000000..56cb2021e
--- /dev/null
+++ b/plato/trainers/gan.py
@@ -0,0 +1,303 @@
+"""
+The training and testing loops for GAN models.
+
+Reference:
+https://pytorch.org/tutorials/beginner/dcgan_faces_tutorial.html
+"""
+import logging
+import math
+import os
+
+import torch
+import torch.nn as nn
+import torchvision
+import numpy as np
+import scipy
+
+from plato.config import Config
+from plato.models import registry as models_registry
+from plato.trainers import basic
+from plato.utils import optimizers
+
+
+class Trainer(basic.Trainer):
+    """A federated learning trainer for GAN models."""
+
+    def __init__(self, model=None):
+        super().__init__()
+
+        if model is None:
+            model = models_registry.get()
+        gan_model = model
+        self.generator = gan_model.generator
+        self.discriminator = gan_model.discriminator
+        self.loss_criterion = gan_model.loss_criterion
+        self.model = gan_model
+
+        # Use the pre-trained InceptionV3 model as a feature extractor
+        # for testing
+        self.inception_model = torchvision.models.inception_v3(
+            pretrained=True, aux_logits=False)
+        # Remove the last output layer of inception
+        self.inception_model.fc = nn.Identity()
+        self.inception_model.eval()
+
+        self.training_start_time = 0
+
+    def save_model(self, filename=None, location=None):
+        """Saving the model to a file. """
+        model_path = Config(
+        ).params['model_path'] if location is None else location
+        model_name = Config().trainer.model_name
+
+        try:
+            if not os.path.exists(model_path):
+                os.makedirs(model_path)
+        except FileExistsError:
+            pass
+
+        if filename is not None:
+            net_gen_path = f'{model_path}/Generator_{filename}'
+            net_disc_path = f'{model_path}/Discriminator_{filename}'
+        else:
+            net_gen_path = f'{model_path}/Generator_{model_name}.pth'
+            net_disc_path = f'{model_path}/Discriminator_{model_name}.pth'
+
+        torch.save(self.generator.state_dict(), net_gen_path)
+        torch.save(self.discriminator.state_dict(), net_disc_path)
+
+        if self.client_id == 0:
+            logging.info("[Server #%d] Generator Model saved to %s.",
+                         os.getpid(), net_gen_path)
+            logging.info("[Server #%d] Discriminator Model saved to %s.",
+                         os.getpid(), net_disc_path)
+        else:
+            logging.info("[Client #%d] Generator Model saved to %s.",
+                         self.client_id, net_gen_path)
+            logging.info("[Client #%d] Discriminator Model saved to %s.",
+                         self.client_id, net_disc_path)
+
+    def load_model(self, filename=None, location=None):
+        """Loading pre-trained model weights from a file. """
+        model_path = Config(
+        ).params['model_path'] if location is None else location
+        model_name = Config().trainer.model_name
+
+        if filename is not None:
+            net_gen_path = f'{model_path}/Generator_{filename}'
+            net_disc_path = f'{model_path}/Discriminator_{filename}'
+        else:
+            net_gen_path = f'{model_path}/Generator_{model_name}.pth'
+            net_disc_path = f'{model_path}/Discriminator_{model_name}.pth'
+
+        if self.client_id == 0:
+            logging.info("[Server #%d] Loading a Generator model from %s.",
+                         os.getpid(), net_gen_path)
+            logging.info("[Server #%d] Loading a Discriminator model from %s.",
+                         os.getpid(), net_disc_path)
+        else:
+            logging.info("[Client #%d] Loading a Generator model from %s.",
+                         self.client_id, net_gen_path)
+            logging.info("[Client #%d] Loading a Discriminator model from %s.",
+                         self.client_id, net_disc_path)
+
+        self.generator.load_state_dict(torch.load(net_gen_path))
+        self.discriminator.load_state_dict(torch.load(net_disc_path))
+
+    def train_loop(self, config, trainset, sampler, cut_layer=None) -> float:
+        """The main training loop in a federated learning workload.
+
+        Arguments:
+        trainset: The training dataset.
+        sampler: the sampler that extracts a partition for this client.
+        cut_layer (optional): The layer which training should start from.
+
+        Returns:
+        float: The training time.
+        """
+        batch_size = config['batch_size']
+        log_interval = 10
+
+        logging.info("[Client #%d] Loading the dataset.", self.client_id)
+
+        train_loader = torch.utils.data.DataLoader(dataset=trainset,
+                                                   shuffle=False,
+                                                   batch_size=batch_size,
+                                                   sampler=sampler)
+
+        self.model.to(self.device)
+        self.model.train()
+
+        # self.generator.apply(self.model.weights_init)
+        # self.discriminator.apply(self.model.weights_init)
+
+        optimizer_gen = optimizers.get_optimizer(self.generator)
+        optimizer_disc = optimizers.get_optimizer(self.discriminator)
+
+        real_label = 1.
+        fake_label = 0.
+
+        epochs = config['epochs']
+        for epoch in range(1, epochs + 1):
+            # Here we assume the data samples still have labels attached to them,
+            # but GAN training does not need labels, so we'll just discard them
+            for batch_id, (examples, _) in enumerate(train_loader):
+                cur_batch_size = len(examples)
+                examples = examples.to(self.device)
+                label = torch.full((cur_batch_size, ),
+                                   real_label,
+                                   dtype=torch.float)
+                label = label.to(self.device)
+                ############################
+                # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
+                ###########################
+                ## Train with all-real batch
+                optimizer_disc.zero_grad()
+                # Forward pass real batch through D
+                output = self.discriminator(examples).view(-1)
+                # Calculate loss on all-real batch
+                err_disc_real = self.loss_criterion(output, label)
+                # Calculate gradients for D in backward pass
+                err_disc_real.backward()
+
+                ## Train with all-fake batch
+                # Generate batch of latent vectors
+                noise = torch.randn(cur_batch_size,
+                                    self.model.nz,
+                                    1,
+                                    1,
+                                    device=self.device)
+                # Generate fake image batch with G
+                fake = self.generator(noise)
+                label.fill_(fake_label)
+                # Classify all fake batch with D
+                output = self.discriminator(fake.detach()).view(-1)
+                # Calculate D's loss on the all-fake batch
+                err_disc_fake = self.loss_criterion(output, label)
+                # Calculate the gradients for this batch, accumulated (summed)
+                # with previous gradients
+                err_disc_fake.backward()
+                # Compute error of D as sum over the fake and the real batches
+                err_disc_total = err_disc_real + err_disc_fake
+                # Update D
+                optimizer_disc.step()
+
+                ############################
+                # (2) Update G network: maximize log(D(G(z)))
+                ###########################
+                optimizer_gen.zero_grad()
+                label.fill_(
+                    real_label)  # fake labels are real for generator cost
+                # Since we just updated D, perform another forward pass of all-fake batch through D
+                output = self.discriminator(fake).view(-1)
+                # Calculate G's loss based on this output
+                err_gen = self.loss_criterion(output, label)
+                # Calculate gradients for G
+                err_gen.backward()
+                # Update G
+                optimizer_gen.step()
+
+                if batch_id % log_interval == 0:
+                    if self.client_id == 0:
+                        logging.info(
+                            "[Server #%d] Epoch: [%d/%d][%d/%d]\tGenerator Loss: %.6f\t"
+                            "Discriminator Loss: %.6f", os.getpid(), epoch,
+                            epochs, batch_id, len(train_loader),
+                            err_gen.data.item(), err_disc_total.data.item())
+                    else:
+                        logging.info(
+                            "[Client #%d] Epoch: [%d/%d][%d/%d]\tGenerator Loss: %.6f\t"
+                            "Discriminator Loss: %.6f", self.client_id, epoch,
+                            epochs, batch_id, len(train_loader),
+                            err_gen.data.item(), err_disc_total.data.item())
+
+    def test_model(self, config, testset):
+        """Test the Generator model with the Frechet Inception Distance metric."""
+
+        self.model.to(self.device)
+        self.model.eval()
+
+        perplexity = -1
+
+        test_loader = torch.utils.data.DataLoader(
+            testset, batch_size=config['batch_size'], shuffle=True)
+
+        real_features, fake_features = [], []
+        with torch.no_grad():
+            for real_examples, _ in test_loader:
+                real_examples = real_examples.to(self.device)
+
+                noise = torch.randn(config['batch_size'],
+                                    self.model.nz,
+                                    1,
+                                    1,
+                                    device=self.device)
+                fake_examples = self.generator(noise)
+
+                # Extract the feature of real and synthetic data with
+                # InceptionV3 model pre-trained on ImageNet
+                self.inception_model.to(self.device)
+                feature_real = self.feature_extractor(real_examples)
+                feature_fake = self.feature_extractor(fake_examples)
+
+                # Store the feature of every real and synthetic data
+                real_features.extend(list(feature_real))
+                fake_features.extend(list(feature_fake))
+
+            real_features, fake_features = np.stack(real_features), np.stack(
+                fake_features)
+            # Calculate the Frechet Distance between the feature distribution
+            # of real data from testset and the feature distribution of data
+            # generated by the generator.
+            perplexity = self.calculate_fid(real_features, fake_features)
+
+        return perplexity
+
+    def feature_extractor(self, inputs):
+        """Extract the feature of input data with InceptionV3.
+
+        The feature extracted from each input is a NumPy array
+        of length 2048.
+        """
+        # Since the input to InceptionV3 needs to be at least 75x75,
+        # we will pad the input image if needed.
+        hpad = math.ceil((75 - inputs.size(dim=-2)) / 2)
+        vpad = math.ceil((75 - inputs.size(dim=-1)) / 2)
+        hpad, vpad = max(0, hpad), max(0, vpad)
+        pad = nn.ZeroPad2d((hpad, hpad, vpad, vpad))
+        inputs = pad(inputs)
+
+        # Extract feature with InceptionV3
+        features = None
+        with torch.no_grad():
+            features = self.inception_model(inputs)
+        features = features.cpu()
+        features = np.array(features)
+
+        return features
+
+    def calculate_fid(self, real_features, fake_features):
+        """Calculate the Frechet Inception Distance (FID) between the
+        given real data feature and the synthetic data feature.
+
+        A lower FID indicates a better Generator model.
+
+        The implementation is borrowed from the following link:
+        https://wandb.ai/ayush-thakur/gan-evaluation/reports/How-to-Evaluate-GANs-using-Frechet-Inception-Distance-FID---Vmlldzo0MTAxOTI
+        """
+        # calculate mean and covariance statistics
+        mu1, sigma1 = real_features.mean(axis=0), np.cov(real_features,
+                                                         rowvar=False)
+        mu2, sigma2 = fake_features.mean(axis=0), np.cov(fake_features,
+                                                         rowvar=False)
+        # calculate sum squared difference between means
+        ssdiff = np.sum((mu1 - mu2)**2.0)
+        # calculate sqrt of product between cov
+        covmean = scipy.linalg.sqrtm(sigma1.dot(sigma2))
+        # check and correct imaginary numbers from sqrt
+        if np.iscomplexobj(covmean):
+            covmean = covmean.real
+        # calculate score
+        fid = ssdiff + np.trace(sigma1 + sigma2 - 2.0 * covmean)
+
+        return fid
diff --git a/plato/trainers/registry.py b/plato/trainers/registry.py
index 9839e5a2d..ee74c732e 100644
--- a/plato/trainers/registry.py
+++ b/plato/trainers/registry.py
@@ -26,11 +26,13 @@
         basic,
         diff_privacy,
         pascal_voc,
+        gan
     )
     registered_trainers = OrderedDict([
         ('basic', basic.Trainer),
         ('diff_privacy', diff_privacy.Trainer),
         ('pascal_voc', pascal_voc.Trainer),
+        ('gan', gan.Trainer),
     ])
 
 
diff --git a/run_workloads b/run_workloads
index 0bf440f7f..d14576f4f 100755
--- a/run_workloads
+++ b/run_workloads
@@ -37,7 +37,7 @@ if __name__ == '__main__':
         except FileExistsError:
             pass
         with open(os.path.join(log_folder, log_name), 'a',
-                  encoding='utf8') as fp:
+                  encoding='utf-8') as fp:
             fp.write(f'{t}\n\n')
             fp.flush()
             subprocess.call(['./run', f'--config={t}.yml'],


Test patch:
diff --git a/tests/TestsConfig/coco.yml b/tests/TestsConfig/coco.yml
index 9223a20f8..b938ba534 100644
--- a/tests/TestsConfig/coco.yml
+++ b/tests/TestsConfig/coco.yml
@@ -37,11 +37,6 @@ data:
   # The random seed for sampling data
   random_seed: 1
 
-  multi_modal_pipeliner:
-    rgb: !include Pipeline/kinetics_csn_rgb.yml
-    flow: !include Pipeline/kinetics_tsn_flow.yml
-    audio: !include Pipeline/kinetics_audio_feature.yml
-
 # Train settings
 
 # optimizer
@@ -92,10 +87,6 @@ trainer:
   log_config: *log_config
   checkpoint_config: *checkpoint_config
 
-model:
-  model_name: rgb_flow_audio_model
-  model_config: !include Models/kinetics_full_models.yml
-
 evaluation:
   interval: 5
   metrics:
diff --git a/tests/TestsConfig/distribution_noniid_sampler.yml b/tests/TestsConfig/distribution_noniid_sampler.yml
index b330079ee..dacef0e9d 100644
--- a/tests/TestsConfig/distribution_noniid_sampler.yml
+++ b/tests/TestsConfig/distribution_noniid_sampler.yml
@@ -38,11 +38,6 @@ data:
   # The random seed for sampling data
   random_seed: 1
 
-  multi_modal_pipeliner:
-    rgb: !include Pipeline/kinetics_csn_rgb.yml
-    flow: !include Pipeline/kinetics_tsn_flow.yml
-    audio: !include Pipeline/kinetics_audio_feature.yml
-
 # Train settings
 
 # optimizer
@@ -93,10 +88,6 @@ trainer:
   log_config: *log_config
   checkpoint_config: *checkpoint_config
 
-model:
-  model_name: rgb_flow_audio_model
-  model_config: !include Models/kinetics_full_models.yml
-
 evaluation:
   interval: 5
   metrics:
diff --git a/tests/TestsConfig/flickr30k_entities.yml b/tests/TestsConfig/flickr30k_entities.yml
index d3a819a5c..a73e9c0f9 100644
--- a/tests/TestsConfig/flickr30k_entities.yml
+++ b/tests/TestsConfig/flickr30k_entities.yml
@@ -19,7 +19,7 @@ data:
   datasource: Flickr30K
   data_path: data
 
-  download_file_id: 1mJ-Y0gLFyCLbB6C-M49dUtEsbfVkK4F-
+  download_url: http://iqua.ece.toronto.edu/baochun/Flickr30KEntitiesRaw.zip
 
   downloader:
     num_workers: 4
@@ -35,11 +35,6 @@ data:
   # The random seed for sampling data
   random_seed: 1
 
-  multi_modal_pipeliner:
-    rgb: !include Pipeline/kinetics_csn_rgb.yml
-    flow: !include Pipeline/kinetics_tsn_flow.yml
-    audio: !include Pipeline/kinetics_audio_feature.yml
-
 # Train settings
 
 # optimizer
@@ -90,10 +85,6 @@ trainer:
   log_config: *log_config
   checkpoint_config: *checkpoint_config
 
-model:
-  model_name: rgb_flow_audio_model
-  model_config: !include Models/kinetics_full_models.yml
-
 evaluation:
   interval: 5
   metrics:
diff --git a/tests/TestsConfig/gym.yml b/tests/TestsConfig/gym.yml
index 58eec1c8f..3daf6993d 100644
--- a/tests/TestsConfig/gym.yml
+++ b/tests/TestsConfig/gym.yml
@@ -35,11 +35,6 @@ data:
   # The random seed for sampling data
   random_seed: 1
 
-  multi_modal_configs:
-    rgb: !include Pipeline/gym_csn_rgb.yml
-    flow: !include Pipeline/gym_tsn_flow.yml
-    audio_feature: !include Pipeline/gym_audio_feature.yml
-
 # Train settings
 
 # optimizer
@@ -90,10 +85,6 @@ trainer:
   log_config: *log_config
   checkpoint_config: *checkpoint_config
 
-model:
-  model_name: rgb_flow_audio_model
-  model_config: !include Models/kinetics_full_models.yml
-
 evaluation:
   interval: 5
   metrics:
diff --git a/tests/TestsConfig/kinetics.yml b/tests/TestsConfig/kinetics.yml
index dc2e1bfe0..0c07e1cf9 100644
--- a/tests/TestsConfig/kinetics.yml
+++ b/tests/TestsConfig/kinetics.yml
@@ -39,11 +39,6 @@ data:
   # The random seed for sampling data
   random_seed: 1
 
-  multi_modal_configs:
-    rgb: !include Pipeline/kinetics_csn_rgb.yml
-    flow: !include Pipeline/kinetics_tsn_flow.yml
-    audio_feature: !include Pipeline/kinetics_audio_feature.yml
-
 # Train settings
 
 # optimizer
@@ -94,10 +89,6 @@ trainer:
   log_config: *log_config
   checkpoint_config: *checkpoint_config
 
-model:
-  model_name: rgb_flow_audio_model
-  model_config: !include Models/kinetics_full_models.yml
-
 evaluation:
   interval: 5
   metrics:
diff --git a/tests/TestsConfig/label_quantity_noniid_sampler.yml b/tests/TestsConfig/label_quantity_noniid_sampler.yml
index 5a5ead9fe..459487d5e 100644
--- a/tests/TestsConfig/label_quantity_noniid_sampler.yml
+++ b/tests/TestsConfig/label_quantity_noniid_sampler.yml
@@ -37,10 +37,6 @@ data:
   # The random seed for sampling data
   random_seed: 1
 
-  multi_modal_pipeliner:
-    rgb: !include Pipeline/kinetics_csn_rgb.yml
-    flow: !include Pipeline/kinetics_tsn_flow.yml
-    audio: !include Pipeline/kinetics_audio_feature.yml
 
 # Train settings
 
@@ -92,10 +88,6 @@ trainer:
   log_config: *log_config
   checkpoint_config: *checkpoint_config
 
-model:
-  model_name: rgb_flow_audio_model
-  model_config: !include Models/kinetics_full_models.yml
-
 evaluation:
   interval: 5
   metrics:
diff --git a/tests/TestsConfig/mixed_label_quantity_noniid_sampler.yml b/tests/TestsConfig/mixed_label_quantity_noniid_sampler.yml
index a07b6473f..4f305a3b7 100644
--- a/tests/TestsConfig/mixed_label_quantity_noniid_sampler.yml
+++ b/tests/TestsConfig/mixed_label_quantity_noniid_sampler.yml
@@ -43,11 +43,6 @@ data:
   # The random seed for sampling data
   random_seed: 1
 
-  multi_modal_pipeliner:
-    rgb: !include Pipeline/kinetics_csn_rgb.yml
-    flow: !include Pipeline/kinetics_tsn_flow.yml
-    audio: !include Pipeline/kinetics_audio_feature.yml
-
 # Train settings
 
 # optimizer
@@ -98,10 +93,6 @@ trainer:
   log_config: *log_config
   checkpoint_config: *checkpoint_config
 
-model:
-  model_name: rgb_flow_audio_model
-  model_config: !include Models/kinetics_full_models.yml
-
 evaluation:
   interval: 5
   metrics:
diff --git a/tests/TestsConfig/referitgame.yml b/tests/TestsConfig/referitgame.yml
index b30c1e9a9..a18627596 100644
--- a/tests/TestsConfig/referitgame.yml
+++ b/tests/TestsConfig/referitgame.yml
@@ -40,11 +40,6 @@ data:
   # The random seed for sampling data
   random_seed: 1
 
-  multi_modal_pipeliner:
-    rgb: !include Pipeline/kinetics_csn_rgb.yml
-    flow: !include Pipeline/kinetics_tsn_flow.yml
-    audio: !include Pipeline/kinetics_audio_feature.yml
-
 # Train settings
 
 # optimizer
@@ -95,10 +90,6 @@ trainer:
   log_config: *log_config
   checkpoint_config: *checkpoint_config
 
-model:
-  model_name: rgb_flow_audio_model
-  model_config: !include Models/kinetics_full_models.yml
-
 evaluation:
   interval: 5
   metrics:
diff --git a/tests/TestsConfig/sample_quantity_noniid_sampler.yml b/tests/TestsConfig/sample_quantity_noniid_sampler.yml
index 873048ad2..d613f3baf 100644
--- a/tests/TestsConfig/sample_quantity_noniid_sampler.yml
+++ b/tests/TestsConfig/sample_quantity_noniid_sampler.yml
@@ -37,11 +37,6 @@ data:
   # The random seed for sampling data
   random_seed: 1
 
-  multi_modal_pipeliner:
-    rgb: !include Pipeline/kinetics_csn_rgb.yml
-    flow: !include Pipeline/kinetics_tsn_flow.yml
-    audio: !include Pipeline/kinetics_audio_feature.yml
-
 # Train settings
 
 # optimizer
@@ -92,10 +87,6 @@ trainer:
   log_config: *log_config
   checkpoint_config: *checkpoint_config
 
-model:
-  model_name: rgb_flow_audio_model
-  model_config: !include Models/kinetics_full_models.yml
-
 evaluation:
   interval: 5
   metrics:
diff --git a/tests/config_tests.py b/tests/config_tests.py
index d23c1db08..f613e7fed 100644
--- a/tests/config_tests.py
+++ b/tests/config_tests.py
@@ -1,50 +1,48 @@
 """
 Testing a federated learning configuration.
+
+How to run the tests:
+
+ 1 Run the following command in the root directory.
+    python tests/config_tests.py
+
 """
 
 import os
 import unittest
 import warnings
 
-os.environ['config_file'] = 'configs/Kinetics/kinetics_mm.yml'
-
 from plato.config import Config
 
+os.environ['config_file'] = 'configs/Kinetics/kinetics_mm.yml'
+
 
 class ConfigTest(unittest.TestCase):
     """ Testing the correctness of loading a configuration file. """
-    def setup(self):
-        super().setup()
-
-        self.addTypeEqualityFunc(Config, 'assertConfigEqual')
 
-        self.defined_config = Config()
+    def setUp(self):
+        super().setUp()
 
-        # define several example parameters
+        # define several example parameters that will be used as
+        #   a demo to test the loaded configuration file.
         data_params_config = {
             "downloader": {
                 "num_workers": 4
             },
-            "multi_modal_pipeliner": {
+            "multi_modal_configs": {
                 "rgb": {
-                    "rgb_data": {
-                        "train": {
-                            "type": "RawframeDataset"
-                        }
+                    "train": {
+                        "type": "RawframeDataset"
                     }
                 },
                 "flow": {
-                    "flow_data": {
-                        "train": {
-                            "type": "RawframeDataset"
-                        }
+                    "train": {
+                        "type": "RawframeDataset"
                     }
                 },
                 "audio": {
-                    "audio_data": {
-                        "train": {
-                            "type": "AudioFeatureDataset"
-                        }
+                    "train": {
+                        "type": "AudioFeatureDataset"
                     }
                 }
             }
@@ -61,64 +59,54 @@ def setup(self):
         self.data_config = Config.namedtuple_from_dict(data_params_config)
         self.model_config = Config.namedtuple_from_dict(model_params_config)
 
-    def assertAttrContained(self, src_config, dst_key):
-        """ The attrs of dst_key should be included in the src_config"""
-        assert hasattr(src_config, dst_key)
+    def test_dataconfig(self):
+        """ Test the structure and necessary parameters of the data configuration.
 
-    def assertAttrValueEqual(self, src_value, dst_value):
-        """ The two values should be equal with each other """
-        self.assertEqual(src_value, dst_value)
+            This function tests the core part of the configuration file. It tests:
+                downloader, num_workers, multi_modal_configs, and whether there are
+                three modalities, including rgb, flow, and audio and the corresponding
+                data types are correct.
 
-    def test_dataconfig(self):
-        """ Test the structure and necessary parameters of the data configuration """
+        """
         with warnings.catch_warnings():
             warnings.filterwarnings("ignore", category=UserWarning)
 
             defined_data_config = Config().data
 
-            self.assertAttrContained(src_config=defined_data_config,
-                                     dst_key="downloader")
-            self.assertAttrContained(src_config=defined_data_config.downloader,
-                                     dst_key="num_workers")
-            self.assertAttrValueEqual(
-                src_value=defined_data_config.downloader.num_workers,
-                dst_value=self.data_config.downloader.num_workers)
-
-            self.assertAttrContained(src_config=defined_data_config,
-                                     dst_key="multi_modal_pipeliner")
-            self.assertAttrContained(
-                src_config=defined_data_config.multi_modal_pipeliner,
-                dst_key="rgb")
-            self.assertAttrContained(
-                src_config=defined_data_config.multi_modal_pipeliner,
-                dst_key="flow")
-            self.assertAttrContained(
-                src_config=defined_data_config.multi_modal_pipeliner,
-                dst_key="audio")
-
-            self.assertAttrContained(
-                src_config=defined_data_config.multi_modal_pipeliner.rgb,
-                dst_key="rgb_data")
-            self.assertAttrValueEqual(
-                src_value=defined_data_config.multi_modal_pipeliner.rgb.
-                rgb_data.train.type,
-                dst_value=self.data_config.multi_modal_pipeliner.rgb.rgb_data.
-                train.type)
+            hasattr(defined_data_config, "downloader")
+            hasattr(defined_data_config.downloader, "num_workers")
+
+            self.assertEqual(defined_data_config.downloader.num_workers,
+                             self.data_config.downloader.num_workers)
+
+            hasattr(defined_data_config, "multi_modal_configs")
+            hasattr(defined_data_config.multi_modal_configs, "rgb")
+            hasattr(defined_data_config.multi_modal_configs, "flow")
+            hasattr(defined_data_config.multi_modal_configs, "audio")
+            hasattr(defined_data_config.multi_modal_configs.rgb, "train")
+
+            self.assertEqual(
+                defined_data_config.multi_modal_configs.rgb.train.type,
+                self.data_config.multi_modal_configs.rgb.train.type)
 
     def test_modelconfig(self):
-        """ Test the structure and necessary parameters of the model configuration """
+        """ Test the structure and necessary parameters of the model configuration.
+
+            This function tests the core part of the configuration file. It tests whether
+            the model is defined and whether the configurations for the defined model is
+            correct.
+        """
         with warnings.catch_warnings():
             warnings.filterwarnings("ignore", category=UserWarning)
 
             defined_model_config = Config().model
-            self.assertAttrContained(src_config=defined_model_config,
-                                     dst_key="model_name")
-            self.assertAttrValueEqual(
-                src_value=defined_model_config.model_name,
-                dst_value="rgb_flow_audio_model")
-            self.assertAttrValueEqual(
-                src_value=defined_model_config.model_config.rgb_model.type,
-                dst_value=self.model_config.model_config.rgb_model.type)
+
+            hasattr(defined_model_config, "model_name")
+
+            self.assertEqual(defined_model_config.model_name,
+                             "rgb_flow_audio_model")
+            self.assertEqual(defined_model_config.model_config.rgb_model.type,
+                             self.model_config.model_config.rgb_model.type)
 
 
 if __name__ == '__main__':
diff --git a/tests/data_tests.py b/tests/data_tests.py
index ffdc9f08f..33744ecb2 100644
--- a/tests/data_tests.py
+++ b/tests/data_tests.py
@@ -1,16 +1,38 @@
 """
 Testing multi-modal datasources of Plato framework.
+
+How to run the tests:
+
+For example, when you want to test the ReferItGame dataset.
+ 1 Uncomment the 'import' code for the dataset you want to test
+    # from plato.datasources.referitgame import DataSource as refer_Datasource
+    to   from plato.datasources.referitgame import DataSource as refer_Datasource
+
+ 2 Uncomment the configuration file for the dataset you want to test
+    # os.environ['config_file'] = 'tests/TestsConfig/referitgame.yml'
+    to   os.environ['config_file'] = 'tests/TestsConfig/referitgame.yml'
+
+ 3 Uncomment the corresponding function that conducts the dataset test.
+    i.e. test_ref_datasource() for the ReferItGame dataset
+
+ 4 Run the following command in the root directory.
+    python tests/data_tests.py
+
+
+Note, before testing the referitgame dataset, you need to first download the
+    COCO data.
+
 """
 
 import os
 import unittest
 
+# os.environ['config_file'] = 'tests/TestsConfig/referitgame.yml'
+
 os.environ['config_file'] = 'tests/TestsConfig/flickr30k_entities.yml'
 
 # os.environ['config_file'] = 'tests/TestsConfig/coco.yml'
 
-# os.environ['config_file'] = 'tests/TestsConfig/referitgame.yml'
-
 # os.environ['config_file'] = 'tests/TestsConfig/kinetics.yml'
 
 # os.environ['config_file'] = 'tests/TestsConfig/gym.yml'
@@ -21,6 +43,8 @@
 import numpy as np
 import torch
 
+from sampler_test_utils import define_sampler
+
 from plato.config import Config
 
 from plato.datasources.flickr30k_entities import DataSource as f30ke_DataSource
@@ -32,11 +56,10 @@
 
 from plato.samplers import modality_iid
 
-from sampler_test_utils import define_sampler
-
 
 class DatasetsTest(unittest.TestCase):
     """ Aiming to test the correcness of implemented samplers and datasets """
+
     def setUp(self):
         super().setUp()
 
@@ -51,7 +74,7 @@ def setUp(self):
         # the datasource being tested
         self.utest_datasource = None
 
-    def assertDataSourceDefinition(self, data_source):
+    def assert_datasource_definition(self, data_source):
         """ Test whether the dataset can be correctly defined.
             This verifies:
             1.1- The datasource can be defined
@@ -86,29 +109,28 @@ def assertDataSourceDefinition(self, data_source):
             batch_size=batch_size,
             sampler=defined_sampler.get())
 
-        obtained_mm_sample = next(iter(testset_loader))
-        print("obtained_sample: ", obtained_mm_sample)
+        _ = next(iter(testset_loader))
 
         return True
 
     def test_f30ke_datasource(self):
         """ Test the flickr30k entities dataset. """
         self.utest_datasource = f30ke_DataSource()
-        assert self.assertDataSourceDefinition(self.utest_datasource)
+        assert self.assert_datasource_definition(self.utest_datasource)
 
     # def test_coco_datasource(self):
     #     """ Test the MSCOCO dataset. """
     #     # set the specific
 
     #     self.utest_datasource = coco_Datasource()
-    #     # assert self.assertDataSourceDefinition(self.utest_datasource)
+    #     # assert self.assert_datasource_definition(self.utest_datasource)
 
     # def test_ref_datasource(self):
     #     """ Test the ReferItGmae dataset. """
     #     # set the specific
 
     #     self.utest_datasource = refer_Datasource()
-    #     assert self.assertDataSourceDefinition(self.utest_datasource)
+    #     assert self.assert_datasource_definition(self.utest_datasource)
 
     # def test_kinetics_datasource(self):
     #     """ Test the kinetics700 dataset. """
@@ -123,14 +145,14 @@ def test_f30ke_datasource(self):
     #     print(iter_data["rgb"]["imgs"].shape)
     #     print(iter_data["rgb"]["label"].shape)
 
-    #     assert self.assertDataSourceDefinition(self.utest_datasource)
+    #     assert self.assert_datasource_definition(self.utest_datasource)
 
     # def test_gym_datasource(self):
     #     """ Test the Gym dataset. """
     #     # set the specific
 
     #     self.utest_datasource = GymDataSource()
-    #     assert self.assertDataSourceDefinition(self.utest_datasource)
+    #     assert self.assert_datasource_definition(self.utest_datasource)
 
 
 if __name__ == '__main__':
diff --git a/tests/lr_schedule_tests.py b/tests/lr_schedule_tests.py
index 32f9f1ba2..f2b1e2574 100644
--- a/tests/lr_schedule_tests.py
+++ b/tests/lr_schedule_tests.py
@@ -10,6 +10,7 @@
 
 
 class LrSchedulerTest(unittest.TestCase):
+
     def setUp(self):
         super().setUp()
         __ = Config()
@@ -25,7 +26,7 @@ def setUp(self):
         self.model = models_registry.get()
         self.optimizer = optimizers.get_optimizer(self.model)
 
-    def assertLrEquals(self, lr):
+    def assert_lr_equal(self, lr):
         self.assertEqual(np.round(self.optimizer.param_groups[0]['lr'], 10),
                          np.round(lr, 10))
 
@@ -41,14 +42,14 @@ def test_vanilla(self):
             Config().trainer = namedtuple('trainer', fields)(*params)
 
             lrs = optimizers.get_lr_schedule(self.optimizer, 10)
-            self.assertLrEquals(0.1)
+            self.assert_lr_equal(0.1)
             for _ in range(100):
                 lrs.step()
-            self.assertLrEquals(0.1)
-            self.assertLrEquals(0.1)
+            self.assert_lr_equal(0.1)
+            self.assert_lr_equal(0.1)
 
     def test_milestones(self):
-        self.assertLrEquals(0.1)
+        self.assert_lr_equal(0.1)
 
         with warnings.catch_warnings():
             warnings.filterwarnings("ignore", category=UserWarning)
@@ -60,42 +61,42 @@ def test_milestones(self):
             params = ['SGD', 'LambdaLR', 0.1, 0.5, 0.0, 0.1, '2ep,4ep,7ep,8ep']
 
             Config().trainer = namedtuple('trainer', fields)(*params)
-            self.assertLrEquals(0.1)
+            self.assert_lr_equal(0.1)
 
             lrs = optimizers.get_lr_schedule(self.optimizer, 10)
 
-            self.assertLrEquals(0.1)
+            self.assert_lr_equal(0.1)
             for _ in range(19):
                 lrs.step()
-            self.assertLrEquals(1e-1)
+            self.assert_lr_equal(1e-1)
 
             for _ in range(1):
                 lrs.step()
-            self.assertLrEquals(1e-2)
+            self.assert_lr_equal(1e-2)
             for _ in range(19):
                 lrs.step()
-            self.assertLrEquals(1e-2)
+            self.assert_lr_equal(1e-2)
 
             for _ in range(1):
                 lrs.step()
-            self.assertLrEquals(1e-3)
+            self.assert_lr_equal(1e-3)
             for _ in range(29):
                 lrs.step()
-            self.assertLrEquals(1e-3)
+            self.assert_lr_equal(1e-3)
 
             for _ in range(1):
                 lrs.step()
-            self.assertLrEquals(1e-4)
+            self.assert_lr_equal(1e-4)
             for _ in range(9):
                 lrs.step()
-            self.assertLrEquals(1e-4)
+            self.assert_lr_equal(1e-4)
 
             for _ in range(1):
                 lrs.step()
-            self.assertLrEquals(1e-5)
+            self.assert_lr_equal(1e-5)
             for _ in range(100):
                 lrs.step()
-            self.assertLrEquals(1e-5)
+            self.assert_lr_equal(1e-5)
 
     def test_warmup(self):
         with warnings.catch_warnings():
@@ -112,12 +113,12 @@ def test_warmup(self):
             lrs = optimizers.get_lr_schedule(self.optimizer, 10)
 
             for i in range(20):
-                self.assertLrEquals(i / 20 * 0.1)
+                self.assert_lr_equal(i / 20 * 0.1)
                 lrs.step()
-            self.assertLrEquals(0.1)
+            self.assert_lr_equal(0.1)
             for i in range(100):
                 lrs.step()
-                self.assertLrEquals(0.1)
+                self.assert_lr_equal(0.1)
 
 
 if __name__ == '__main__':
diff --git a/tests/sampler_tests.py b/tests/sampler_tests.py
index 2285b4415..8476ec5c9 100644
--- a/tests/sampler_tests.py
+++ b/tests/sampler_tests.py
@@ -1,5 +1,19 @@
 """
 Testing samplers in Plato framework.
+
+How to run the tests:
+
+For example, when you want to test the label_quantity_noniid_sampler.
+
+ 1 Uncomment the configuration file for the dataset you want to test
+    # os.environ[
+    'config_file'] = 'tests/TestsConfig/label_quantity_noniid_sampler.yml'
+    to   os.environ[
+    'config_file'] = 'tests/TestsConfig/label_quantity_noniid_sampler.yml'
+
+ 2 Run the following command in the root directory.
+    python tests/sampler_tests.py
+
 """
 
 import os
@@ -7,25 +21,26 @@
 
 # os.environ['config_file'] = 'tests/TestsConfig/distribution_noniid_sampler.yml'
 
-os.environ[
-    'config_file'] = 'tests/TestsConfig/label_quantity_noniid_sampler.yml'
+# os.environ['config_file'] = 'tests/TestsConfig/label_quantity_noniid_sampler.yml'
 
 # os.environ[
 #     'config_file'] = 'tests/TestsConfig/mixed_label_quantity_noniid_sampler.yml'
 
-# os.environ[
-#     'config_file'] = 'tests/TestsConfig/sample_quantity_noniid_sampler.yml'
+os.environ[
+    'config_file'] = 'tests/TestsConfig/sample_quantity_noniid_sampler.yml'
 
-import sampler_test_utils
 import numpy as np
 
 from plato.config import Config
 from plato.datasources.cifar10 import DataSource
 from plato.samplers import registry as samplers_registry
 
+import sampler_test_utils
+
 
 class SamplersTest(unittest.TestCase):
     """ Testing the correctness of implemented samplers. """
+
     def setUp(self):
         super().setUp()