def test_unpack_typerror_crypten_model():
dummy_input = th.rand(1, 28 * 28)
expected_crypten_model = crypten.nn.from_pytorch(ExampleNet(), dummy_input)
packed = utils.pack_values(expected_crypten_model)
with pytest.raises(TypeError):
utils.unpack_values(packed)
def test_pack_crypten_model():
class ExampleNet(th.nn.Module):
def __init__(self):
super(ExampleNet, self).__init__()
self.fc = th.nn.Linear(28 * 28, 2)
def forward(self, x):
out = self.fc(x)
return out
dummy_input = th.rand(1, 28 * 28)
expected_crypten_model = crypten.nn.from_pytorch(ExampleNet(), dummy_input)
expected_out = expected_crypten_model(dummy_input)
packed = utils.pack_values(expected_crypten_model)
# zero all model's parameters
with th.no_grad():
for p in expected_crypten_model.parameters():
assert isinstance(p, th.Tensor)
p.set_(th.zeros_like(p))
crypten_model = utils.unpack_values(packed, model=expected_crypten_model)
out = crypten_model(dummy_input)
assert th.all(expected_out == out)
def test_pack_typerror_crypten_model():
"""
Testing if we throw an error when trying to unpack an encrypted model,
we should be able to unpack models that are not encrypted
"""
dummy_input = th.rand(1, 28 * 28)
expected_crypten_model = crypten.nn.from_pytorch(ExampleNet(), dummy_input)
expected_crypten_model.encrypted = True
with pytest.raises(TypeError):
packed = utils.pack_values(expected_crypten_model)
def test_pack_tensors(tensors):
packed = utils.pack_values(tensors)
unpacked = utils.unpack_values(packed)
if isinstance(unpacked, tuple): # return tensor1, tensor2 ...
assert len(tensors) == len(unpacked)
for unpacked_tensor, tensor in zip(unpacked, tensors):
assert th.all(unpacked_tensor == tensor)
else: # return tensor
assert th.all(unpacked == tensors)
def test_pack_crypten_model():
dummy_input = th.rand(1, 28 * 28)
expected_crypten_model = crypten.nn.from_pytorch(ExampleNet(), dummy_input)
expected_out = expected_crypten_model(dummy_input)
packed = utils.pack_values(expected_crypten_model)
# zero all model's parameters
with th.no_grad():
for p in expected_crypten_model.parameters():
assert isinstance(p, th.Tensor)
p.set_(th.zeros_like(p))
crypten_model = utils.unpack_values(packed, model=expected_crypten_model)
out = crypten_model(dummy_input)
assert th.all(expected_out == out)
def _launch(func, rank, world_size, master_addr, master_port, queue, func_args,
func_kwargs):
communicator_args = {
"RANK": rank,
"WORLD_SIZE": world_size,
"RENDEZVOUS": "env://",
"MASTER_ADDR": master_addr,
"MASTER_PORT": master_port,
"DISTRIBUTED_BACKEND": "gloo",
}
for key, val in communicator_args.items():
os.environ[key] = str(val)
crypten.init()
return_value = func(*func_args, **func_kwargs)
crypten.uninit()
return_value = utils.pack_values(return_value)
queue.put(return_value)
def test_pack_other():
expected_value = utils.pack_values(42)
assert 42 == utils.unpack_values(expected_value)