plato-learn 1.1__py3-none-any.whl

plato/trainers/yolov8.py

@@ -0,0 +1,41 @@
+"""The YOLOV8 model for PyTorch."""
+
+import logging
+
+from plato.config import Config
+from plato.trainers import basic
+
+
+class Trainer(basic.Trainer):
+    """The YOLOV8 trainer."""
+
+    # pylint: disable=unused-argument
+    def train_model(self, config, trainset, sampler, **kwargs):
+        """The training loop for YOLOv8.
+
+        Arguments:
+        config: A dictionary of configuration parameters.
+        trainset: The training dataset.
+        """
+        self.model.train(
+            data=Config().data.data_params,
+            epochs=Config().trainer.epochs,
+        )
+
+        self.train_run_end(config)
+        self.callback_handler.call_event("on_train_run_end", self, config)
+
+    def test_model(self, config, testset, sampler=None, **kwargs):
+        """The test loop for YOLOv8.
+
+        Arguments:
+        config: A dictionary of configuration parameters.
+        testset: The test dataset.
+        """
+
+        logging.info("[%s] Started model testing.", self)
+        metrics = self.model.val(
+            data=Config().data.data_params,
+        )
+
+        return metrics.box.map50

plato/utils/count_parameters.py

@@ -0,0 +1,30 @@
+from prettytable import PrettyTable
+import torch
+
+
+def count_parameters(model):
+    table = PrettyTable(["Modules", "Parameters"])
+    total_params = 0
+    for name, parameter in model.named_parameters():
+        if not parameter.requires_grad:
+            continue
+        params = parameter.numel()
+        table.add_row([name, params])
+        total_params += params
+    print(table)
+    print(f"Total Trainable Params: {total_params}")
+    return total_params
+
+
+resnet18 = torch.hub.load("pytorch/vision:v0.10.0", "resnet18", pretrained=True)
+mobilenet = torch.hub.load("pytorch/vision:v0.10.0", "mobilenet_v2", pretrained=True)
+alexnet = torch.hub.load("pytorch/vision:v0.10.0", "alexnet", pretrained=True)
+
+print("The size of ResNet-18:")
+count_parameters(resnet18)
+
+print("\nThe size of MobileNet:")
+count_parameters(mobilenet)
+
+print("\nThe size of AlexNet:")
+count_parameters(alexnet)

plato/utils/csv_processor.py

@@ -0,0 +1,26 @@
+"""
+Utility functions that write results into a CSV file.
+"""
+
+import csv
+import os
+from typing import List
+
+
+def initialize_csv(result_csv_file: str, logged_items: List, result_path: str) -> None:
+    """Create a CSV file and writer the first row."""
+    # Create a new directory if it does not exist
+    if not os.path.exists(result_path):
+        os.makedirs(result_path)
+
+    with open(result_csv_file, "w", encoding="utf-8") as result_file:
+        result_writer = csv.writer(result_file)
+        header_row = logged_items
+        result_writer.writerow(header_row)
+
+
+def write_csv(result_csv_file: str, new_row: List) -> None:
+    """Write the results of current round."""
+    with open(result_csv_file, "a", encoding="utf-8") as result_file:
+        result_writer = csv.writer(result_file)
+        result_writer.writerow(new_row)

plato/utils/data_loaders.py

