plato-learn 1.1__py3-none-any.whl

plato/trainers/self_supervised_learning.py

@@ -0,0 +1,302 @@
+"""
+A self-supervised learning (SSL) trainer for SSL training and testing.
+
+Federated learning with SSL trains the global model based on the data loader and
+objective function of SSL algorithms. For this unsupervised learning process, we
+cannot test the model directly as the model only extracts features from the
+data. Therefore, we use KNN as a classifier to get the accuracy of the global
+model during the regular federated training process.
+
+In the personalization process, each client trains a linear layer locally, based
+on the features extracted by the trained global model.
+
+The accuracy obtained by KNN during the regular federated training rounds may
+not be used to compare with the accuracy in supervised learning methods.
+"""
+
+import logging
+from collections import UserList
+
+import torch
+from lightly.data.multi_view_collate import MultiViewCollate
+
+from plato.config import Config
+from plato.trainers import basic
+from plato.models import registry as models_registry
+from plato.trainers import optimizers, lr_schedulers, loss_criterion
+
+
+class SSLSamples(UserList):
+    """A container for SSL sample, which contains multiple views as a list."""
+
+    def to(self, device):
+        """Assign a list of views into the specific device."""
+        for view_idx, view in enumerate(self.data):
+            if isinstance(view, torch.Tensor):
+                view = view.to(device)
+
+            self[view_idx] = view
+
+        return self.data
+
+
+class MultiViewCollateWrapper(MultiViewCollate):
+    """
+    An interface to connect collate from lightly with Plato's data loading mechanism.
+    """
+
+    def __call__(self, batch):
+        """Turn a batch of tuples into a single tuple."""
+        # Add a fname to each sample to make the batch compatible with lightly
+        batch = [batch[i] + (" ",) for i in range(len(batch))]
+
+        # Process first two parts with the lightly collate
+        views, labels, _ = super().__call__(batch)
+
+        # Assign views, which is a list of tensors, into SSLSamples
+        samples = SSLSamples(views)
+        return samples, labels
+
+
+class Trainer(basic.Trainer):
+    """A federated SSL trainer."""
+
+    def __init__(self, model=None, callbacks=None):
+        """Initialize the trainer."""
+        super().__init__(model=model, callbacks=callbacks)
+
+        # Datasets for personalization.
+        self.personalized_trainset = None
+        self.personalized_testset = None
+
+        # Define the personalized model
+        model_params = Config().parameters.personalization.model._asdict()
+        model_params["input_dim"] = self.model.encoder.encoding_dim
+        model_params["output_dim"] = model_params["num_classes"]
+        self.local_layers = models_registry.get(
+            model_name=Config().algorithm.personalization.model_name,
+            model_type=Config().algorithm.personalization.model_type,
+            model_params=model_params,
+        )
+
+    def set_personalized_datasets(self, trainset, testset):
+        """Set the personalized trainset."""
+        self.personalized_trainset = trainset
+        self.personalized_testset = testset
+
+    def get_train_loader(self, batch_size, trainset, sampler, **kwargs):
+        """Obtain the training loader based on the learning mode."""
+        # Get the trainloader for personalization
+        if self.current_round > Config().trainer.rounds:
+            return torch.utils.data.DataLoader(
+                dataset=self.personalized_trainset,
+                shuffle=False,
+                batch_size=batch_size,
+                sampler=sampler,
+            )
+        else:
+            collate_fn = MultiViewCollateWrapper()
+            return torch.utils.data.DataLoader(
+                dataset=trainset,
+                shuffle=False,
+                batch_size=batch_size,
+                sampler=sampler,
+                collate_fn=collate_fn,
+            )
+
+    def get_optimizer(self, model):
+        """Return the optimizer for SSL and personalization."""
+        if self.current_round <= Config().trainer.rounds:
+            return super().get_optimizer(model)
+        # Define the optimizer for the personalized model
+        optimizer_name = Config().algorithm.personalization.optimizer
+        optimizer_params = Config().parameters.personalization.optimizer._asdict()
+        return optimizers.get(
+            self.local_layers,
+            optimizer_name=optimizer_name,
+            optimizer_params=optimizer_params,
+        )
+
+    def get_ssl_criterion(self):
+        """
+        Get the loss criterion for SSL. Some SSL algorithms, for example,
+        BYOL, will overwrite this function for specific loss functions.
+        """
+
+        # Get loss criterion for the SSL
+        ssl_loss_function = loss_criterion.get()
+
+        # We need to wrap the loss function to make it compatible
+        # with different types of outputs
+        # The types of the outputs can vary from Tensor to a list of Tensors
+        def compute_loss(outputs, __):
+            if isinstance(outputs, (list, tuple)):
+                return ssl_loss_function(*outputs)
+
+            return ssl_loss_function(outputs)
+
+        return compute_loss
+
+    def get_loss_criterion(self):
+        """Return the loss criterion for SSL."""
+        # Get loss criterion for the subsequent training process
+        if self.current_round > Config().trainer.rounds:
+            loss_criterion_type = Config().algorithm.personalization.loss_criterion
+            loss_criterion_params = {}
+            if hasattr(Config().parameters.personalization, "loss_criterion"):
+                loss_criterion_params = (
+                    Config().parameters.personalization.loss_criterion._asdict()
+                )
+            return loss_criterion.get(
+                loss_criterion=loss_criterion_type,
+                loss_criterion_params=loss_criterion_params,
+            )
+
+        return self.get_ssl_criterion()
+
+    def get_lr_scheduler(self, config, optimizer):
+        # Get the lr scheduler for personalization
+        if self.current_round > Config().trainer.rounds:
+            lr_scheduler = Config().algorithm.personalization.lr_scheduler
+            lr_params = Config().parameters.personalization.learning_rate._asdict()
+
+            return lr_schedulers.get(
+                optimizer,
+                len(self.train_loader),
+                lr_scheduler=lr_scheduler,
+                lr_params=lr_params,
+            )
+        # Get the lr scheduler for SSL
+        return super().get_lr_scheduler(config, optimizer)
+
+    def train_run_start(self, config):
+        """Set the config before training."""
+        if self.current_round > Config().trainer.rounds:
+            # Set the config for the personalization
+            config["batch_size"] = Config().algorithm.personalization.batch_size
+            config["epochs"] = Config().algorithm.personalization.epochs
+
+            # Move the local layers to the device and set it to train mode
+            self.local_layers.to(self.device)
+            self.local_layers.train()
+
+    def perform_forward_and_backward_passes(self, config, examples, labels):
+        """Perform forward and backward passes in the training loop.
+        This function needs to reuse the optimization code of Plato as
+        during personalization, the encoder of the self.model will be used to
+        extract features into the local layers.
+        """
+
+        # Perform SSL training in the first `Config().trainer.rounds`` rounds
+        if not self.current_round > Config().trainer.rounds:
+            return super().perform_forward_and_backward_passes(config, examples, labels)
+
+        # Perform personalization after the final round
+        # Perform the local update on self.local_layers
+        self.optimizer.zero_grad()
+
+        # Use the trained encoder to output features.
+        # No optimizer for this basic encoder
+        features = self.model.encoder(examples)
+        outputs = self.local_layers(features)
+
+        loss = self._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()
+
+        self.optimizer.step()
+
+        return loss
+
+    def collect_encodings(self, data_loader):
+        """Collect encodings of the data by using self.model."""
+        samples_encoding = None
+        samples_label = None
+        self.model.eval()
+        self.model.to(self.device)
+        for examples, labels in data_loader:
+            examples, labels = examples.to(self.device), labels.to(self.device)
+            with torch.no_grad():
+                features = self.model.encoder(examples)
+                if samples_encoding is None:
+                    samples_encoding = features
+                else:
+                    samples_encoding = torch.cat([samples_encoding, features], dim=0)
+                if samples_label is None:
+                    samples_label = labels
+                else:
+                    samples_label = torch.cat([samples_label, labels], dim=0)
+
+        return samples_encoding, samples_label
+
+    def test_model(self, config, testset, sampler=None, **kwargs):
+        """Test the model to report the accuracy in each round."""
+        batch_size = config["batch_size"]
+        if self.current_round > Config().trainer.rounds:
+            # Test the personalized model after the final round.
+            self.local_layers.eval()
+            self.local_layers.to(self.device)
+
+            self.model.eval()
+            self.model.to(self.device)
+
+            test_loader = torch.utils.data.DataLoader(
+                testset, batch_size=batch_size, shuffle=False, sampler=sampler
+            )
+
+            correct = 0
+            total = 0
+            accuracy = 0
+            with torch.no_grad():
+                for examples, labels in test_loader:
+                    examples, labels = examples.to(self.device), labels.to(self.device)
+
+                    features = self.model.encoder(examples)
+                    outputs = self.local_layers(features)
+
+                    _, predicted = torch.max(outputs.data, 1)
+                    total += labels.size(0)
+                    correct += (predicted == labels).sum().item()
+
+            accuracy = correct / total
+
+            return accuracy
+        else:
+            # Test the personalized model in each round.
+
+            # For SSL, the way to test the trained model before personalization is
+            # to use the KNN as a classifier to evaluate the extracted features.
+
+            logging.info("[Client #%d] Testing the model with KNN.", self.client_id)
+
+            # Get the training loader and test loader
+            train_loader = torch.utils.data.DataLoader(
+                dataset=self.personalized_trainset,
+                shuffle=False,
+                batch_size=batch_size,
+                sampler=sampler,
+            )
+            test_loader = torch.utils.data.DataLoader(
+                testset, batch_size=batch_size, shuffle=False, sampler=sampler
+            )
+            # For evaluating self-supervised performance, we need to calculate
+            # distance between training samples and testing samples.
+            train_encodings, train_labels = self.collect_encodings(train_loader)
+            test_encodings, test_labels = self.collect_encodings(test_loader)
+
+            # Build KNN and perform the prediction
+            distances = torch.cdist(test_encodings, train_encodings, p=2)
+            knn = distances.topk(1, largest=False)
+            nearest_idx = knn.indices
+            predicted_labels = train_labels[nearest_idx].view(-1)
+            test_labels = test_labels.view(-1)
+
+            # Compute the accuracy
+            num_correct = torch.sum(predicted_labels == test_labels).item()
+            accuracy = num_correct / len(test_labels)
+
+            return accuracy

