plato-learn 1.1__py3-none-any.whl

plato/trainers/diff_privacy.py

@@ -0,0 +1,178 @@
+"""
+The training and testing loops for PyTorch.
+"""
+
+import logging
+import time
+
+from opacus import GradSampleModule
+from opacus.privacy_engine import PrivacyEngine
+from opacus.utils.batch_memory_manager import BatchMemoryManager
+from opacus.validators import ModuleValidator
+from torch.utils.data import Subset
+
+from plato.config import Config
+from plato.trainers import basic
+
+
+class Trainer(basic.Trainer):
+    """A differentially private federated learning trainer, used by the client."""
+
+    def __init__(self, model=None, **kwargs):
+        """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)
+        if len(errors) > 0:
+            self.model = ModuleValidator.fix(self.model)
+            errors = ModuleValidator.validate(self.model, strict=False)
+            assert len(errors) == 0
+
+    # pylint: disable=unused-argument
+    def train_model(self, config, trainset, sampler, **kwargs):
+        """The default training loop that supports differential privacy."""
+        batch_size = config["batch_size"]
+        self.sampler = sampler
+        tic = time.perf_counter()
+
+        self.train_run_start(config)
+        self.callback_handler.call_event("on_train_run_start", self, config)
+
+        # We have to use poisson sampling to sample the data, rather than the provided sampler.
+        # Replacing the poisson sampler with the provided sampler is problematic since it may
+        # violate the basic theory of DP-SGD. Therefore, we need to first obtain the train subset
+        # based on the provided sampler, and then create a simple dataloader on the train subset
+        # without the sampler. We will finally use Opacus to recreate the dataloader from the
+        # simple dataloader (with poisson sampling).
+        trainset = Subset(trainset, list(sampler))
+        self.train_loader = self.get_train_loader(batch_size, trainset, sampler=None)
+
+        # Initializing the loss criterion
+        _loss_criterion = self.get_loss_criterion()
+
+        # Initializing the optimizer
+        optimizer = self.get_optimizer(self.model)
+        self.lr_scheduler = self.get_lr_scheduler(config, optimizer)
+        optimizer = self._adjust_lr(config, self.lr_scheduler, optimizer)
+
+        self.model.to(self.device)
+        total_epochs = config["epochs"]
+
+        logging.info(
+            "[Client #%s] Using differential privacy during training.",
+            self.client_id,
+        )
+
+        privacy_engine = PrivacyEngine(accountant="rdp", secure_mode=False)
+
+        self.model, optimizer, train_loader = privacy_engine.make_private_with_epsilon(
+            module=self.model,
+            optimizer=optimizer,
+            data_loader=self.train_loader,
+            target_epsilon=config["dp_epsilon"] if "dp_epsilon" in config else 10.0,
+            target_delta=config["dp_delta"] if "dp_delta" in config else 1e-5,
+            epochs=total_epochs,
+            max_grad_norm=config["dp_max_grad_norm"]
+            if "max_grad_norm" in config
+            else 1.0,
+        )
+
+        self.model.train()
+
+        for self.current_epoch in range(1, total_epochs + 1):
+            with BatchMemoryManager(
+                data_loader=train_loader,
+                max_physical_batch_size=self.max_physical_batch_size,
+                optimizer=optimizer,
+            ) as memory_safe_train_loader:
+                self._loss_tracker.reset()
+                self.train_epoch_start(config)
+                self.callback_handler.call_event("on_train_epoch_start", self, config)
+
+                for batch_id, (examples, labels) in enumerate(memory_safe_train_loader):
+                    examples, labels = (
+                        examples.to(self.device),
+                        labels.to(self.device),
+                    )
+                    optimizer.zero_grad(set_to_none=True)
+
+                    outputs = self.model(examples)
+
+                    loss = _loss_criterion(outputs, labels)
+                    self._loss_tracker.update(loss, labels.size(0))
+
+                    if "create_graph" in config:
+                        loss.backward(create_graph=config["create_graph"])
+                    else:
+                        loss.backward()
+
+                    optimizer.step()
+
+                    self.train_step_end(config, batch=batch_id, loss=loss)
+                    self.callback_handler.call_event(
+                        "on_train_step_end",
+                        self,
+                        config,
+                        batch=batch_id,
+                        loss=loss,
+                    )
+
+            self.lr_scheduler_step()
+
+            if hasattr(optimizer, "params_state_update"):
+                optimizer.params_state_update()
+
+            # Simulate client's speed
+            if (
+                self.client_id != 0
+                and hasattr(Config().clients, "speed_simulation")
+                and Config().clients.speed_simulation
+            ):
+                self.simulate_sleep_time()
+
+            # Saving the model at the end of this epoch to a file so that
+            # it can later be retrieved to respond to server requests
+            # in asynchronous mode when the wall clock time is simulated
+            if (
+                hasattr(Config().server, "request_update")
+                and Config().server.request_update
+            ):
+                self.model.cpu()
+                training_time = time.perf_counter() - tic
+                filename = f"{self.client_id}_{self.current_epoch}_{training_time}.pth"
+                self.save_model(filename)
+                self.model.to(self.device)
+
+            self.run_history.update_metric("train_loss", self._loss_tracker.average)
+            self.train_epoch_end(config)
+            self.callback_handler.call_event("on_train_epoch_end", self, config)
+
+        self.train_run_end(config)
+        self.callback_handler.call_event("on_train_run_end", self, config)
+
+    def train_run_start(self, config):
+        """
+        Method called at the start of training run.
+        """
+        self.model = GradSampleModule(self.model)
+
+    def train_run_end(self, config):
+        """
+        Method called at the end of a training run.
+        """
+        # 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()
+        }

plato/trainers/gan.py

@@ -0,0 +1,330 @@
+"""
+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.trainers import optimizers
+
+
+class Trainer(basic.Trainer):
+    """A federated learning trainer for GAN models."""
+
+    def __init__(self, model=None, **kwargs):
+        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))
+
+    # pylint: disable=unused-argument
+    def train_model(self, config, trainset, sampler, **kwargs):
+        """The main training loop in a federated learning workload.
+
+        Arguments:
+        trainset: The training dataset.
+        sampler: the sampler that extracts a partition for this client.
+
+        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(self.generator)
+        optimizer_disc = optimizers.get(self.discriminator)
+
+        real_label = 1.0
+        fake_label = 0.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, sampler=None, **kwargs):
+        """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

