viceroy

git clone git://git.codymlewis.com/viceroy.git
Log | Files | Refs | README

commit cdb1da6482162028a53af7b8b6a12752e62233ea
parent 3015c5ccf62fb988601bc9d8de5c183cd7b48ec2
Author: Cody Lewis <cody@codymlewis.com>
Date:   Tue, 20 Oct 2020 17:01:05 +1100

Removed sockets due to limitations

Diffstat:
MClient.py | 60++++++------------------------------------------------------
MGlobalModel.py | 30+++++++++++-------------------
MServer.py | 80++++++++++++++++++++++++-------------------------------------------------------
MSoftMaxModel.py | 35+++++++++++++++++++++--------------
Aoptions.json | 6++++++
Mutils.py | 19++++++++++---------
6 files changed, 78 insertions(+), 152 deletions(-)

diff --git a/Client.py b/Client.py @@ -4,62 +4,14 @@ Classes and functions for the client networking aspect of federated learning Author: Cody Lewis """ -import socket -import pickle - -import torch -import torch.nn as nn - import SoftMaxModel -import utils class Client: - def __init__(self, x, y): - dims = utils.get_dims(x.shape, y.shape) - self.net = SoftMaxModel.SoftMaxModel(dims['x'], dims['y']) - self.x = x - self.y = y - self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - - def connect(self, host, port): - """Connect to the host:port federated learning server""" - self.socket.connect((host, port)) - - def fit(self): - criterion = nn.BCELoss() - e = 0 - while (msg := self.socket.recv(4096)) != b'DONE': - self.net.copy_params(pickle.loads(msg)) - e += 1 - history, grads = self.net.fit(self.x, self.y, 1, verbose=False) - self.socket.sendall(pickle.dumps(grads)) - print( - f"Epoch: {e}, Loss: {criterion(self.net(X), Y)}", - end="\r" - ) - # An improvement would be to save grads as a backlog and concurrent - # send them when the server is ready - print() - + def __init__(self, data): + self.net = SoftMaxModel.SoftMaxModel(data['x_dim'], data['y_dim']) + self.x = data['x'] + self.y = data['y'] -if __name__ == '__main__': - # X = torch.tensor([ - # [0, 0], - # [0, 1], - # [1, 0], - # [1, 1] - # ], dtype=torch.float32) - # Y = torch.tensor([ - # [1, 0], - # [1, 0], - # [1, 0], - # [0, 1] - # ], dtype=torch.float32) - X, Y = utils.load_data("mnist") - client = Client(X, Y) - HOST, PORT = '127.0.0.1', 5000 - print(f"Connecting to {HOST}:{PORT}") - client.connect(HOST, PORT) - client.fit() + def fit(self, verbose=True): + return self.net.fit(self.x, self.y, 8, 1, verbose=verbose) diff --git a/GlobalModel.py b/GlobalModel.py @@ -8,8 +8,10 @@ import torch import torch.nn as nn import SoftMaxModel +import utils +# TODO: Maybe remove num_clients? class GlobalModel: """The central global model for use within federated learning""" def __init__(self, num_in, num_out): @@ -80,34 +82,24 @@ def foolsgold(histories, num_clients, kappa, grads): if __name__ == '__main__': - x = torch.tensor([ - [0, 0], - [0, 1], - [1, 0], - [1, 1] - ], dtype=torch.float32) - y = torch.tensor([ - [1, 0], - [1, 0], - [1, 0], - [0, 1] - ], dtype=torch.float32) - server = GlobalModel(len(x[0]), len(y[0])) - client = SoftMaxModel.SoftMaxModel(len(x[0]), len(y[0])) + data = utils.load_data('mnist') + server = GlobalModel(data['x_dim'], data['y_dim']) + client = SoftMaxModel.SoftMaxModel(data['x_dim'], data['y_dim']) client.copy_params(server.get_params()) server.add_client() epochs = 5000 - criterion = nn.BCELoss() + criterion = nn.CrossEntropyLoss() for i in range(epochs): - history, grads = client.fit(x, y, 1, verbose=False) + client.copy_params(server.net.get_params()) + loss, grads = client.fit(data['x'], data['y'], 8, 1, verbose=False) grads = {0: grads} server.fit(1, grads) print( - f"Epoch {i + 1}/{epochs}, client loss: {history['loss'][-1]}, server loss {criterion(server.predict(x), y)}", + f"Epoch {i + 1}/{epochs}, client loss: {loss}, server loss: {criterion(server.predict(data['x']), data['y'][0])}", end="\r" ) print() print() - print(f"Client's final predictions:\n{client(x)}") + print(f"Client's final predictions:\n{client(data['x'])}") print() - print(f"Server's final predictions:\n{server.predict(x)}") + print(f"Server's final predictions:\n{server.predict(data['x'])}") diff --git a/Server.py b/Server.py @@ -4,84 +4,52 @@ Classes and functions for the server networking aspect of federated learning Author: Cody Lewis """ -import socket -import pickle import torch.nn as nn import GlobalModel +import Client import utils class Server: """Federated learning server class""" - def __init__(self, num_in, num_out, port=5000): + def __init__(self, num_in, num_out): self.net = GlobalModel.GlobalModel(num_in, num_out) self.num_clients = 0 - self.address = ('', port) - self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) - self.socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) - self.socket.bind(self.address) self.clients = [] - def accept_client(self, s): - """Accept a client and update the model accordingly""" - res = s.accept() - self.num_clients += 1 - self.net.add_client() - return res - - def accept_clients(self, num_clients): - """Accept some clients to the system""" - self.socket.listen(num_clients) - self.clients.extend([ - (c, addr) for c, addr in [ - self.accept_client(self.socket) for _ in range(num_clients) - ] - ]) - for c, _ in self.clients: - c.send(pickle.dumps(self.net.get_params())) - def fit(self, X, Y, epochs): - criterion = nn.BCELoss() + criterion = nn.CrossEntropyLoss() for e in range(epochs): grads = dict() - for i, (c, _) in enumerate(self.clients): - grads[i] = pickle.loads(c.recv(4096)) + losses = dict() + for i, c in enumerate(self.clients): + c.net.copy_params(self.net.get_params()) + losses[i], grads[i] = c.fit(verbose=False) self.net.fit(1, grads) - for c, _ in self.clients: - c.send(pickle.dumps(self.net.get_params())) print( - f"Epoch: {e + 1}/{epochs}, Loss: {criterion(server.net.predict(X), Y)}", + f"Epoch: {e + 1}/{epochs}, Loss: {criterion(self.net.predict(X), Y[0])}", end="\r" ) print() - def close(self): - for c, _ in self.clients: - c.send(b'DONE') - c.close() - self.clients = [] + def add_clients(self, clients): + """Add clients to the server""" + self.num_clients += len(clients) + self.net.add_client() + self.clients.extend(clients) + +def fit_client(client): + return client.fit(verbose=True) if __name__ == '__main__': - PORT = 5000 - X, Y = utils.load_data("mnist", train=False) - dims = utils.get_dims(X.shape, Y.shape) - server = Server(dims['x'], dims['y'], PORT) - print(f"Starting server on port {PORT}") - # X = torch.tensor([ - # [0, 0], - # [0, 1], - # [1, 0], - # [1, 1] - # ], dtype=torch.float32) - # Y = torch.tensor([ - # [1, 0], - # [1, 0], - # [1, 0], - # [0, 1] - # ], dtype=torch.float32) - server.accept_clients(1) - server.fit(X, Y, 5000) - server.close() + print("Loading Datasets and Creating Models...") + train_data = utils.load_data("mnist", train=True) + val_data = utils.load_data("mnist", train=False) + server = Server(val_data['x_dim'], val_data['y_dim']) + server.add_clients([Client.Client(train_data) for _ in range(5)]) + print("Done.") + print("Starting training...") + server.fit(val_data['x'], val_data['y'], 2) diff --git a/SoftMaxModel.py b/SoftMaxModel.py @@ -4,6 +4,7 @@ Pytorch implementation of a softmax perceptron Author: Cody Lewis """ +import torch import torch.nn as nn import torch.optim as optim @@ -16,7 +17,9 @@ class SoftMaxModel(nn.Module): super(SoftMaxModel, self).__init__() self.features = nn.ModuleList([ - nn.Linear(num_in, num_out), + nn.Linear(num_in, num_in * 10), + nn.Sigmoid(), + nn.Linear(num_in * 10, num_out), nn.Softmax(dim=1) ]).eval() self.lr = lr @@ -26,7 +29,7 @@ class SoftMaxModel(nn.Module): x = feature(x) return x - def fit(self, x, y, epochs, verbose=True): + def fit(self, x, y, batch_size=0, epochs=1, verbose=True): """ Fit the model for some epochs, return history of loss values and the gradients of the changed parameters @@ -38,24 +41,29 @@ class SoftMaxModel(nn.Module): verbose -- output training stats if True """ optimizer = optim.SGD(self.parameters(), lr=self.lr, momentum=0.9) - # criterion = nn.MSELoss() criterion = nn.CrossEntropyLoss() - history = {'loss': []} - # TODO: take a random batch of the data + n, _ = x.shape for i in range(epochs): optimizer.zero_grad() - output = self(x) - history['loss'].append(criterion(output, y)) + if 0 < batch_size < n: + ids = torch.randperm(n)[:batch_size] + sample_x = x[ids] + sample_y = y[0][ids] + else: + sample_x = x + sample_y = y + output = self(sample_x) + loss = criterion(output, sample_y) if verbose: print( - f"Epoch {i + 1}/{epochs} loss: {history['loss'][-1]}", + f"Epoch {i + 1}/{epochs} loss: {loss}", end="\r" ) - history['loss'][-1].backward() + loss.backward() optimizer.step() if verbose: print() - return history, { + return loss, { "params": [p.grad for p in self.parameters()], "data_count": len(x) } @@ -71,7 +79,6 @@ class SoftMaxModel(nn.Module): if __name__ == '__main__': - X, Y = utils.load_data("mnist") - dims = utils.get_dims(X.shape, Y.shape) - net = SoftMaxModel(dims['x'], dims['y']) - net.fit(X, Y, 500) + data = utils.load_data("mnist") + net = SoftMaxModel(data['x_dim'], data['y_dim']) + net.fit(data['x'], data['y'], 64, 50000) diff --git a/options.json b/options.json @@ -0,0 +1,6 @@ +{ + "epochs": 5000, + "users": 2, + "batch_size": 8, + "result_log_file": "./results.log" +} diff --git a/utils.py b/utils.py @@ -1,7 +1,7 @@ """ Utility functions for use on other classes in this project -Author: Cody +Author: Cody Lewis """ import torch @@ -24,16 +24,17 @@ def load_data(ds_name, train=True): return torch.tensor(), torch.tensor() data = chosen_set(f"./data/{ds_name}", train=train, download=True) X = data.data + Y = data.targets.long().unsqueeze(dim=0) if len(X.shape) == 3: X = X.reshape(X.shape[0], X.shape[1] * X.shape[2]) - # return X.float(), nn.functional.one_hot(data.targets) - return X.float(), data.targets.long().unsqueeze(dim=0) - - -def get_dims(x_shape, y_shape): - """Get the dimensions for a dataset based on its shapes""" return { - "x": x_shape[-1] if len(x_shape) > 1 else 1, - "y": y_shape[-1] if len(y_shape) > 1 else 1, + "x": X.float(), + "y": Y, + "x_dim": X.shape[-1] if len(X.shape) > 1 else 1, + "y_dim": int(torch.max(Y)) + 1, } +# def find_acc(model, X, Y): +# torch.argmax(server.net.predict(val_data['x']), dim=1) == val_data['y'][0] +# count instances of true in above divide by length +