plato/trainers/split_learning.py

@@ -0,0 +1,305 @@
+"""
+A federated learning trainer using split learning.
+
+Reference:
+
+Vepakomma, et al., "Split Learning for Health: Distributed Deep Learning without Sharing
+Raw Patient Data," in Proc. AI for Social Good Workshop, affiliated with ICLR 2018.
+
+https://arxiv.org/pdf/1812.00564.pdf
+
+Chopra, Ayush, et al. "AdaSplit: Adaptive Trade-offs for Resource-constrained Distributed
+Deep Learning." arXiv preprint arXiv:2112.01637 (2021).
+
+https://arxiv.org/pdf/2112.01637.pdf
+"""
+
+import logging
+import os
+
+import torch
+from plato.config import Config
+
+from plato.trainers import basic
+from plato.datasources import feature
+from plato.samplers import all_inclusive
+
+
+# pylint:disable=too-many-instance-attributes
+class Trainer(basic.Trainer):
+    """The split learning trainer."""
+
+    def __init__(self, model=None, callbacks=None):
+        """Initializing the trainer with the provided model.
+
+        Arguments:
+        model: The model to train.
+        callbacks: The callbacks that this trainer uses.
+        """
+        super().__init__(model=model, callbacks=callbacks)
+        self.last_client_id = None
+        self.last_optimizer = None
+
+        # Client side variables
+        self.training_samples = None
+        self.gradients = None
+        self.data_loader = None
+
+        # Server side variables
+        self.cut_layer_grad = []
+
+    def get_train_loader(self, batch_size, trainset, sampler, **kwargs):
+        """
+        Creates an instance of the trainloader.
+
+        Arguments:
+        batch_size: the batch size.
+        trainset: the training dataset.
+        sampler: the sampler for the trainloader to use.
+        """
+        return trainset
+
+    def perform_forward_and_backward_passes(self, config, examples, labels):
+        """Perform forward and backward passes in the training loop.
+
+        Arguments:
+        config: the configuration.
+        examples: data samples in the current batch.
+        labels: labels in the current batch.
+
+        Returns: loss values after the current batch has been processed.
+        """
+        if self.client_id == 0:
+            return self._server_train_loop(config, examples, labels)
+
+        return self._client_train_loop(examples)
+
+    def train_run_end(self, config):
+        """Additional tasks after training."""
+        if self.client_id == 0:
+            # Server needs to save gradients, clients not
+            self.save_gradients(config)
+
+    def get_optimizer(self, model):
+        """Return the optimizer used in the last round to avoid reconfiguration."""
+        if self.last_optimizer is None or self.last_client_id != self.client_id:
+            self.last_optimizer = super().get_optimizer(model)
+            self.last_client_id = self.client_id
+
+        return self.last_optimizer
+
+    def get_train_samples(self, batch_size, trainset, sampler):
+        """
+        Get a batch of training samples to extract feature, the trainer has to save these
+        samples to complete training later.
+        """
+        data_loader = torch.utils.data.DataLoader(
+            dataset=trainset, shuffle=False, batch_size=batch_size, sampler=sampler
+        )
+        data_loader = iter(data_loader)
+        self.training_samples = next(data_loader)
+        # Wrap the training samples with datasource and sampler to be fed into Plato trainer
+        self.training_samples = self.process_training_samples_before_retrieving(
+            self.training_samples
+        )
+        return self.training_samples
+
+    def retrieve_train_samples(self):
+        """Retrieve the training samples to complete client training."""
+        samples = feature.DataSource([[self.training_samples]])
+        sampler = all_inclusive.Sampler(samples)
+
+        return samples, sampler
+
+    def load_gradients(self, gradients):
+        """Load the gradients which will be used to complete client training."""
+        self.gradients = gradients
+
+    def _client_train_loop(self, examples):
+        """Complete the client side training with gradients from server."""
+        self.optimizer.zero_grad()
+        examples, batch_size = self.process_samples_before_client_forwarding(examples)
+        outputs = self.model.forward_to(examples)
+
+        # Backpropagate with gradients from the server
+        gradients = self.gradients[0]
+        if gradients is None:
+            logging.warning("[Client #%d] Gradients from server is None.", os.getpid())
+        else:
+            gradients = gradients.to(self.device)
+            outputs.backward(gradients)
+            self.optimizer.step()
+
+        # No loss value on the client side
+        loss = torch.zeros(1)
+        self._loss_tracker.update(loss, batch_size)
+        return loss
+
+    def _server_train_loop(self, config, examples, labels):
+        """The training loop on the server."""
+        self.optimizer.zero_grad()
+        loss, grad, batch_size = self.server_forward_from((examples, labels), config)
+        loss = loss.cpu().detach()
+        self._loss_tracker.update(loss, batch_size)
+
+        # Record gradients within the cut layer
+        if grad is not None:
+            grad = grad.cpu().clone().detach()
+        self.cut_layer_grad = [grad]
+        self.optimizer.step()
+
+        logging.warning(
+            "[Server #%d] Gradients computed with training loss: %.4f",
+            os.getpid(),
+            loss,
+        )
+
+        return loss
+
+    def save_gradients(self, config):
+        """Server saves recorded gradients to a file."""
+        model_name = config["model_name"]
+        model_path = Config().params["model_path"]
+
+        if not os.path.exists(model_path):
+            os.makedirs(model_path)
+
+        if "/" in model_name:
+            model_name = model_name.replace("/", "_")
+
+        model_gradients_path = f"{model_path}/{model_name}_gradients.pth"
+        torch.save(self.cut_layer_grad, model_gradients_path)
+
+        logging.info(
+            "[Server #%d] Gradients saved to %s.", os.getpid(), model_gradients_path
+        )
+
+    def get_gradients(self):
+        """Read gradients from a file."""
+        model_path = Config().params["model_path"]
+        model_name = Config().trainer.model_name
+
+        if "/" in model_name:
+            model_name = model_name.replace("/", "_")
+
+        model_gradients_path = f"{model_path}/{model_name}_gradients.pth"
+        logging.info(
+            "[Server #%d] Loading gradients from %s.", os.getpid(), model_gradients_path
+        )
+
+        return torch.load(model_gradients_path)
+
+    def test_model(self, config, testset, sampler=None, **kwargs):
+        """
+        Evaluates the model with the provided test dataset and test sampler.
+
+        Arguments:
+        testset: the test dataset.
+        sampler: the test sampler. The default is None.
+        kwargs (optional): Additional keyword arguments.
+        """
+        batch_size = config["batch_size"]
+        accuracy = self.test_model_split_learning(batch_size, testset, sampler)
+        return accuracy
+
+    # API functions for split learning
+    def process_training_samples_before_retrieving(self, training_samples) -> ...:
+        """Process training samples before completing retrieving samples."""
+        return training_samples
+
+    def process_samples_before_client_forwarding(self, examples) -> ...:
+        """Process the examples before client conducting forwarding."""
+        return examples, examples.size(0)
+
+    # pylint:disable=unused-argument
+    def server_forward_from(self, batch, config) -> (..., ..., int):
+        """
+        The event for server completing training by forwarding from intermediate features.
+        Users may override this function for training different models with split learning.
+
+        Inputs:
+            batch: the batch of inputs for forwarding.
+            config: training configuration.
+        Returns:
+            loss: the calculated loss.
+            grad: the gradients over the intermediate feature.
+            batch_size: the batch size of the current sample.
+        """
+
+        inputs, target = batch
+        batch_size = inputs.size(0)
+        inputs = inputs.detach().requires_grad_(True)
+        outputs = self.model.forward_from(inputs)
+        loss = self._loss_criterion(outputs, target)
+        loss.backward()
+        grad = inputs.grad
+        return loss, grad, batch_size
+
+    def update_weights_before_cut(self, current_weights, weights) -> ...:
+        """
+        Update the weights before cut layer, called when testing accuracy in trainer.
+        Inputs:
+        current_weights: the current weights extracted by the algorithm.
+        weights: the weights to load.
+        Output:
+        current_weights: the updated current weights of the model.
+        """
+        cut_layer_idx = self.model.layers.index(self.model.cut_layer)
+
+        for i in range(0, cut_layer_idx):
+            weight_name = f"{self.model.layers[i]}.weight"
+            bias_name = f"{self.model.layers[i]}.bias"
+
+            if weight_name in current_weights:
+                current_weights[weight_name] = weights[weight_name]
+
+            if bias_name in current_weights:
+                current_weights[bias_name] = weights[bias_name]
+
+        return current_weights
+
+    def forward_to_intermediate_feature(self, inputs, targets) -> (..., ...):
+        """
+        The process to forward to get intermediate feature on the client.
+        Arguments:
+        inputs: the inputs for the model on the clients.
+        targets: the targets to get of the whole model.
+
+        Return:
+        outputs: the intermediate feature.
+        targets: the targets to get of the whole model.
+        """
+        with torch.no_grad():
+            logits = self.model.forward_to(inputs)
+
+        outputs = logits.detach().cpu()
+        targets = targets.detach().cpu()
+        return outputs, targets
+
+    def test_model_split_learning(self, batch_size, testset, sampler=None) -> ...:
+        """
+        The test model process for split learning.
+
+        Returns:
+        accuracy: the metrics for evaluating the model.
+        """
+        test_loader = torch.utils.data.DataLoader(
+            testset, batch_size=batch_size, shuffle=False, sampler=sampler
+        )
+        correct = 0
+        total = 0
+
+        self.model.to(self.device)
+        with torch.no_grad():
+            for examples, labels in test_loader:
+                examples, labels = examples.to(self.device), labels.to(self.device)
+
+                outputs = self.model(examples)
+
+                outputs = self.process_outputs(outputs)
+
+                _, predicted = torch.max(outputs.data, 1)
+                total += labels.size(0)
+                correct += (predicted == labels).sum().item()
+
+        return correct / total

