async def _fit_model_on_worker(self, worker, curr_round):
"""Send the model to the worker and fit the model on the worker's training data.
Args:
worker (class:`syft.workers.WebsocketClientWorker`): Remote worker, where
the model shall be trained.
Returns:
A tuple containing:
* worker_id: Union[int, str], id of the worker.
* improved model: torch.jit.ScriptModule, model after remote training.
* loss: Loss on last training batch, torch.tensor.
"""
train_config = sy.TrainConfig(
model=self.model,
loss_fn=self.loss_fn,
batch_size=self.config.train_batch_size,
shuffle=True,
max_nr_batches=self.config.fed_after_n_batches,
epochs=1,
optimizer=self.config.optimizer,
optimizer_args=self.config.optimizer_params,
)
train_config.send(worker)
loss = await worker.async_fit(
dataset_key=self.config.train_dataset_name, return_ids=[0])
logger.info("Training round: %s, worker: %s, avg_loss: %.4f" %
(curr_round, worker.id, loss.mean().item()))
model = train_config.model_ptr.get().obj
return worker.id, model, loss
def _evaluate_model_on_worker(self, model_identifier, worker, model):
"""
Method to evaluate models on remote workers. It logs model performance.
Args:
model_identifier (str): The model which
worker (class:`syft.workers.WebsocketClientWorker`): Remote worker, where
the model shall be evaluated. Note that, all workers have the same test
set so we always use the first.
model (:class:`torch.nn.Module`): Neural network defined in Torch.
"""
model.eval()
# Create and send train config
train_config = sy.TrainConfig(model=model,
loss_fn=self.loss_fn,
batch_size=self.config.test_batch_size,
optimizer_args=None,
epochs=1)
train_config.send(worker)
test_loss, eval_metrics, ys = worker.evaluate(
dataset_key=self.config.test_dataset_name,
metrics=self.config.metrics,
regression=self.config.regression)
logger.info("%s: Test set: Average loss: %.4f" %
(model_identifier, test_loss))
print_metrics(eval_metrics, ys)
def test_send_model_and_loss_fn(workers):
train_config = sy.TrainConfig(batch_size=2,
id="send_model_and_loss_fn_tc",
model=None,
loss_fn=None)
alice = workers["alice"]
orig_func = sy.ID_PROVIDER.pop
model_id = 44
model_id_at_location = 44000
loss_fn_id = 55
loss_fn_id_at_location = 55000
sy.ID_PROVIDER.pop = mock.Mock(side_effect=[
model_id, model_id_at_location, loss_fn_id, loss_fn_id_at_location
])
train_config.send(alice)
assert alice.train_config.id == train_config.id
assert alice.train_config._model_id == train_config._model_id
assert alice.train_config._loss_fn_id == train_config._loss_fn_id
assert alice.train_config.batch_size == train_config.batch_size
assert alice.train_config.epochs == train_config.epochs
assert alice.train_config.optimizer == train_config.optimizer
assert alice.train_config.optimizer_args == train_config.optimizer_args
assert alice.train_config.location == train_config.location
assert alice.train_config._model_id == model_id
assert alice.train_config._loss_fn_id == loss_fn_id
sy.ID_PROVIDER.pop = orig_func
def _detail_train_config(worker: AbstractWorker,
train_config_tuple: tuple) -> sy.TrainConfig:
"""This function reconstructs a TrainConfig object given it's attributes in the form of a tuple.
Args:
worker: the worker doing the deserialization
train_config_tuple: a tuple holding the attributes of the TrainConfig
Returns:
train_config: A TrainConfig object
"""
model_id, loss_fn_id, batch_size, epochs, optimizer, lr, id, max_nr_batches, shuffle = (
train_config_tuple)
id = _detail(worker, id)
detailed_optimizer = _detail(worker, optimizer)
train_config = syft.TrainConfig(
model=None,
loss_fn=None,
owner=worker,
id=id,
model_id=model_id,
loss_fn_id=loss_fn_id,
batch_size=batch_size,
epochs=epochs,
optimizer=detailed_optimizer,
lr=lr,
max_nr_batches=max_nr_batches,
shuffle=shuffle,
)
return train_config
def test___str__():
train_config = sy.TrainConfig(batch_size=2, id=99887766, model=None, loss_fn=None)
train_config_str = str(train_config)
str_expected = "<TrainConfig id:99887766 owner:me epochs: 1 batch_size: 2 lr: 0.1>"
assert str_expected == train_config_str
async def fit_model_on_worker(worker: WebsocketClientWorker,
traced_model: torch.jit.ScriptModule,
batch_size: int, curr_round: int, lr: float,
no_federated_epochs: int):
train_config = syft.TrainConfig(
model=traced_model,
loss_fn=loss_fn,
batch_size=batch_size,
shuffle=True,
epochs=no_federated_epochs,
optimizer="SGD",
optimizer_args={"lr": lr},
)
#send the training config
train_config.send(worker)
#Call async fit on worker - async fit calls the method calls self fit method
print("Training round {}, calling fit on worker: {}".format(
curr_round, worker.id))
loss = await worker.async_fit(dataset_key="targeted", return_ids=[0])
print("Training round: {}, worker: {}, avg_loss: {}".format(
curr_round, worker.id, loss.item()))
#Call back to the model
model = train_config.get_model().obj
return worker.id, model, loss
async def fit_model_on_worker(
args,
worker: websocket_client.WebsocketClientWorker,
traced_model: torch.jit.ScriptModule,
curr_round: int,
lr: float,
):
train_config = sy.TrainConfig(model=traced_model,
loss_fn=get_serializable_loss(args.loss),
optimizer=args.optimizer,
batch_size=args.batch_size,
optimizer_args={"lr": lr},
epochs=args.federate_after_n_batches,
shuffle=True,
max_nr_batches=args.max_nr_batches
)
train_config.send(worker)
loss = await worker.async_fit(dataset_key=args.dataset, return_ids=[0])
model = train_config.model_ptr.get().obj
return worker.id, model,loss
def test_train_config_with_jit_trace(hook, workers): # pragma: no cover
alice = workers["alice"]
me = workers["me"]
data = torch.tensor([[-1, 2.0], [0, 1.1], [-1, 2.1], [0, 1.2]],
requires_grad=True)
target = torch.tensor([[1], [0], [1], [0]])
dataset = sy.BaseDataset(data, target)
alice.add_dataset(dataset, key="vectors")
@hook.torch.jit.script
def loss_fn(real, pred):
return ((real.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)
model_with_id = pointers.ObjectWrapper(model, sy.ID_PROVIDER.pop())
loss_fn_with_id = pointers.ObjectWrapper(loss_fn, sy.ID_PROVIDER.pop())
model_ptr = me.send(model_with_id, alice)
loss_fn_ptr = me.send(loss_fn_with_id, alice)
print("Evaluation before training")
pred = model(data)
loss_before = loss_fn(real=target, pred=pred)
print("Loss: {}".format(loss_before))
# Create and send train config
train_config = sy.TrainConfig(model=model, loss_fn=loss_fn, batch_size=2)
train_config.send(alice)
for epoch in range(5):
loss = alice.fit(dataset="vectors")
print("-" * 50)
print("Iteration %s: alice's loss: %s" % (epoch, loss))
print("Evaluation after training:")
new_model = model_ptr.get()
pred = new_model.obj(data)
loss_after = loss_fn(real=target, pred=pred)
print("Loss: {}".format(loss_after))
assert loss_after < loss_before
def test_train_config_with_jit_trace(hook, workers): # pragma: no cover
alice = workers["alice"]
data = torch.tensor([[-1, 2.0], [0, 1.1], [-1, 2.1], [0, 1.2]],
requires_grad=True)
target = torch.tensor([[1], [0], [1], [0]])
dataset = sy.BaseDataset(data, target)
alice.add_dataset(dataset, key="gaussian_mixture")
@hook.torch.jit.script
def loss_fn(pred, target):
return ((target.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)
if PRINT_IN_UNITTESTS:
print("Evaluation before training")
pred = model(data)
loss_before = loss_fn(target=target, pred=pred)
if PRINT_IN_UNITTESTS:
print("Loss: {}".format(loss_before))
# Create and send train config
train_config = sy.TrainConfig(model=model, loss_fn=loss_fn, batch_size=2)
train_config.send(alice)
for epoch in range(5):
loss = alice.fit(dataset_key="gaussian_mixture")
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()
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))
assert loss_after < loss_before
def test_train_config_with_jit_trace_send_twice_with_fit(
hook, workers): # pragma: no cover
torch.manual_seed(0)
alice = workers["alice"]
model, loss_fn, data, target, loss_before, dataset_key = prepare_training(
hook, alice)
# Create and send train config
train_config_0 = sy.TrainConfig(model=model, loss_fn=loss_fn, batch_size=2)
train_config_0.send(alice)
for epoch in range(5):
loss = alice.fit(dataset_key=dataset_key)
if PRINT_IN_UNITTESTS: # pragma: no cover:
print("-" * 50)
print("TrainConfig 0, iteration %s: alice's loss: %s" %
(epoch, loss))
new_model = train_config_0.model_ptr.get()
pred = new_model.obj(data)
loss_after_0 = loss_fn(pred=pred, target=target)
assert loss_after_0 < loss_before
train_config = sy.TrainConfig(model=model, loss_fn=loss_fn, batch_size=2)
train_config.send(alice)
for epoch in range(5):
loss = alice.fit(dataset_key=dataset_key)
if PRINT_IN_UNITTESTS: # pragma: no cover:
print("-" * 50)
print("TrainConfig 1, iteration %s: alice's loss: %s" %
(epoch, loss))
new_model = train_config.model_ptr.get()
pred = new_model.obj(data)
loss_after = loss_fn(pred=pred, target=target)
if PRINT_IN_UNITTESTS: # pragma: no cover:
print(
"Loss after training with TrainConfig 0: {}".format(loss_after_0))
print(
"Loss after training with TrainConfig 1: {}".format(loss_after))
assert loss_after < loss_before
def test___str__():
train_config = sy.TrainConfig(batch_size=2, id=99887766, model=None, loss_fn=None)
train_config_str = str(train_config)
str_expected = (
"<TrainConfig id:99887766 owner:me epochs: 1 batch_size: 2 optimizer_args: {'lr': 0.1}>"
)
assert str_expected == train_config_str
def test_fit():
data = torch.tensor([[-1, 2.0], [0, 1.1], [-1, 2.1], [0, 1.2]], requires_grad=True)
target = torch.tensor([[1], [0], [1], [0]])
fed_client = federated.FederatedClient()
dataset = sy.BaseDataset(data, target)
fed_client.add_dataset(dataset, key="vectors")
def loss_fn(real, pred):
return ((real.float() - pred.float()) ** 2).mean()
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)
self.fc3 = torch.nn.Linear(2, 1)
def forward(self, x):
x = torch.nn.functional.relu(self.fc1(x))
x = torch.nn.functional.relu(self.fc2(x))
x = self.fc3(x)
return x
model_untraced = Net()
model = torch.jit.trace(model_untraced, data)
model_id = 0
model_ow = pointers.ObjectWrapper(obj=model, id=model_id)
loss_id = 1
loss_ow = pointers.ObjectWrapper(obj=loss_fn, id=loss_id)
print("Evaluation before training")
pred = model(data)
loss_before = loss_fn(real=target, pred=pred)
print("Loss: {}".format(loss_before))
# Create and send train config
train_config = sy.TrainConfig(
batch_size=1, model=None, loss_fn=None, model_id=model_id, loss_fn_id=loss_id
)
fed_client.set_obj(model_ow)
fed_client.set_obj(loss_ow)
fed_client.set_obj(train_config)
for epoch in range(5):
loss = fed_client.fit(dataset_key="vectors")
print("-" * 50)
print("Iteration %s: alice's loss: %s" % (epoch, loss))
print("Evaluation after training:")
new_model = fed_client.get_obj(model_id)
pred = new_model.obj(data)
loss_after = loss_fn(real=target, pred=pred)
print("Loss: {}".format(loss_after))
assert loss_after < loss_before
def evaluate_model_on_worker(
model_identifier,
worker,
dataset_key,
model,
nr_bins,
batch_size,
device,
print_target_hist=False,
):
model.eval()
# Create and send train config
train_config = sy.TrainConfig(batch_size=batch_size,
model=model,
loss_fn=loss_fn,
optimizer_args=None,
epochs=1)
train_config.send(worker)
result = worker.evaluate(
dataset_key=dataset_key,
return_histograms=True,
nr_bins=nr_bins,
return_loss=True,
return_raw_accuracy=True,
device=device,
)
test_loss = result["loss"]
correct = result["nr_correct_predictions"]
len_dataset = result["nr_predictions"]
hist_pred = result["histogram_predictions"]
hist_target = result["histogram_target"]
if print_target_hist:
logger.info("Target histogram: %s", hist_target)
percentage_0_3 = int(100 * sum(hist_pred[0:4]) / len_dataset)
percentage_4_6 = int(100 * sum(hist_pred[4:7]) / len_dataset)
percentage_7_9 = int(100 * sum(hist_pred[7:10]) / len_dataset)
logger.info(
"%s: Percentage numbers 0-3: %s%%, 4-6: %s%%, 7-9: %s%%",
model_identifier,
percentage_0_3,
percentage_4_6,
percentage_7_9,
)
logger.info(
"%s: Average loss: %s, Accuracy: %s/%s (%s%%)",
model_identifier,
f"{test_loss:.4f}",
correct,
len_dataset,
f"{100.0 * correct / len_dataset:.2f}",
)
def test_send(workers):
alice = workers["alice"]
train_config = sy.TrainConfig(batch_size=2, id="id", model=None, loss_fn=None)
train_config.send(alice)
assert alice.train_config.id == train_config.id
assert alice.train_config._model_id == train_config._model_id
assert alice.train_config._loss_fn_id == train_config._loss_fn_id
assert alice.train_config.batch_size == train_config.batch_size
assert alice.train_config.epochs == train_config.epochs
assert alice.train_config.optimizer == train_config.optimizer
assert alice.train_config.optimizer_args == train_config.optimizer_args
assert alice.train_config.location == train_config.location
def test_train_config_with_jit_trace_send_twice_with_fit(
hook, workers): # pragma: no cover
alice = workers["alice"]
model, loss_fn, data, target, loss_before = prepare_training(hook, alice)
# Create and send train config
train_config_0 = sy.TrainConfig(model=model, loss_fn=loss_fn, batch_size=2)
train_config_0.send(alice)
for epoch in range(5):
loss = alice.fit(dataset_key="vectors")
print("-" * 50)
print("Iteration %s: alice's loss: %s" % (epoch, loss))
print("Evaluation after training train_config_0:")
new_model = train_config_0.model_ptr.get()
pred = new_model.obj(data)
loss_after = loss_fn(real=target, pred=pred)
print("Loss: {}".format(loss_after))
assert loss_after < loss_before
train_config = sy.TrainConfig(model=model, loss_fn=loss_fn, batch_size=2)
train_config.send(alice)
for epoch in range(5):
loss = alice.fit(dataset_key="vectors")
print("-" * 50)
print("Iteration %s: alice's loss: %s" % (epoch, loss))
print("Evaluation after training:")
new_model = train_config.model_ptr.get()
pred = new_model.obj(data)
loss_after = loss_fn(real=target, pred=pred)
print("Loss: {}".format(loss_after))
assert loss_after < loss_before
def evaluate_model_on_worker(model_identifier,
worker,
dataset_key: str,
model: Model,
nr_bins,
batch_size,
print_target_hist=False):
"""
Evaluate model on testing set. The testing set is
located on the remote server worker called `testing` -- the trusted
third party aggregartor.
Loss and accuracy over the test set are returned.
"""
model.eval()
# Create and send train config
train_config = sy.TrainConfig(
batch_size=batch_size,
model=model,
loss_fn=loss_fn,
optimizer_args=None,
epochs=1
) #evaluate full model on 'testing' remote server worker (trusted third party)
train_config.send(worker)
#evaluate aggregate model on validation set
result = worker.evaluate(
dataset_key=dataset_key,
nr_bins=nr_bins,
return_loss=True,
return_raw_accuracy=True,
)
test_loss = result["loss"]
correct = result["nr_correct_predictions"]
len_dataset = result["nr_predictions"]
logger.info(
"%s: Test set: Average loss: %s, Accuracy: %s/%s (%s)",
model_identifier,
"{:.4f}".format(test_loss),
correct,
len_dataset,
"{:.2f}".format(100.0 * correct / len_dataset),
)
def evaluate_model_on_worker(
model_identifier, worker, dataset_key, model, nr_bins, batch_size, print_target_hist=False
):
model.eval()
# Create and send train config
train_config = sy.TrainConfig(
batch_size=batch_size, model=model, loss_fn=loss_fn, optimizer_args=None, epochs=1
)
train_config.send(worker)
result = worker.evaluate(
dataset_key=dataset_key,
return_histograms=True,
nr_bins=nr_bins,
return_loss=True,
return_raw_accuracy=True,
)
test_loss = result["loss"]
correct = result["nr_correct_predictions"]
len_dataset = result["nr_predictions"]
hist_pred = result["histogram_predictions"]
hist_target = result["histogram_target"]
if print_target_hist:
logger.info("Target histogram: %s", hist_target)
logger.info("%s: Prediction hist.: %s", model_identifier, hist_pred)
logger.info(
"%s: Percentage numbers 0-3: %s", model_identifier, sum(hist_pred[0:4]) / len_dataset
)
logger.info(
"%s: Percentage numbers 4-6: %s", model_identifier, sum(hist_pred[4:7]) / len_dataset
)
logger.info(
"%s: Percentage numbers 7-9: %s", model_identifier, sum(hist_pred[7:10]) / len_dataset
)
logger.info(
"%s: Test set: Average loss: %s, Accuracy: %s/%s (%s)",
model_identifier,
"{:.4f}".format(test_loss),
correct,
len_dataset,
"{:.2f}".format(100.0 * correct / len_dataset),
)
def evaluate_model_on_worker(worker, dataset_key, model, batch_size):
""" Evaluate a model on worker"""
# Create and send train config
train_config = sy.TrainConfig(
batch_size=batch_size, model=model, loss_fn=loss_fn, optimizer_args={}, epochs=1
)
train_config.send(worker)
result = worker.evaluate(
dataset_key=dataset_key,
return_histograms=False,
return_loss=True,
return_raw_accuracy=False,
return_confusion_matrix=True,
example_inputs=torch.rand(1, 10),
)
return result['loss'], result['confusion_matrix']
async def fit_model_on_worker(
worker: workers.WebsocketClientWorker,
traced_model: torch.jit.ScriptModule,
batch_size: int,
curr_round: int,
max_nr_batches: int,
lr: float,
):
"""Send the model to the worker and fit the model on the worker's training data.
Args:
worker: Remote location, where the model shall be trained.
traced_model: Model which shall be trained.
batch_size: Batch size of each training step.
curr_round: Index of the current training round (for logging purposes).
max_nr_batches: If > 0, training on worker will stop at min(max_nr_batches, nr_available_batches).
lr: Learning rate of each training step.
Returns:
A tuple containing:
* worker_id: Union[int, str], id of the worker.
* improved model: torch.jit.ScriptModule, model after training at the worker.
* loss: Loss on last training batch, torch.tensor.
"""
train_config = sy.TrainConfig(
model=traced_model,
loss_fn=loss_fn,
batch_size=batch_size,
shuffle=True,
max_nr_batches=max_nr_batches,
epochs=1,
lr=lr,
)
train_config.send(worker)
logger.info(
"Training round %s, calling fit on worker: %s, lr = %s",
curr_round,
worker.id,
"{:.3f}".format(train_config.lr),
)
loss = await worker.async_fit(dataset_key="mnist", return_ids=[0])
logger.info("Training round: %s, worker: %s, avg_loss: %s", curr_round,
worker.id, loss.mean())
model = train_config.model_ptr.get().obj
return worker.id, model, loss
async def fit_model_on_worker(worker, traced_model, optimizer, batch_size, epochs, lr, dataset_key, shuffle):
""" Fit model on a worker """
print('Training on "%s" ...' % (worker.id,))
# Create and send train config
train_config = sy.TrainConfig(
model=traced_model,
loss_fn=loss_fn,
batch_size=batch_size,
shuffle=shuffle,
epochs=epochs,
optimizer=optimizer,
optimizer_args={'lr': lr},
)
train_config.send(worker)
await worker.async_fit(dataset_key=dataset_key, return_ids=[0])
model = train_config.model_ptr.get().obj
return worker.id, model
async def fit_model_on_worker(worker: WebsocketClientWorker,
traced_model: torch.jit.ScriptModule,
batch_size: int, curr_round: int, lr: float,
no_federated_epochs: int):
train_config = syft.TrainConfig(
model=traced_model,
loss_fn=loss_fn,
batch_size=batch_size,
shuffle=True,
epochs=no_federated_epochs,
optimizer="SGD",
optimizer_args={"lr": lr},
)
#send monitoring command
message = worker.create_message_execute_command(
command_name="start_monitoring", command_owner="self")
serialized_message = syft.serde.serialize(message)
worker._recv_msg(serialized_message)
#send the training config
train_config.send(worker)
#Call async fit on worker - async fit calls the method calls self fit method
print("Training round {}, calling fit on worker: {}".format(
curr_round, worker.id))
loss = await worker.async_fit(dataset_key="targeted", return_ids=[0])
print("Training round: {}, worker: {}, avg_loss: {}".format(
curr_round, worker.id, loss.item()))
#Call back to the model
model = train_config.get_model().obj
#Stop monitoring command
message = worker.create_message_execute_command(
command_name="stop_monitoring", command_owner="self")
serialized_message = syft.serde.serialize(message)
network_info = worker._recv_msg(serialized_message)
#Deserialize the response recieved
network_info = syft.serde.deserialize(network_info)
return worker.id, model, loss, network_info
async def fit_model_on_worker(
worker: websocket_client.WebsocketClientWorker,
traced_model: torch.jit.ScriptModule,
batch_size: int,
#curr_round: int,
epoch: int,
max_nr_batches: int,
lr: float,
):
train_config = sy.TrainConfig(
model=traced_model,
loss_fn=loss_fn,
batch_size=batch_size,
shuffle=True,
#max_nr_batches=max_nr_batches,
max_nr_batches=-1,
epochs=1,
optimizer="Adam",
optimizer_args={"lr": lr},
)
train_config.send(worker)
loss = await worker.async_fit(dataset_key="dga", return_ids=[0])
model = train_config.model_ptr.get().obj
return worker.id, model, loss
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
mock_data = torch.zeros(1)
traced_model = torch.jit.trace(model,mock_data)
optimizer = "SGD"
batch_size = 4
optimizer_args = {"lr" : 0.1, "weight_decay" : 0.01}
max_nr_batches = -1 # not used in this example
shuffle = True
train_config = syft.TrainConfig(model=traced_model,
loss_fn=loss_fn,
optimizer=optimizer,
batch_size=batch_size,
optimizer_args=optimizer_args,
epochs=5,
shuffle=shuffle)
arw = {"host":"10.0.0.1","hook":hook}
h1 = WebsocketClientWorker(id="h1",port=8778,**arw)
train_config.send(h1)
message = h1.create_message_execute_command(command_name="start_monitoring",command_owner="self")
serialized_message = syft.serde.serialize(message)
h1._recv_msg(serialized_message)
time.sleep(3)
def test_evaluate(hook, start_proc): # pragma: no cover
sy.local_worker.clear_objects()
sy.frameworks.torch.hook.hook_args.hook_method_args_functions = {}
sy.frameworks.torch.hook.hook_args.hook_method_response_functions = {}
sy.frameworks.torch.hook.hook_args.get_tensor_type_functions = {}
sy.frameworks.torch.hook.hook_args.register_response_functions = {}
data, target = utils.iris_data_partial()
dataset = sy.BaseDataset(data=data, targets=target)
kwargs = {
"id": "evaluate_remote",
"host": "localhost",
"port": 8780,
"hook": hook
}
dataset_key = "iris"
# TODO: check why unit test sometimes fails when WebsocketServerWorker is started from the unit test. Fails when run after test_federated_client.py
# process_remote_worker = start_proc(WebsocketServerWorker, dataset=(dataset, dataset_key), verbose=True, **kwargs)
local_worker = instantiate_websocket_client_worker(**kwargs)
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(4, 3)
torch.nn.init.xavier_normal_(self.fc1.weight)
def forward(self, x):
x = torch.nn.functional.relu(self.fc1(x))
return x
model_untraced = Net()
model = torch.jit.trace(model_untraced, data)
loss_traced = torch.jit.trace(
loss_fn, (torch.tensor([[0.3, 0.5, 0.2]]), torch.tensor([1])))
pred = model(data)
loss_before = loss_fn(target=target, pred=pred)
if PRINT_IN_UNITTESTS: # pragma: no cover
print("Loss: {}".format(loss_before))
# Create and send train config
train_config = sy.TrainConfig(
batch_size=4,
model=model,
loss_fn=loss_traced,
model_id=None,
loss_fn_id=None,
optimizer_args=None,
epochs=1,
)
train_config.send(local_worker)
result = local_worker.evaluate(dataset_key=dataset_key,
calculate_histograms=True,
nr_bins=3,
calculate_loss=True)
test_loss_before, correct_before, len_dataset, hist_pred_before, hist_target = result
if PRINT_IN_UNITTESTS: # pragma: no cover
print("Evaluation result before training: {}".format(result))
assert len_dataset == 30
assert (hist_target == [10, 10, 10]).all()
local_worker.close()
local_worker.remove_worker_from_local_worker_registry()
def test_evaluate(): # pragma: no cover
data, target = utils.iris_data_partial()
fed_client = FederatedClient()
dataset = sy.BaseDataset(data, target)
dataset_key = "iris"
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(4, 3)
def forward(self, x):
x = torch.nn.functional.relu(self.fc1(x))
return x
model_untraced = Net()
with torch.no_grad():
model_untraced.fc1.weight.set_(
torch.tensor(
[
[0.0160, 1.3753, -0.1202, -0.9129],
[0.1539, 0.3092, 0.0749, 0.2142],
[0.0984, 0.6248, 0.0274, 0.1735],
]
)
)
model_untraced.fc1.bias.set_(torch.tensor([0.3477, 0.2970, -0.0799]))
model = torch.jit.trace(model_untraced, data)
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)
loss_before = loss_fn(target=target, 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=None,
epochs=1,
)
fed_client.set_obj(model_ow)
fed_client.set_obj(loss_ow)
fed_client.set_obj(train_config)
fed_client.optimizer = None
result = fed_client.evaluate(
dataset_key=dataset_key, return_histograms=True, nr_bins=3, return_loss=True
)
test_loss_before = result["loss"]
correct_before = result["nr_correct_predictions"]
len_dataset = result["nr_predictions"]
hist_pred_before = result["histogram_predictions"]
hist_target = result["histogram_target"]
if PRINT_IN_UNITTESTS: # pragma: no cover
print("Evaluation result before training: {}".format(result))
assert len_dataset == 30
assert (hist_target == [10, 10, 10]).all()
train_config = sy.TrainConfig(
batch_size=8,
model=None,
loss_fn=None,
model_id=model_id,
loss_fn_id=loss_id,
optimizer="SGD",
optimizer_args={"lr": 0.01},
shuffle=True,
epochs=2,
)
fed_client.set_obj(train_config)
train_model(
fed_client, dataset_key, available_dataset_key=dataset_key, nr_rounds=50, device="cpu"
)
result = fed_client.evaluate(
dataset_key=dataset_key, return_histograms=True, nr_bins=3, return_loss=True
)
test_loss_after = result["loss"]
correct_after = result["nr_correct_predictions"]
len_dataset = result["nr_predictions"]
hist_pred_after = result["histogram_predictions"]
hist_target = result["histogram_target"]
if PRINT_IN_UNITTESTS: # pragma: no cover
print("Evaluation result: {}".format(result))
assert len_dataset == 30
assert (hist_target == [10, 10, 10]).all()
assert correct_after > correct_before
assert torch.norm(torch.tensor(hist_target - hist_pred_after)) < torch.norm(
torch.tensor(hist_target - hist_pred_before)
)
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
worker.clear_objects_remote()
traced_model = trace(Classifier(), torch.rand(1, 10))
train_losses = []
test_losses = []
confusion_matrices = []
for epoch in range(config_epochs):
traced_model.train()
train_config = sy.TrainConfig(
model=traced_model,
loss_fn=loss_fn,
batch_size=config_batch,
shuffle=config_shuffle,
max_nr_batches=-1,
epochs=1,
optimizer=config_optimizer,
optimizer_args={'lr': config_lr},
)
worker.clear_objects_remote()
train_config.send(worker)
loss = worker.fit(dataset_key='train')
traced_model = train_config.model_ptr.get().obj
# Evaluation
traced_model.eval()
# Evaluate on train set
train_config = sy.TrainConfig(
def test_fit():
data, target = utils.create_gaussian_mixture_toy_data(nr_samples=100)
fed_client = federated.FederatedClient()
dataset = sy.BaseDataset(data, target)
dataset_key = "gaussian_mixture"
fed_client.add_dataset(dataset, key=dataset_key)
def loss_fn(target, pred):
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
model_untraced = Net()
model = torch.jit.trace(model_untraced, data)
model_id = 0
model_ow = pointers.ObjectWrapper(obj=model, id=model_id)
loss_id = 1
loss_ow = pointers.ObjectWrapper(obj=loss_fn, id=loss_id)
pred = model(data)
loss_before = loss_fn(target=target, 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,
lr=0.05,
weight_decay=0.01,
)
fed_client.set_obj(model_ow)
fed_client.set_obj(loss_ow)
fed_client.set_obj(train_config)
fed_client.optimizer = None
for curr_round in range(12):
loss = fed_client.fit(dataset_key=dataset_key)
if PRINT_IN_UNITTESTS and curr_round % 4 == 0: # pragma: no cover
print("-" * 50)
print("Iteration %s: alice's loss: %s" % (curr_round, loss))
new_model = fed_client.get_obj(model_id)
pred = new_model.obj(data)
loss_after = loss_fn(target=target, pred=pred)
if PRINT_IN_UNITTESTS: # pragma: no cover:
print("Loss after training: {}".format(loss_after))
assert loss_after < loss_before
