Supported federated learning with generative adversarial networks (GANs). (#189)

Co-authored-by: Baochun Li <bli@ece.toronto.edu>

Author: cuiboyuan

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

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

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)

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),
     ])
 
 

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

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()

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),
     ])