Federated Learning

import torch

def model_to_vector(model):
  """
  Convert a PyTorch model to a vector of its weights.

  Args:
    model: A PyTorch model.

  Returns:
    A flattened vector of the model's weights.
  """
  vector = []
  for param in model.parameters():
    vector.append(param.data.view(-1))
  return torch.cat(vector)


def vector_to_model(vector, model):
  """
  Convert a vector of weights back to a PyTorch model.

  Args:
    vector: A flattened vector of weights.
    model: A PyTorch model (with the same architecture as the model used to create the vector).

  Returns:
    The PyTorch model with its weights updated from the vector.
  """
  offset = 0
  for param in model.parameters():
    param_size = param.data.numel()
    param.data = vector[offset:offset + param_size].view(param.data.shape)
    offset += param_size
  return model

# let's create some test code
from torchvision.models import resnet18

model = resnet18()
vector = model_to_vector(model)
print(vector.shape)
model_recon = vector_to_model(vector, model)
x = torch.rand(1, 3, 224, 224)
y = model_recon(x)
print(y.shape)

import matplotlib.pyplot as plt
plt.imshow(x[0].permute(1,2,0))

import matplotlib.pyplot as plt
# download the image
!wget https://www.stockvault.net/data/2011/02/06/117454/thumb16.jpg -O test.jpg

img = plt.imread('test.jpg')
print(img.shape)
print("Image height: {}, Image width: {}, #channels: {}".format(img.shape[0],img.shape[1],img.shape[2]))
print("Range: {}-{}".format(img.min(), img.max()))
plt.imshow(img)

# Let's start with a model
from torch import nn

class LogisiticRegression(nn.Module):
  def __init__(self, input_dim, output_dim, device = "cuda"):
    super(LogisiticRegression, self).__init__()
    self.linear = nn.Linear(input_dim, output_dim) # W,b Wx+b 784 -> 100
    self.device = device # GPU
    self.to(device)# pass the param to gpu

  def forward(self, x): #y = f(x)
    x = x.to(self.device) # pass the input to gpu
    x = nn.Flatten()(x) # make x into a 1d vector
    out = self.linear(x) # pass it through a linear layer, i.e. Wx+b

    return out # softmax will be applied later by loss function

def test_model():
  model = LogisiticRegression(784, 10) #784 = 28*28*1
  x = torch.randn(4, 28, 28)
  out = model(x)
  return out
  # print(out)


y = test_model()
print(y, y.shape)

import torch
from torchvision import datasets, transforms
from torch.utils.data import random_split, DataLoader

# Load MNIST dataset
train_dataset = datasets.MNIST(root='./data', train=True, download=True, transform=transforms.ToTensor())

# Calculate the size of each partition
partition_size = len(train_dataset) // 10

# Split the dataset into 10 non-intersecting partitions
PARTITIONS = random_split(train_dataset, [partition_size] * 10 + [len(train_dataset) - partition_size * 10])[:-1]

# Print the size of each partition
for i, partition in enumerate(PARTITIONS):
  print(f"Partition {i+1}: {len(PARTITIONS[i])} samples")

from __future__ import annotations
# client class
class Client:
  def __init__(self, client_id, local_batch_size, local_iteration, device="cuda"):
    self.client_id = client_id
    self.local_bs = local_batch_size
    self.local_iter = local_iteration
    self.device = device
    self.data_loader = DataLoader(PARTITIONS[client_id], self.local_bs, shuffle = True)
    self.model = LogisiticRegression(784, 10, device)
    self.optimizer = torch.optim.SGD(self.model.parameters(), lr=0.1)
    self.loss_fn = nn.CrossEntropyLoss().to(device)
    self.losses = []

  def connect(self, server: Server):
    server.append_(self)

  # method to reconstruct the model from weight update
  def receive_update(self, update: torch.Tensor):
    self.model = vector_to_model(update, self.model)

  # method to locally train the data for local_iter round
  def local_train(self) -> None:
    iter = 0
    for batch_idx, (image, label) in enumerate(self.data_loader):
      self.optimizer.zero_grad()
      label = label.to(self.device)
      output = self.model(image).to(self.device)
      loss = self.loss_fn(output, label)
      loss.backward()
      self.optimizer.step()
      iter += 1
      if iter == self.local_iter:
        local_loss = loss.item()
        self.losses.append(local_loss)
        break

  def send_local_model(self)->torch.Tensor:
    return model_to_vector(self.model)

from tqdm import tqdm
class Server:
  def __init__(self, rounds, device="cuda"):
    self.device = device
    self.model = LogisiticRegression(784, 10, device)
    self.clients = []
    self.rounds = rounds

  def append_(self, client: Client):
    self.clients.append(client)

  def train_round_(self):
    weights = []
    for client in self.clients:
      client.local_train()
      weights.append(client.send_local_model())
    new_update = self.aggregate(weights)
    for client in self.clients:
      client.receive_update(new_update)
    return new_update

  def aggregate(self, weights: torch.Tensor):
    return torch.mean(torch.stack(weights), dim=0)

  def train(self):
    for i in tqdm(range(self.rounds)):
      last_update= self.train_round_()
    self.model = vector_to_model(last_update, self.model)

server = Server(100)
for i in range(10):
  client = Client(client_id = i, local_batch_size = 32, local_iteration = 4)
  client.connect(server)
server.train()

import matplotlib.pyplot as plt
# plot clients loss progress
for client in server.clients:
  plt.plot(client.losses)
plt.title("Clients Loss Progress")
plt.xlabel("Local Iterations")
plt.ylabel("Loss")
plt.show()

import random
# visualize an example
test_mnist = datasets.MNIST(root='./data', train=False, download=True, transform=transforms.ToTensor())
test_loader = torch.utils.data.DataLoader(dataset=test_mnist, batch_size=1, shuffle=False)

index = random.randint(0, len(test_mnist))
image, label = test_mnist[index]
prediction = server.model(image.to("cuda")).argmax().item()
print(f"Label: {label}, prediction: {prediction}")
plt.imshow(image.squeeze(), cmap='gray')
plt.show()

def accuracy(model):
  model.eval()
  correct = 0
  total = 0
  for images, labels in test_loader:
    output = model(images.to("cuda"))
    labels = labels.to("cuda")
    _, predicted = torch.max(output.data, 1)
    total += labels.size(0)
    correct += (predicted == labels).sum().item()
  return 100 * correct / total
acc = accuracy(server.model)

print("Final model testing accuracy: {}%".format(acc))