plato/trainers/huggingface.py

@@ -0,0 +1,173 @@
+"""
+Training and testing loops for HuggingFace's transformer models for natural
+language processing.
+"""
+
+import math
+from typing import Optional
+
+from torch.utils.data import RandomSampler, Sampler
+
+from transformers import (
+    AutoConfig,
+    AutoTokenizer,
+    HfArgumentParser,
+    TrainerCallback,
+    LlamaTokenizer,
+)
+from transformers import Trainer as HuggingFaceTrainer
+from transformers import TrainingArguments, default_data_collator
+
+from plato.config import Config
+from plato.trainers import basic
+
+
+class SampledHuggingFaceTrainer(HuggingFaceTrainer):
+    """
+    Training and testing loops for HuggingFace's transformer models for natural
+    language processing.
+    """
+
+    def __init__(
+        self,
+        model,
+        args,
+        train_dataset,
+        eval_dataset,
+        tokenizer,
+        data_collator,
+        sampler,
+        callbacks,
+    ):
+        super().__init__(
+            model=model,
+            args=args,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+            tokenizer=tokenizer,
+            data_collator=data_collator,
+            callbacks=callbacks,
+        )
+        self.sampler = sampler
+
+    def _get_train_sampler(self) -> Optional[Sampler]:
+        if self.sampler is None:
+            return RandomSampler(self.train_dataset)
+
+        return self.sampler
+
+    def _get_eval_sampler(self, eval_dataset) -> Optional[Sampler]:
+        if self.sampler is None:
+            return super()._get_eval_sampler(eval_dataset)
+
+        return self.sampler
+
+
+class Trainer(basic.Trainer):
+    """The trainer for HuggingFace transformer models for natural language processing."""
+
+    def __init__(self, model=None, callbacks=None):
+        super().__init__(model)
+
+        self.trainer = None
+        self.trainer_callbacks = []
+        if callbacks:
+            # Huggingface needs to check callback types
+            self.add_callbacks(callbacks)
+
+        self.model.train()
+
+        parser = HfArgumentParser(TrainingArguments)
+        (self.training_args,) = parser.parse_args_into_dataclasses(
+            args=[
+                "--output_dir=" + Config.params["checkpoint_path"],
+                "--report_to=none",
+            ]
+        )
+
+        model_name = Config().trainer.model_name
+        config_kwargs = {
+            "cache_dir": None,
+            "revision": "main",
+            "use_auth_token": None,
+        }
+        self.config = AutoConfig.from_pretrained(model_name, **config_kwargs)
+
+        tokenizer_kwargs = {
+            "cache_dir": None,
+            "use_fast": True,
+            "revision": "main",
+            "use_auth_token": None,
+        }
+        if "llama" in model_name:
+            self.tokenizer = LlamaTokenizer.from_pretrained(
+                model_name, config=self.config, **tokenizer_kwargs
+            )
+        else:
+            self.tokenizer = AutoTokenizer.from_pretrained(
+                model_name, config=self.config, **tokenizer_kwargs
+            )
+
+    # pylint: disable=unused-argument
+    def train_model(self, config, trainset, sampler, **kwargs):
+        """The training loop for HuggingFace models.
+
+        Arguments:
+        config: A dictionary of configuration parameters.
+        trainset: The training dataset.
+        sampler: the sampler that extracts a partition for this client.
+        """
+
+        self.training_args.num_train_epochs = config["epochs"]
+        self.training_args.per_device_train_batch_size = config["batch_size"]
+
+        self.trainer = SampledHuggingFaceTrainer(
+            model=self.model,
+            args=self.training_args,
+            train_dataset=trainset,
+            eval_dataset=None,
+            tokenizer=self.tokenizer,
+            data_collator=default_data_collator,
+            sampler=sampler,
+            callbacks=self.trainer_callbacks,
+        )
+
+        self.trainer.train()
+
+    def test_model(self, config, testset, sampler=None, **kwargs):  # pylint: disable=unused-argument
+        """The testing loop for HuggingFace models.
+
+        Arguments:
+            config: Configuration parameters as a dictionary.
+            testset: The test dataset.
+        """
+        self.training_args.per_device_eval_batch_size = config["batch_size"]
+
+        self.trainer = SampledHuggingFaceTrainer(
+            model=self.model,
+            args=self.training_args,
+            train_dataset=None,
+            eval_dataset=testset,
+            tokenizer=self.tokenizer,
+            data_collator=default_data_collator,
+            sampler=sampler,
+            callbacks=None,
+        )
+
+        metrics = self.trainer.evaluate()
+
+        try:
+            perplexity = math.exp(metrics["eval_loss"])
+        except OverflowError:
+            perplexity = float("inf")
+
+        return perplexity
+
+    def add_callbacks(self, callbacks):
+        """Callbacks will be handled by Huggingface instead of Plato."""
+        for callback in callbacks:
+            if not issubclass(callback, TrainerCallback):
+                raise ValueError(
+                    f"Huggingface trainer expects subclass of {TrainerCallback}, got {callback} instead."
+                )
+        self.trainer_callbacks.extend(callbacks)