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)