@@ -0,0 +1,148 @@
+"""
+The implementation of various wrappers to support flexible combinations of data loaders.
+
+Two types of data loader wrappers are supported:
+
+    - ParallelDataLoader
+    - SequentialDataLoader
+
+One specific utilization condition is self-supervised learning, where datasets,
+such as STL10, contains trainsets with and without labels, and the desired data
+loader first loads the trainset with labels and then the one without labels.
+We can use SequentialDataLoader for this purpose.
+
+"""
+
+import numpy as np
+
+
+class ParallelIterator:
+    """An iterator to support iterating along each data loader simultaneously to generate
+    one batch."""
+
+    def __init__(self, defined_compound_loader):
+        self.defined_compound_loader = defined_compound_loader
+        self.compound_loaders = self.defined_compound_loader.loaders
+        self.loader_iters = [iter(loader) for loader in self.compound_loaders]
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        # When the shortest loader (the one with minimum number of batches)
+        # terminates, this iterator will terminates.
+        # The `StopIteration` raised inside that shortest loader's `__next__`
+        # method will in turn gets out of this `__next__` method.
+        batches = [next(loader_iter) for loader_iter in self.loader_iters]
+        return self.defined_compound_loader.combine_batch(batches)
+
+    def __len__(self):
+        return len(self.defined_compound_loader)
+
+
+class ParallelDataLoader:
+    """This class wraps several pytorch DataLoader objects, allowing each time
+    taking a batch from each of them and then combining these several batches
+    into one. This class mimics the `for batch in loader:` interface of
+    pytorch `DataLoader`.
+
+    :param defined_loaders: a list or tuple of pytorch DataLoader objects
+
+    [For example]
+    There are two dataloaders A and B.
+    With ParallelDataLoader, one iter will extract one batch of samples 'A_b'
+    from A and one batch of samples 'B_b' from B. Thus, the loaded term is a
+    list containing [A_b, B_b].
+
+    The size of this dataloader equals to the minimum length of the dataloader
+    within input defined loaders.
+    """
+
+    def __init__(self, defined_loaders):
+        self.loaders = [loader for loader in defined_loaders if loader is not None]
+
+    def __iter__(self):
+        return ParallelIterator(self)
+
+    def __len__(self):
+        return min(len(loader) for loader in self.loaders)
+
+    def combine_batch(self, batches):
+        """Customize the behavior of combining batches here."""
+        return batches
+
+
+class SequentialIterator:
+    """An iterator to support iterating through each data loader sequentially.
+
+    For example, there are three data loaders, A, B, and, C:
+        the iteration will start from A, once A finished, B will start; then C will start.
+
+    Thus, the length of this iterator is:
+        len(A) + len(B) + len(C)
+    """
+
+    def __init__(self, defined_compound_loader):
+        # only utilize the vaild loaders
+
+        self.defined_compound_loader = defined_compound_loader
+        self.compound_loaders = self.defined_compound_loader.loaders
+        self.loader_iters = [iter(loader) for loader in self.compound_loaders]
+
+        self.loaders_len = [len(loader) for loader in self.compound_loaders]
+        self.loaders_batch_bound = np.cumsum(self.loaders_len, axis=0)
+
+        self.num_loaders = len(self.loaders_len)
+        self.batch_idx = 0
+
+    def __iter__(self):
+        return self
+
+    def __next__(self):
+        # When the final loader (the last loader in the input loaders)
+        # terminates, this iterator will terminates.
+        # The `StopIteration` raised inside that shortest loader's `__next__`
+        # method will in turn gets out of this `__next__` method.
+        cur_loader_idx = np.digitize(self.batch_idx, self.loaders_batch_bound)
+
+        # if completed the final loader, we just recycle to the final loader
+        # then, this loader will be terminated because:
+        # The `StopIteration` raised inside that final loader's `__next__`
+        if cur_loader_idx == self.num_loaders:
+            cur_loader_idx -= 1
+
+        loader_iter = self.loader_iters[cur_loader_idx]
+        batch = next(loader_iter)
+
+        self.batch_idx += 1
+
+        return self.defined_compound_loader.process_batch(batch)
+
+    def __len__(self):
+        return len(self.target_loader)
+
+
+class SequentialDataLoader:
+    """This class wraps several pytorch DataLoader objects, allowing each time
+    taking a batch from each of them and then combining these several batches
+    into one. This class mimics the `for batch in loader:` interface of
+    pytorch `DataLoader`.
+
+    :param defined_loaders: A list or tuple containing pytorch DataLoader objects
+
+    The size of this dataloader equals to the minimum length of the dataloader
+    within input defined loaders.
+    """
+
+    def __init__(self, defined_loaders):
+        self.loaders = [loader for loader in defined_loaders if loader is not None]
+
+    def __iter__(self):
+        return SequentialIterator(self)
+
+    def __len__(self):
+        return sum(len(loader) for loader in self.loaders)
+
+    def process_batch(self, batch):
+        """Customize the behavior of combining batches here."""
+        return batch

plato/utils/decorators.py

@@ -0,0 +1,24 @@
+"""Useful decorators."""
+
+import time
+from functools import wraps
+
+
+def timeit(func_timed):
+    """Measures the time elapsed for a particular function 'func_timed'."""
+
+    @wraps(func_timed)
+    def timed(*args, **kwargs):
+        started = time.perf_counter()
+        output = func_timed(*args, **kwargs)
+        ended = time.perf_counter()
+        elapsed = ended - started
+        print(
+            '"{}" took {:.2f} seconds to execute.'.format(func_timed.__name__, elapsed)
+        )
+        if output is None:
+            return elapsed
+        else:
+            return output, elapsed
+
+    return timed