plato/trainers/tracking.py

@@ -0,0 +1,96 @@
+"""
+Keeping a history of metrics during the training run.
+"""
+
+from collections import defaultdict
+from typing import Iterable
+
+
+class RunHistory:
+    """
+    A simple history of metrics during a training or evaluation run.
+    """
+
+    def __init__(self):
+        self._metrics = defaultdict(list)
+
+    def get_metric_names(self) -> Iterable:
+        """
+        Returns an iterable set containing of all unique metric names which are
+        being tracked.
+
+        :return: an iterable of the unique metric names.
+        """
+        return set(self._metrics.keys())
+
+    def get_metric_values(self, metric_name) -> Iterable:
+        """
+        Returns an ordered iterable list of values that has been stored since
+        the last reset corresponding to the provided metric name.
+
+        :param metric_name: the name of the metric being tracked.
+        :return: an ordered iterable of values that have been recorded for that metric.
+        """
+        return self._metrics[metric_name]
+
+    def get_latest_metric(self, metric_name):
+        """
+        Returns the most recent value that has been recorded for the given metric.
+
+        :param metric_name: the name of the metric being tracked.
+        :return: the last recorded value.
+        """
+        if len(self._metrics[metric_name]) > 0:
+            return self._metrics[metric_name][-1]
+        else:
+            raise ValueError(
+                f"No values have been recorded for the metric {metric_name}"
+            )
+
+    def update_metric(self, metric_name, metric_value):
+        """
+        Records a new value for the given metric.
+
+        :param metric_name: the name of the metric being tracked.
+        :param metric_value: the value to record.
+        """
+        self._metrics[metric_name].append(metric_value)
+
+    def reset(self):
+        """
+        Resets the state of the :class:`RunHistory`.
+
+        """
+        self._metrics = defaultdict(list)
+
+
+class LossTracker:
+    """A simple tracker for computing the average loss."""
+
+    def __init__(self):
+        self.loss_value = 0
+        self._average = 0
+        self.total_loss = 0
+        self.running_count = 0
+
+    def reset(self):
+        """Resets this loss tracker."""
+
+        self.loss_value = 0
+        self._average = 0
+        self.total_loss = 0
+        self.running_count = 0
+
+    def update(self, loss_batch_value, batch_size=1):
+        """Updates the loss tracker with another loss value from a batch."""
+
+        self.loss_value = loss_batch_value
+        self.total_loss += loss_batch_value * batch_size
+        self.running_count += batch_size
+        self._average = self.total_loss / self.running_count
+
+    @property
+    def average(self):
+        """Returns the computed average of loss values tracked."""
+
+        return self._average.cpu().detach().mean().item()