Add LeNet, Add Sharding

src/decentralizepy/datasets/CIFAR10.py

 import logging
-import os
 
-import numpy as np
 import torch
 import torch.nn.functional as F
 import torchvision
@@ -9,9 +7,8 @@ import torchvision.transforms as transforms
 from torch import nn
 from torch.utils.data import DataLoader
 
-from decentralizepy.datasets.Data import Data
 from decentralizepy.datasets.Dataset import Dataset
-from decentralizepy.datasets.Partitioner import DataPartitioner, SimpleDataPartitioner
+from decentralizepy.datasets.Partitioner import DataPartitioner, KShardDataPartitioner
 from decentralizepy.mappings.Mapping import Mapping
 from decentralizepy.models.Model import Model
 
@@ -53,9 +50,9 @@ class CIFAR10(Dataset):
             all_trainset = []
             for y, x in train_data.items():
                 all_trainset.extend([(a, y) for a in x])
-            self.trainset = SimpleDataPartitioner(all_trainset, self.sizes).use(
-                self.uid
-            )
+            self.trainset = KShardDataPartitioner(
+                all_trainset, self.sizes, shards=self.shards
+            ).use(self.uid)
 
     def load_testset(self):
         """
@@ -79,6 +76,7 @@ class CIFAR10(Dataset):
         sizes="",
         test_batch_size=1024,
         partition_niid=False,
+        shards=1,
     ):
         """
         Constructor which reads the data files, instantiates and partitions the dataset
@@ -117,6 +115,7 @@ class CIFAR10(Dataset):
         )
 
         self.partition_niid = partition_niid
+        self.shards = shards
         self.transform = transforms.Compose(
             [
                 transforms.ToTensor(),
@@ -274,3 +273,49 @@ class CNN(Model):
         x = F.relu(self.fc2(x))
         x = self.fc3(x)
         return x
+
+
+class LeNet(Model):
+    """
+    Class for a LeNet Model for CIFAR10
+    Inspired by original LeNet network for MNIST: https://ieeexplore.ieee.org/abstract/document/726791
+
+    """
+
+    def __init__(self):
+        """
+        Constructor. Instantiates the CNN Model
+            with 10 output classes
+
+        """
+        super().__init__()
+        self.conv1 = nn.Conv2d(3, 32, 5, padding="same")
+        self.pool = nn.MaxPool2d(2, 2)
+        self.gn1 = nn.GroupNorm(2, 32)
+        self.conv2 = nn.Conv2d(32, 32, 5, padding="same")
+        self.gn2 = nn.GroupNorm(2, 32)
+        self.conv3 = nn.Conv2d(32, 64, 5, padding="same")
+        self.gn3 = nn.GroupNorm(2, 64)
+        self.fc1 = nn.Linear(64 * 4 * 4, NUM_CLASSES)
+
+    def forward(self, x):
+        """
+        Forward pass of the model
+
+        Parameters
+        ----------
+        x : torch.tensor
+            The input torch tensor
+
+        Returns
+        -------
+        torch.tensor
+            The output torch tensor
+
+        """
+        x = self.pool(F.relu(self.gn1(self.conv1(x))))
+        x = self.pool(F.relu(self.gn2(self.conv2(x))))
+        x = self.pool(F.relu(self.gn3(self.conv3(x))))
+        x = torch.flatten(x, 1)
+        x = self.fc1(x)
+        return x

src/decentralizepy/datasets/Partitioner.py

@@ -131,3 +131,50 @@ class SimpleDataPartitioner(DataPartitioner):
             part_len = int(frac * data_len)
             self.partitions.append(indexes[0:part_len])
             indexes = indexes[part_len:]
+
+
+class KShardDataPartitioner(DataPartitioner):
+    """
+    Class to partition the dataset
+
+    """
+
+    def __init__(self, data, sizes=[1.0], shards=1, seed=1234):
+        """
+        Constructor. Partitions the data according the parameters
+
+        Parameters
+        ----------
+        data : indexable
+            An indexable list of data items
+        sizes : list(float)
+            A list of fractions for each process
+        shards : int
+            Number of shards to allot to process
+        seed : int, optional
+            Seed for generating a random subset
+
+        """
+        self.data = data
+        self.partitions = []
+        data_len = len(data)
+        indexes = [x for x in range(0, data_len)]
+        rng = Random()
+        rng.seed(seed)
+
+        for frac in sizes:
+            self.partitions.append([])
+            for _ in range(shards):
+                start = rng.randint(0, len(indexes) - 1)
+                part_len = int(frac * data_len) // shards
+                if start + part_len > len(indexes):
+                    self.partitions[-1].extend(indexes[start:])
+                    self.partitions[-1].extend(
+                        indexes[: (start + part_len - len(indexes))]
+                    )
+                    indexes = indexes[(start + part_len - len(indexes)) : start]
+                else:
+                    self.partitions[-1].extend(indexes[start : start + part_len])
+                    index_start = indexes[:start]
+                    index_start.extend(indexes[start + part_len :])
+                    indexes = index_start