def setup(workers): seed_max = 2 ** 32 paired_workers = list(zip(workers, workers[1:])) paired_workers.append((workers[-1], workers[0])) workers_ptr = defaultdict(dict) for cur_worker, next_worker in paired_workers: if cur_worker == syft.local_worker: ptr = cur_worker.torch.randint(-seed_max, seed_max - 1, size=(1,)) ptr_next = ptr.send(next_worker) else: ptr = cur_worker.remote.torch.randint(-seed_max, seed_max - 1, size=(1,)) ptr_next = ptr.copy().move(next_worker) workers_ptr[cur_worker]["cur_seed"] = ptr workers_ptr[next_worker]["prev_seed"] = ptr_next for worker, seeds in workers_ptr.items(): cur_seed = seeds["cur_seed"] prev_seed = seeds["prev_seed"] if worker == syft.local_worker: func = _initialize_generators else: func = remote(_initialize_generators, location=worker) func(cur_seed, prev_seed)
def gen_alpha_3of3(worker): func = None if worker == syft.local_worker: func = _generate_alpha_3of3 else: func = remote(_generate_alpha_3of3, location=worker) return func(worker.id)
def get_random(name_generator, shape, worker): func = None if worker == syft.local_worker: func = _get_random_tensor(name_generator, shape, worker.id) else: func = remote(_get_random_tensor, location=worker) return func(name_generator, shape, worker.id)
def setup(cls, players): seed_map = cls.generate_and_share_seeds(players) for worker, seeds in seed_map.items(): if worker == syft.local_worker: initialize_generators = _initialize_generators else: initialize_generators = remote(_initialize_generators, location=worker) initialize_generators(*seeds)
def test_remote(workers, return_value): alice = workers["alice"] x = th.tensor([1.0]) expected = my_awesome_computation(x) p = x.send(alice) args = (p,) results = remote(my_awesome_computation, location=alice)( *args, return_value=return_value, return_arity=2 ) if not return_value: results = tuple(result.get() for result in results) assert results == expected
def test_remote_wrong_arity(workers, return_value): """ Identical to test_remote except the use didn't set return_arity to be the correct number of return values. Here it should be 2, not 1. """ alice = workers["alice"] x = th.tensor([1.0]) expected = my_awesome_computation(x) p = x.send(alice) args = (p,) results = remote(my_awesome_computation, location=alice)( *args, return_value=return_value, return_arity=1 ) if not return_value: results = tuple(result.get() for result in results) assert results == expected
def fss_op(x1, x2, op="eq"): """ Define the workflow for a binary operation using Function Secret Sharing Currently supported operand are = & <=, respectively corresponding to op = 'eq' and 'comp' Args: x1: first AST x2: second AST op: type of operation to perform, should be 'eq' or 'comp' Returns: shares of the comparison """ if isinstance(x1, sy.AdditiveSharingTensor): locations = x1.locations class_attributes = x1.get_class_attributes() else: locations = x2.locations class_attributes = x2.get_class_attributes() dtype = class_attributes.get("dtype") asynchronous = isinstance(locations[0], WebsocketClientWorker) workers_args = [ ( x1.child[location.id] if isinstance(x1, sy.AdditiveSharingTensor) else (x1 if i == 0 else 0), x2.child[location.id] if isinstance(x2, sy.AdditiveSharingTensor) else (x2 if i == 0 else 0), op, ) for i, location in enumerate(locations) ] try: shares = [] for i, location in enumerate(locations): share = remote(mask_builder, location=location)(*workers_args[i], return_value=True) shares.append(share) except EmptyCryptoPrimitiveStoreError as e: if sy.local_worker.crypto_store.force_preprocessing: raise sy.local_worker.crypto_store.provide_primitives(workers=locations, **e.kwargs_) return fss_op(x1, x2, op) # async has a cost which is too expensive for this command # shares = asyncio.run(sy.local_worker.async_dispatch( # workers=locations, # commands=[ # (full_name(mask_builder), None, workers_args[i], {}) # for i in [0, 1] # ], # return_value=True # )) mask_value = sum(shares) % 2 ** n for location, share in zip(locations, shares): location.de_register_obj(share) del share workers_args = [(th.IntTensor([i]), mask_value, op, dtype) for i in range(2)] if not asynchronous: shares = [] for i, location in enumerate(locations): share = remote(evaluate, location=location)(*workers_args[i], return_value=False) shares.append(share) else: print("async") shares = asyncio.run( sy.local_worker.async_dispatch( workers=locations, commands=[(full_name(evaluate), None, workers_args[i], {}) for i in [0, 1]], ) ) shares = {loc.id: share for loc, share in zip(locations, shares)} response = sy.AdditiveSharingTensor(shares, **class_attributes) return response
def spdz_mul(cmd, x, y, crypto_provider, dtype, torch_dtype, field): """Abstractly multiplies two tensors (mul or matmul) Args: cmd: a callable of the equation to be computed (mul or matmul) x (AdditiveSharingTensor): the left part of the operation y (AdditiveSharingTensor): the right part of the operation crypto_provider (AbstractWorker): an AbstractWorker which is used to generate triples dtype (str): denotes the dtype of the shares, should be 'long' (default), 'int' or 'custom' torch_dtype (torch.dtype): the real type of the shares, should be th.int64 (default) or th.int32 field (int): an integer denoting the size of the field, default is 2**64 Return: an AdditiveSharingTensor """ op = cmd locations = x.locations # Experimental results don't show real improvements with asynchronous = True asynchronous = False # isinstance(locations[0], WebsocketClientWorker) try: shares_delta, shares_epsilon = [], [] for location in locations: args = (x.child[location.id], y.child[location.id], op, dtype, torch_dtype, field, location) share_delta, share_epsilon = remote(spdz_mask, location=location)( *args, return_value=True, return_arity=2) shares_delta.append(share_delta) shares_epsilon.append(share_epsilon) except EmptyCryptoPrimitiveStoreError as e: if sy.local_worker.crypto_store.force_preprocessing: raise crypto_provider.crypto_store.provide_primitives(workers=locations, **e.kwargs_) return spdz_mul(cmd, x, y, crypto_provider, dtype, torch_dtype, field) delta = sum(shares_delta) epsilon = sum(shares_epsilon) for location, share_delta, share_epsilon in zip(locations, shares_delta, shares_epsilon): location.de_register_obj(share_delta) location.de_register_obj(share_epsilon) del share_delta del share_epsilon if not asynchronous: shares = [] for i, location in enumerate(locations): args = (th.LongTensor([i]), delta, epsilon, op, dtype, torch_dtype, field, location) share = remote(spdz_compute, location=location)(*args, return_value=False) shares.append(share) else: shares = asyncio.run( sy.local_worker.async_dispatch( workers=locations, commands=[( full_name(spdz_compute), None, (th.LongTensor([i]), delta, epsilon, op), {}, ) for i in [0, 1]], return_value=False, )) shares = {loc.id: share for loc, share in zip(locations, shares)} response = sy.AdditiveSharingTensor(shares, **x.get_class_attributes()) return response
def _pool2d(input, kernel_size: int = 2, stride: int = 2, padding=0, dilation=1, ceil_mode=None, mode="avg"): if isinstance(kernel_size, tuple): assert kernel_size[0] == kernel_size[1] kernel_size = kernel_size[0] if isinstance(stride, tuple): assert stride[0] == stride[1] stride = stride[0] input_fp = input input = input.child locations = input.locations im_reshaped_shares = {} params = {} for location in locations: input_share = input.child[location.id] im_reshaped_shares[location.id], *params[location.id] = remote( _pre_pool, location=location)(input_share, kernel_size, stride, padding, dilation, return_value=False, return_arity=6) im_reshaped = sy.AdditiveSharingTensor(im_reshaped_shares, **input.get_class_attributes()) if mode == "max": # We have optimisations when the kernel is small, namely a square of size 2 or 3 # to reduce the number of rounds and the total number of comparisons. # See more in Appendice C.3 https://arxiv.org/pdf/2006.04593.pdf def max_half_split(tensor4d, half_size): """ Split the tensor on 2 halves on the last dim and return the maximum half """ left, right = tensor4d[:, :, :, :half_size], tensor4d[:, :, :, half_size:] max_half = left + (right >= left) * (right - left) return max_half if im_reshaped.shape[-1] == 4: # Compute the max as a binary tree: 2 steps are needed for 4 values res = max_half_split(im_reshaped, 2) res = max_half_split(res, 1) elif im_reshaped.shape[-1] == 9: # For 9 values we need 4 steps: we process the 8 first values and then # compute the max with the 9th value res = max_half_split(im_reshaped[:, :, :, :8], 4) res = max_half_split(res, 2) left = max_half_split(res, 1) right = im_reshaped[:, :, :, 8:] res = left + (right >= left) * (right - left) else: res = im_reshaped.max(dim=-1) elif mode == "avg": res = im_reshaped.mean(dim=-1) else: raise ValueError(f"In pool2d, mode should be avg or max, not {mode}.") res_shares = {} for location in locations: res_share = res.child[location.id] res_share = remote(_post_pool, location=location)(res_share, *params[location.id]) res_shares[location.id] = res_share result_fp = sy.FixedPrecisionTensor(**input_fp.get_class_attributes()).on( sy.AdditiveSharingTensor(res_shares, **res.get_class_attributes()), wrap=False) return result_fp
def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1): """ Overloads torch.nn.functional.conv2d to be able to use MPC on convolutional networks. The idea is to unroll the input and weight matrices to compute a matrix multiplication equivalent to the convolution. Args: input: input image weight: convolution kernels bias: optional additive bias stride: stride of the convolution kernels padding: implicit paddings on both sides of the input. dilation: spacing between kernel elements groups: split input into groups, in_channels should be divisible by the number of groups Returns: the result of the convolution as a fixed precision tensor. """ input_fp, weight_fp = input, weight if isinstance(input.child, FrameworkTensor) or isinstance( weight.child, FrameworkTensor): assert isinstance(input.child, FrameworkTensor) assert isinstance(weight.child, FrameworkTensor) im_reshaped, weight_reshaped, *params = _pre_conv( input, weight, bias, stride, padding, dilation, groups) if groups > 1: res = [] chunks_im = torch.chunk(im_reshaped, groups, dim=2) chunks_weights = torch.chunk(weight_reshaped, groups, dim=0) for g in range(groups): tmp = chunks_im[g].matmul(chunks_weights[g]) res.append(tmp) result = torch.cat(res, dim=2) else: result = im_reshaped.matmul(weight_reshaped) result = _post_conv(bias, result, *params) return result.wrap() input, weight = input.child, weight.child if bias is not None: bias = bias.child assert isinstance( bias, sy.AdditiveSharingTensor ), "Have you provided bias as a kwarg? If so, please remove `bias=`." locations = input.locations im_reshaped_shares = {} weight_reshaped_shares = {} params = {} for location in locations: input_share = input.child[location.id] weight_share = weight.child[location.id] bias_share = bias.child[location.id] if bias is not None else None ( im_reshaped_shares[location.id], weight_reshaped_shares[location.id], *params[location.id], ) = remote(_pre_conv, location=location)( input_share, weight_share, bias_share, stride, padding, dilation, groups, return_value=False, return_arity=6, ) im_reshaped = sy.FixedPrecisionTensor( **input_fp.get_class_attributes()).on(sy.AdditiveSharingTensor( im_reshaped_shares, **input.get_class_attributes()), wrap=False) weight_reshaped = sy.FixedPrecisionTensor( **weight_fp.get_class_attributes()).on(sy.AdditiveSharingTensor( weight_reshaped_shares, **input.get_class_attributes()), wrap=False) # Now that everything is set up, we can compute the convolution as a simple matmul if groups > 1: res = [] chunks_im = torch.chunk(im_reshaped, groups, dim=2) chunks_weights = torch.chunk(weight_reshaped, groups, dim=0) for g in range(groups): tmp = chunks_im[g].matmul(chunks_weights[g]) res.append(tmp) res_fp = torch.cat(res, dim=2) res = res_fp.child else: res_fp = im_reshaped.matmul(weight_reshaped) res = res_fp.child # and then we reshape the result res_shares = {} for location in locations: bias_share = bias.child[location.id] if bias is not None else None res_share = res.child[location.id] res_share = remote(_post_conv, location=location)(bias_share, res_share, *params[location.id]) res_shares[location.id] = res_share result_fp = sy.FixedPrecisionTensor(**res_fp.get_class_attributes()).on( sy.AdditiveSharingTensor(res_shares, **res.get_class_attributes()), wrap=False) return result_fp