plato/utils/fonts.py

@@ -0,0 +1,23 @@
+"""
+Colours and fonts for logging messages
+"""
+
+
+def colourize(message, colour="yellow", style="bold"):
+    """Returns the message in input colour and style"""
+    reset = "\033[0m"
+    colours = {
+        "green": "\033[92m",
+        "blue": "\033[94m",
+        "yellow": "\033[93m",
+        "red": "\033[91m",
+    }
+    styles = {"standard": "", "bold": "\033[1m", "underline": "\033[4m"}
+
+    if not (colour in colours and style in styles):
+        raise ValueError(
+            f"Your colour '{colour}' or your style '{style}' is not supported."
+            f"\nThe supported colours are: {', '.join(colours)}. \nThe supported styles are: {', '.join(styles)}."
+        )
+
+    return colours[colour] + styles[style] + message + reset

plato/utils/homo_enc.py

@@ -0,0 +1,187 @@
+"""
+Utility functions for homomorphric encryption with TenSEAL.
+"""
+
+import os
+import pickle
+import zlib
+from typing import OrderedDict
+
+import numpy as np
+import tenseal as ts
+import torch
+
+
+def get_ckks_context():
+    """Obtain a TenSEAL context for encryption and decryption."""
+    context_dir = ".ckks_context/"
+    context_name = "context"
+    try:
+        with open(os.path.join(context_dir, context_name), "rb") as f:
+            return ts.context_from(f.read())
+    except FileNotFoundError:
+        # Create a new context if it does not exist
+        if not os.path.exists(context_dir):
+            os.mkdir(context_dir)
+
+        context = ts.context(
+            ts.SCHEME_TYPE.CKKS,
+            poly_modulus_degree=8192,
+            coeff_mod_bit_sizes=[60, 40, 40, 60],
+        )
+        context.global_scale = 2**40
+
+        with open(os.path.join(context_dir, context_name), "wb") as f:
+            f.write(context.serialize(save_secret_key=True))
+            f.close()
+
+        return context
+
+
+def encrypt_weights(
+    plain_weights,
+    serialize=True,
+    context=None,
+    indices=None,
+):
+    """Flatten the model weights and encrypt the selected ones."""
+    assert context is not None
+
+    # Step 1: flatten all weight tensors to a vector
+    weights_vector = np.array([])
+    for weight in plain_weights.values():
+        weights_vector = np.append(weights_vector, weight)
+
+    # Step 2: set up the indices for encrypted weights
+    encrypt_indices = None
+    if indices is None:
+        encrypt_indices = np.arange(len(weights_vector)).tolist()
+    else:
+        encrypt_indices = indices
+    encrypt_indices.sort()
+
+    # Step 3: separate weights into encrypted and unencrypted ones
+    unencrypted_weights = np.delete(weights_vector, encrypt_indices)
+    weights_to_enc = weights_vector[encrypt_indices]
+
+    if len(weights_to_enc) == 0:
+        encrypted_weights = None
+    else:
+        encrypted_weights = _encrypt(weights_to_enc, context, serialize)
+
+    # Finish by wrapping up the information
+    output = wrap_encrypted_model(
+        unencrypted_weights, encrypted_weights, encrypt_indices
+    )
+
+    return output
+
+
+def _encrypt(data_vector, context, serialize=True):
+    if serialize:
+        return ts.ckks_vector(context, data_vector).serialize()
+    else:
+        return ts.ckks_vector(context, data_vector)
+
+
+def deserialize_weights(serialized_weights, context):
+    """Deserialize the encrypted weights (not decrypted yet)."""
+    deserialized_weights = OrderedDict()
+    for name, weight in serialized_weights.items():
+        if name == "encrypted_weights" and weight is not None:
+            deser_weights_vector = ts.lazy_ckks_vector_from(weight)
+            deser_weights_vector.link_context(context)
+            deserialized_weights[name] = deser_weights_vector
+        else:
+            deserialized_weights[name] = weight
+
+    return deserialized_weights
+
+
+def decrypt_weights(data, weight_shapes=None, para_nums=None):
+    """Decrypt the vector and restore model weights according to the shapes."""
+    vector_length = []
+    for para_num in para_nums.values():
+        vector_length.append(para_num)
+
+    # Step 1: decrypt the encrypted weights
+    plaintext_weights_vector = None
+    unencrypted_weights, encrypted_weights, indices = extract_encrypted_model(data)
+
+    if len(indices) != 0:
+        decrypted_vector = np.array(encrypted_weights.decrypt())
+
+        vector_size = len(unencrypted_weights) + len(indices)
+        plaintext_weights_vector = np.zeros(vector_size)
+        plaintext_weights_vector[indices] = decrypted_vector
+
+        unencrypted_indices = np.delete(range(vector_size), indices)
+        plaintext_weights_vector[unencrypted_indices] = unencrypted_weights
+    else:
+        plaintext_weights_vector = unencrypted_weights
+
+    # Step 2: rebuild the original weight vector
+    decrypted_weights = OrderedDict()
+    plaintext_weights_vector = np.split(
+        plaintext_weights_vector, np.cumsum(vector_length)
+    )[:-1]
+    weight_index = 0
+
+    for name, shape in weight_shapes.items():
+        decrypted_weights[name] = plaintext_weights_vector[weight_index].reshape(shape)
+        try:
+            decrypted_weights[name] = torch.from_numpy(decrypted_weights[name])
+        except Exception:
+            # PyTorch does not exist, just return numpy array and handle it somewhere else.
+            decrypted_weights[name] = decrypted_weights[name]
+        weight_index = weight_index + 1
+
+    return decrypted_weights
+
+
+def wrap_encrypted_model(unencrypted_weights, encrypted_weights, indices):
+    """Wrap up the encrypted model in a dict as the message between server and client."""
+    message = {
+        "unencrypted_weights": unencrypted_weights,
+        "encrypted_weights": encrypted_weights,
+        "indices": indices,
+    }
+
+    return message
+
+
+def extract_encrypted_model(data):
+    """Extract infromation from the message of encrytped model"""
+    unencrypted_weights = data["unencrypted_weights"]
+    encrypted_weights = data["encrypted_weights"]
+    indices = data["indices"]
+
+    return unencrypted_weights, encrypted_weights, indices
+
+
+def indices_to_bitmap(indices):
+    """Turn a list of indices into a bitmap."""
+    if indices == []:
+        # In case of empty list
+        return indices
+    bit_array = np.zeros(np.max(indices) + 1, dtype=np.int8)
+    bit_array[indices] = 1
+    bitmap = np.packbits(bit_array)
+
+    # Compress the bitmap before sending it out
+    compressed_bitmap = zlib.compress(pickle.dumps(bitmap))
+
+    return compressed_bitmap
+
+
+def bitmap_to_indices(bitmap):
+    """Translate a bitmap back to a list of indices."""
+    if bitmap == []:
+        # In case of empty list
+        return bitmap
+
+    decompressed_bitmap = pickle.loads(zlib.decompress(bitmap))
+    bit_array = np.unpackbits(decompressed_bitmap)
+    indices = np.where(bit_array == 1)[0].tolist()
+
+    return indices

