def send_model_to_worker(
worker,
built_model: sy.Plan,
):
"""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.
"""
built_model.id = "GlobalModel"
model_send_start = time.time()
# pdb.set_trace()
built_model.send(worker)
print("[trace] GlobalModelSend duration", worker.id,
time.time() - model_send_start)
return None
def test_plan_execution(client: sy.VirtualMachineClient) -> None:
tensor_pointer1 = th.tensor([1, 2, 3]).send(client)
tensor_pointer2 = th.tensor([4, 5, 6]).send(client)
tensor_pointer3 = th.tensor([7, 8, 9]).send(client)
result_tensor_pointer1 = th.tensor([0, 0, 0]).send(client)
result_tensor_pointer2 = th.tensor([0, 0, 0]).send(client)
result1_uid = result_tensor_pointer1.id_at_location
result2_uid = result_tensor_pointer2.id_at_location
a1 = RunClassMethodAction(
path="torch.Tensor.add",
_self=tensor_pointer1,
args=[tensor_pointer2],
kwargs={},
id_at_location=result1_uid,
address=Address(),
msg_id=UID(),
)
a2 = RunClassMethodAction(
path="torch.Tensor.add",
_self=result_tensor_pointer1,
args=[tensor_pointer3],
kwargs={},
id_at_location=result2_uid,
address=Address(),
msg_id=UID(),
)
plan = Plan([a1, a2])
plan_pointer = plan.send(client)
plan_pointer()
expected_tensor1 = th.tensor([5, 7, 9])
expected_tensor2 = th.tensor([12, 15, 18])
assert all(expected_tensor1 == result_tensor_pointer1.get())
assert all(expected_tensor2 == result_tensor_pointer2.get())
def test_plan_batched_execution(client: sy.VirtualMachineClient) -> None:
# placeholders for our input
input_tensor_pointer1 = th.tensor([0, 0]).send(client)
input_tensor_pointer2 = th.tensor([0, 0]).send(client)
# tensors in our model
model_tensor_pointer1 = th.tensor([1, 2]).send(client)
model_tensor_pointer2 = th.tensor([3, 4]).send(client)
# placeholders for intermediate results
result_tensor_pointer1 = th.tensor([0, 0]).send(client)
result_tensor_pointer2 = th.tensor([0, 0]).send(client)
result_tensor_pointer3 = th.tensor([0, 0]).send(client)
# define plan
a1 = RunClassMethodAction(
path="torch.Tensor.mul",
_self=input_tensor_pointer1,
args=[model_tensor_pointer1],
kwargs={},
id_at_location=result_tensor_pointer1.id_at_location,
address=Address(),
msg_id=UID(),
)
a2 = RunClassMethodAction(
path="torch.Tensor.add",
_self=result_tensor_pointer1,
args=[model_tensor_pointer2],
kwargs={},
id_at_location=result_tensor_pointer2.id_at_location,
address=Address(),
msg_id=UID(),
)
a3 = RunFunctionOrConstructorAction(
path="torch.eq",
args=[result_tensor_pointer2, input_tensor_pointer2],
kwargs={},
id_at_location=result_tensor_pointer3.id_at_location,
address=Address(),
msg_id=UID(),
)
plan = Plan([a1, a2, a3],
inputs={
"x": input_tensor_pointer1,
"y": input_tensor_pointer2
})
plan_pointer = plan.send(client)
# Test
# x is random input, y is the expected model(x)
x_batches = [(th.tensor([1, 1]) + i).send(client) for i in range(2)]
y_batches = [((th.tensor([1, 1]) + i) * th.tensor([1, 2]) +
th.tensor([3, 4])).send(client) for i in range(2)]
for x, y in zip(x_batches, y_batches):
plan_pointer(x=x, y=y)
# checks if (model(x) == y) == [True, True]
assert all(result_tensor_pointer3.get(delete_obj=False))
async def fit_model_on_worker(
worker,
built_model: sy.Plan,
built_loss_fn: sy.Plan,
encrypters,
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.
"""
num_of_parameters = len(built_model.parameters())
built_model.id = "GlobalModel"
# built_loss_fn.id = "LossFunc"
# model_config = sy.ModelConfig(model=built_model,
# loss_fn=built_loss_fn,
# optimizer="SGD",
# batch_size=batch_size,
# optimizer_args={"lr": lr},
# epochs=1,
# max_nr_batches=max_nr_batches)
# model_config_send_start = time.time()
built_model.send(worker)
# model_config_send_end = time.time()
print("[trace] GlobalInformationSend duration", worker.id,
model_config_send_end - model_config_send_start)
return_ids = [0, 1]
for i in range(num_of_parameters):
return_ids.append("p" + str(i))
fit_sagg_start = time.time()
result_list = await worker.async_fit_sagg_mc(dataset_key="mnist",
encrypters=encrypters,
return_ids=return_ids)
fit_sagg_end = time.time()
print("[trace] FitSagg", "duration", worker.id,
fit_sagg_end - fit_sagg_start)
loss = result_list[0]
num_of_training_data = result_list[1]
enc_params = result_list[2:]
print("Iteration %s: %s loss: %s" % (curr_round, worker.id, loss))
return worker.id, enc_params, loss, num_of_training_data