def prepare_training(hook, alice): # pragma: no cover data, target = utils.create_gaussian_mixture_toy_data(nr_samples=100) dataset_key = "gaussian_mixture" dataset = sy.BaseDataset(data, target) alice.add_dataset(dataset, key=dataset_key) @hook.torch.jit.script def loss_fn(pred, target): return ((pred - target.unsqueeze(1))**2).mean() class Net(torch.nn.Module): def __init__(self): super(Net, self).__init__() self.fc1 = nn.Linear(2, 3) self.fc2 = nn.Linear(3, 2) self.fc3 = nn.Linear(2, 1) def forward(self, x): x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x model_untraced = Net() model = torch.jit.trace(model_untraced, data) pred = model(data) loss_before = loss_fn(target=target, pred=pred) return model, loss_fn, data, target, loss_before, dataset_key
def start_websocket_server_worker(id, host, port, hook, verbose, keep_labels=None, training=True): # pragma: no cover """Helper function for spinning up a websocket server and setting up the local datasets.""" server = WebsocketServerWorker(id=id, host=host, port=port, hook=hook, verbose=verbose) # Setup toy data (mnist example) mnist_dataset = datasets.MNIST( root="./data", train=training, download=True, transform=transforms.Compose([ transforms.ToTensor(), transforms.Normalize((0.1307, ), (0.3081, )) ]), ) if training: indices = np.isin(mnist_dataset.targets, keep_labels).astype("uint8") logger.info("number of true indices: %s", indices.sum()) selected_data = (torch.native_masked_select( mnist_dataset.data.transpose(0, 2), torch.tensor(indices)).view(28, 28, -1).transpose(2, 0)) logger.info("after selection: %s", selected_data.shape) selected_targets = torch.native_masked_select(mnist_dataset.targets, torch.tensor(indices)) dataset = sy.BaseDataset(data=selected_data, targets=selected_targets, transform=mnist_dataset.transform) key = "mnist" else: dataset = sy.BaseDataset( data=mnist_dataset.data, targets=mnist_dataset.targets, transform=mnist_dataset.transform, ) key = "mnist_testing" server.add_dataset(dataset, key=key) # Setup toy data (vectors example) data_vectors = torch.tensor([[-1, 2.0], [0, 1.1], [-1, 2.1], [0, 1.2]], requires_grad=True) target_vectors = torch.tensor([[1], [0], [1], [0]]) server.add_dataset(sy.BaseDataset(data_vectors, target_vectors), key="vectors") # Setup toy data (xor example) data_xor = torch.tensor([[0.0, 1.0], [1.0, 0.0], [1.0, 1.0], [0.0, 0.0]], requires_grad=True) target_xor = torch.tensor([1.0, 1.0, 0.0, 0.0], requires_grad=False) server.add_dataset(sy.BaseDataset(data_xor, target_xor), key="xor") # Setup gaussian mixture dataset data, target = utils.create_gaussian_mixture_toy_data(nr_samples=100) server.add_dataset(sy.BaseDataset(data, target), key="gaussian_mixture") # Setup partial iris dataset data, target = utils.iris_data_partial() dataset = sy.BaseDataset(data, target) dataset_key = "iris" server.add_dataset(dataset, key=dataset_key) logger.info("datasets: %s", server.datasets) if training: logger.info("len(datasets[mnist]): %s", len(server.datasets["mnist"])) server.start() return server
def test_fit(fit_dataset_key, epochs, device): if device == "cuda" and not torch.cuda.is_available(): return data, target = utils.create_gaussian_mixture_toy_data(nr_samples=100) fed_client = FederatedClient() dataset = sy.BaseDataset(data, target) dataset_key = "gaussian_mixture" fed_client.add_dataset(dataset, key=dataset_key) def loss_fn(pred, target): return torch.nn.functional.cross_entropy(input=pred, target=target) class Net(torch.nn.Module): def __init__(self): super(Net, self).__init__() self.fc1 = torch.nn.Linear(2, 3) self.fc2 = torch.nn.Linear(3, 2) torch.nn.init.xavier_normal_(self.fc1.weight) torch.nn.init.xavier_normal_(self.fc2.weight) def forward(self, x): x = torch.nn.functional.relu(self.fc1(x)) x = torch.nn.functional.relu(self.fc2(x)) return x data_device = data.to(torch.device(device)) target_device = target.to(torch.device(device)) model_untraced = Net().to(torch.device(device)) model = torch.jit.trace(model_untraced, data_device) model_id = 0 model_ow = ObjectWrapper(obj=model, id=model_id) loss_id = 1 loss_ow = ObjectWrapper(obj=loss_fn, id=loss_id) pred = model(data_device) loss_before = loss_fn(target=target_device, pred=pred) if PRINT_IN_UNITTESTS: # pragma: no cover print("Loss before training: {}".format(loss_before)) # Create and send train config train_config = sy.TrainConfig( batch_size=8, model=None, loss_fn=None, model_id=model_id, loss_fn_id=loss_id, optimizer_args={"lr": 0.05, "weight_decay": 0.01}, epochs=epochs, ) fed_client.set_obj(model_ow) fed_client.set_obj(loss_ow) fed_client.set_obj(train_config) fed_client.optimizer = None train_model( fed_client, fit_dataset_key, available_dataset_key=dataset_key, nr_rounds=3, device=device ) if dataset_key == fit_dataset_key: loss_after = evaluate_model(fed_client, model_id, loss_fn, data_device, target_device) if PRINT_IN_UNITTESTS: # pragma: no cover print("Loss after training: {}".format(loss_after)) if loss_after >= loss_before: # pragma: no cover if PRINT_IN_UNITTESTS: print("Loss not reduced, train more: {}".format(loss_after)) train_model( fed_client, fit_dataset_key, available_dataset_key=dataset_key, nr_rounds=10 ) loss_after = evaluate_model(fed_client, model_id, loss_fn, data, target) assert loss_after < loss_before
def test_train_config_with_jit_trace_sync( hook, start_remote_worker): # pragma: no cover data, target = utils.create_gaussian_mixture_toy_data(100) dataset = sy.BaseDataset(data, target) dataset_key = "gaussian_mixture" server, remote_proxy = start_remote_worker(id="sync_fit", hook=hook, port=9000, dataset=(dataset, dataset_key)) @hook.torch.jit.script def loss_fn(pred, target): return ((target.view(pred.shape).float() - pred.float())**2).mean() class Net(torch.nn.Module): def __init__(self): super(Net, self).__init__() self.fc1 = nn.Linear(2, 3) self.fc2 = nn.Linear(3, 2) self.fc3 = nn.Linear(2, 1) def forward(self, x): x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x model_untraced = Net() model = torch.jit.trace(model_untraced, data) pred = model(data) loss_before = loss_fn(pred=pred, target=target) # Create and send train config train_config = sy.TrainConfig(model=model, loss_fn=loss_fn, batch_size=2, epochs=1) train_config.send(remote_proxy) for epoch in range(5): loss = remote_proxy.fit(dataset_key=dataset_key) if PRINT_IN_UNITTESTS: # pragma: no cover print("-" * 50) print("Iteration %s: alice's loss: %s" % (epoch, loss)) new_model = train_config.model_ptr.get() # assert that the new model has updated (modified) parameters assert not ((model.fc1._parameters["weight"] - new_model.obj.fc1._parameters["weight"]).abs() < 10e-3).all() assert not ((model.fc2._parameters["weight"] - new_model.obj.fc2._parameters["weight"]).abs() < 10e-3).all() assert not ((model.fc3._parameters["weight"] - new_model.obj.fc3._parameters["weight"]).abs() < 10e-3).all() assert not ((model.fc1._parameters["bias"] - new_model.obj.fc1._parameters["bias"]).abs() < 10e-3).all() assert not ((model.fc2._parameters["bias"] - new_model.obj.fc2._parameters["bias"]).abs() < 10e-3).all() assert not ((model.fc3._parameters["bias"] - new_model.obj.fc3._parameters["bias"]).abs() < 10e-3).all() new_model.obj.eval() pred = new_model.obj(data) loss_after = loss_fn(pred=pred, target=target) if PRINT_IN_UNITTESTS: # pragma: no cover print("Loss before training: {}".format(loss_before)) print("Loss after training: {}".format(loss_after)) remote_proxy.close() server.terminate() assert loss_after < loss_before
async def test_train_config_with_jit_trace_async( hook, start_proc): # pragma: no cover kwargs = { "id": "async_fit", "host": "localhost", "port": 8777, "hook": hook } # data = torch.tensor([[0.0, 1.0], [1.0, 0.0], [1.0, 1.0], [0.0, 0.0]], requires_grad=True) # target = torch.tensor([[1.0], [1.0], [0.0], [0.0]], requires_grad=False) # dataset_key = "xor" data, target = utils.create_gaussian_mixture_toy_data(100) dataset_key = "gaussian_mixture" mock_data = torch.zeros(1, 2) # TODO check reason for error (RuntimeError: This event loop is already running) when starting websocket server from pytest-asyncio environment # dataset = sy.BaseDataset(data, target) # server, remote_proxy = start_remote_worker(id="async_fit", port=8777, hook=hook, dataset=(dataset, dataset_key)) # time.sleep(0.1) remote_proxy = WebsocketClientWorker(**kwargs) @hook.torch.jit.script def loss_fn(pred, target): return ((target.view(pred.shape).float() - pred.float())**2).mean() class Net(torch.nn.Module): def __init__(self): super(Net, self).__init__() self.fc1 = nn.Linear(2, 3) self.fc2 = nn.Linear(3, 2) self.fc3 = nn.Linear(2, 1) def forward(self, x): x = F.relu(self.fc1(x)) x = F.relu(self.fc2(x)) x = self.fc3(x) return x model_untraced = Net() model = torch.jit.trace(model_untraced, mock_data) pred = model(data) loss_before = loss_fn(target=target, pred=pred) # Create and send train config train_config = sy.TrainConfig(model=model, loss_fn=loss_fn, batch_size=2, optimizer="SGD", optimizer_args={"lr": 0.1}) train_config.send(remote_proxy) for epoch in range(5): loss = await remote_proxy.async_fit(dataset_key=dataset_key) if PRINT_IN_UNITTESTS: # pragma: no cover print("-" * 50) print("Iteration %s: alice's loss: %s" % (epoch, loss)) new_model = train_config.model_ptr.get() assert not (model.fc1._parameters["weight"] == new_model.obj.fc1._parameters["weight"]).all() assert not (model.fc2._parameters["weight"] == new_model.obj.fc2._parameters["weight"]).all() assert not (model.fc3._parameters["weight"] == new_model.obj.fc3._parameters["weight"]).all() assert not (model.fc1._parameters["bias"] == new_model.obj.fc1._parameters["bias"]).all() assert not (model.fc2._parameters["bias"] == new_model.obj.fc2._parameters["bias"]).all() assert not (model.fc3._parameters["bias"] == new_model.obj.fc3._parameters["bias"]).all() new_model.obj.eval() pred = new_model.obj(data) loss_after = loss_fn(target=target, pred=pred) if PRINT_IN_UNITTESTS: # pragma: no cover print("Loss before training: {}".format(loss_before)) print("Loss after training: {}".format(loss_after)) remote_proxy.close() # server.terminate() assert loss_after < loss_before