plato/utils/reinforcement_learning/policies/base.py

@@ -0,0 +1,161 @@
+import copy
+import logging
+import os
+import random
+from abc import ABC, abstractmethod
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from plato.config import Config
+from torch import nn
+
+
+class ReplayMemory:
+    """A simple example of replay memory buffer."""
+
+    def __init__(self, state_dim, action_dim, capacity, seed):
+        random.seed(seed)
+        self.device = Config().device()
+        self.capacity = int(capacity)
+        self.ptr = 0
+        self.size = 0
+
+        self.state = np.zeros((self.capacity, state_dim))
+        self.action = np.zeros((self.capacity, action_dim))
+        self.reward = np.zeros((self.capacity, 1))
+        self.next_state = np.zeros((self.capacity, state_dim))
+        self.done = np.zeros((self.capacity, 1))
+
+    def push(self, data):
+        self.state[self.ptr] = data[0]
+        self.action[self.ptr] = data[1]
+        self.reward[self.ptr] = data[2]
+        self.next_state[self.ptr] = data[3]
+        self.done[self.ptr] = data[4]
+
+        self.ptr = (self.ptr + 1) % self.capacity
+        self.size = min(self.size + 1, self.capacity)
+
+    def sample(self):
+        ind = np.random.randint(0, self.size, size=int(Config().algorithm.batch_size))
+
+        state = self.state[ind]
+        action = self.action[ind]
+        reward = self.reward[ind]
+        next_state = self.next_state[ind]
+        done = self.done[ind]
+
+        return state, action, reward, next_state, done
+
+    def __len__(self):
+        return self.size
+
+
+class Actor(nn.Module):
+    def __init__(self, state_dim, action_dim, max_action):
+        super(Actor, self).__init__()
+
+        self.l1 = nn.Linear(state_dim, 400)
+        self.l2 = nn.Linear(400, 300)
+        self.l3 = nn.Linear(300, action_dim)
+
+        self.max_action = max_action
+
+    def forward(self, x):
+        x = F.relu(self.l1(x))
+        x = F.relu(self.l2(x))
+        x = self.max_action * torch.tanh(self.l3(x))
+        return x
+
+
+class Critic(nn.Module):
+    def __init__(self, state_dim, action_dim):
+        super(Critic, self).__init__()
+
+        self.l1 = nn.Linear(state_dim + action_dim, 400)
+        self.l2 = nn.Linear(400, 300)
+        self.l3 = nn.Linear(300, 1)
+
+    def forward(self, x, u):
+        x = F.relu(self.l1(torch.cat([x, u], 1)))
+        x = F.relu(self.l2(x))
+        x = self.l3(x)
+        return x
+
+
+class Policy(ABC):
+    """A simple example of DRL policy."""
+
+    def __init__(self, state_dim, action_dim):
+        self.max_action = Config().algorithm.max_action
+        self.device = Config().device()
+        self.total_it = 0
+
+        # Initialize NNs
+        self.actor = Actor(state_dim, action_dim, self.max_action).to(self.device)
+        self.actor_target = copy.deepcopy(self.actor)
+        self.actor_optimizer = torch.optim.Adam(
+            self.actor.parameters(), lr=Config().algorithm.learning_rate
+        )
+
+        self.critic = Critic(state_dim, action_dim).to(self.device)
+        self.critic_target = copy.deepcopy(self.critic)
+        self.critic_optimizer = torch.optim.Adam(
+            self.critic.parameters(), lr=Config().algorithm.learning_rate
+        )
+        # Initialize replay memory
+        self.replay_buffer = ReplayMemory(
+            state_dim,
+            action_dim,
+            Config().algorithm.replay_size,
+            Config().algorithm.replay_seed,
+        )
+
+    def save_model(self, ep=None):
+        """Saving the model to a file."""
+        model_name = Config().algorithm.model_name
+        model_path = f"./models/{model_name}/"
+        if not os.path.exists(model_path):
+            os.makedirs(model_path)
+        if ep is not None:
+            model_path += "iter" + str(ep) + "_"
+
+        torch.save(self.actor.state_dict(), model_path + "actor.pth")
+        torch.save(
+            self.actor_optimizer.state_dict(), model_path + "actor_optimizer.pth"
+        )
+        torch.save(self.critic.state_dict(), model_path + "critic.pth")
+        torch.save(
+            self.critic_optimizer.state_dict(), model_path + "critic_optimizer.pth"
+        )
+
+        logging.info("[RL Agent] Model saved to %s.", model_path)
+
+    def load_model(self, ep=None):
+        """Loading pre-trained model weights from a file."""
+        model_name = Config().algorithm.model_name
+        model_path = f"./models/{model_name}/"
+        if ep is not None:
+            model_path += "iter" + str(ep) + "_"
+
+        logging.info("[RL Agent] Loading a model from %s.", model_path)
+
+        self.actor.load_state_dict(torch.load(model_path + "actor.pth"))
+        self.actor_optimizer.load_state_dict(
+            torch.load(model_path + "actor_optimizer.pth")
+        )
+        self.critic.load_state_dict(torch.load(model_path + "critic.pth"))
+        self.critic_optimizer.load_state_dict(
+            torch.load(model_path + "critic_optimizer.pth")
+        )
+
+    @abstractmethod
+    def select_action(self, state, hidden=None, test=False):
+        """Select action from policy."""
+        raise NotImplementedError("Please Implement this method")
+
+    @abstractmethod
+    def update(self):
+        """Update policy."""
+        raise NotImplementedError("Please Implement this method")