def test_autograd_registation(self): """Tests registration of new autograd function.""" # check that get_grad_fn() returns correct functions: for func_name, reference_func in gradients.FUNCTION_REGISTRY.items(): grad_fn = gradients.get_grad_fn(func_name) self.assertEqual(grad_fn, reference_func) self.assertEqual(grad_fn.name, func_name) # check that non-existing functions return None: for invalid_func_name in ["bfobofb", "djhfhr"]: func = gradients.get_grad_fn(invalid_func_name) self.assertIsNone(func) # check that registering new classes works: for func_name in ["mock_func1", "mock_func2", "mock_func3"]: cls = type("%sName" % func_name, (AutogradFunction,), {}) gradients.register_function(func_name)(cls) grad_fn = gradients.get_grad_fn(func_name) self.assertEqual(grad_fn, cls) self.assertEqual(grad_fn.name, func_name) # check that existing functions cannot be overwritten: for func_name in ["add", "sub", "view"]: cls = type("%sName" % func_name, (AutogradFunction,), {}) with self.assertRaises(ValueError): gradients.register_function(func_name)(cls)
def test_autograd_func_take(self): """Tests the part of autograd take that does not have a torch equivalent""" tensor_size = [5, 5, 5, 5] index = torch.tensor([[[1, 2], [3, 4]], [[4, 2], [1, 3]]], dtype=torch.long) # Test when dimension!=None for dimension in range(0, 4): tensor = get_random_test_tensor(size=tensor_size, is_float=True) ref_forward = torch.from_numpy(tensor.numpy().take(index, dimension)) encrypted_tensor = crypten.cryptensor(tensor) encr_inputs = [encrypted_tensor, index, dimension] # test forward ctx = AutogradContext() grad_fn_take = gradients.get_grad_fn("take") encr_output = grad_fn_take.forward(ctx, encr_inputs) self._check(encr_output, ref_forward, "take forward failed: dimension set") # test backward: # first, recreate take forward function with only torch operations tensor2 = get_random_test_tensor(size=tensor_size, is_float=True) tensor2.requires_grad = True all_indices = [slice(0, x) for x in tensor2.size()] all_indices[dimension] = index ref_forward_torch = tensor2[all_indices] grad_output = torch.ones(ref_forward_torch.size()) ref_forward_torch.backward(grad_output) # next, do backward pass on encrypted tensor encr_grad_output = encr_output.new(grad_output) encr_grad = grad_fn_take.backward(ctx, encr_grad_output) # finally, compare values self._check(encr_grad, tensor2.grad, "take backward failed: dimension set")
def test_batchnorm(self): """ Tests batchnorm forward and backward steps with training on / off. """ # sizes for 1D, 2D, and 3D batchnorm # batch_size (dim=0) > 500 and increase tolerance to avoid flaky precision # errors in inv_var, which involves sqrt and reciprocal sizes = [(800, 5), (500, 8, 15), (600, 10, 3, 15)] tolerance = 0.5 for size in sizes: for is_trainning in (False, True): tensor = get_random_test_tensor(size=size, is_float=True) tensor.requires_grad = True encrypted_input = crypten.cryptensor(tensor) C = size[1] weight = get_random_test_tensor(size=[C], max_value=1, is_float=True) bias = get_random_test_tensor(size=[C], max_value=1, is_float=True) weight.requires_grad = True bias.requires_grad = True # dimensions for mean and variance stats_dimensions = list(range(tensor.dim())) # perform on C dimension for tensor of shape (N, C, +) stats_dimensions.pop(1) running_mean = tensor.mean(stats_dimensions).detach() running_var = tensor.var(stats_dimensions).detach() enc_running_mean = encrypted_input.mean(stats_dimensions) enc_running_var = encrypted_input.var(stats_dimensions) reference = torch.nn.functional.batch_norm(tensor, running_mean, running_var, weight=weight, bias=bias) encrypted_input = AutogradCrypTensor(encrypted_input) ctx = AutogradContext() batch_norm_fn = gradients.get_grad_fn("batchnorm") encrypted_out = batch_norm_fn.forward( ctx, (encrypted_input, weight, bias), training=is_trainning, running_mean=enc_running_mean, running_var=enc_running_var, ) # check forward self._check( encrypted_out, reference, "batchnorm forward failed with trainning " f"{is_trainning} on {tensor.dim()}-D", tolerance=tolerance, ) # check backward (input, weight, and bias gradients) reference.backward(reference) encrypted_grad = batch_norm_fn.backward(ctx, encrypted_out) TorchGrad = namedtuple("TorchGrad", ["name", "value"]) torch_gradients = [ TorchGrad("input gradient", tensor.grad), TorchGrad("weight gradient", weight.grad), TorchGrad("bias gradient", bias.grad), ] for i, torch_gradient in enumerate(torch_gradients): self._check( encrypted_grad[i], torch_gradient.value, f"batchnorm backward {torch_gradient.name} failed" f"with training {is_trainning} on {tensor.dim()}-D", tolerance=tolerance, )
def test_autograd_functions(self): """Tests individual autograd functions without testing autograd.""" # input sizes for tests of autograd functions: input_size = { "t": (2, 4), "transpose": (4, 8, 3), "flip": (2, 3, 7, 2), "view": (8, 6), "reshape": (8, 6), "flatten": (8, 6), "narrow": (10, 7), "take": (5, 10, 15), # NOTE: this only tests the pytorch take # functionality. The remaining take functionality # is tested separately "gather": (2, 2), "scatter": (3, 5), "roll": (4, 8), "squeeze": (12, 1, 6), "unsqueeze": (7, 3), "__getitem__": (6, 6), "neg": (8, 4), "relu": (3, 7), "tanh": (4, 3), "add": (10, 7), "sub": (9, 2), "mul": (3, 5), "matmul": (7, 7), "div": (5, 4), "pow": (4, 3), "square": (8, 5), "sqrt": (5, 6), "exp": (5, 2), "log": (3, 7), "dot": (8, ), "ger": (12, ), "sin": (5, 4), "cos": (9, 3), "abs": (8, 5), "sign": (8, 5), "norm": (3, 2), # NOTE: Flaky because sqrt only works for values up to 200. "sum": (4, 3), "cumsum": (13, 7), "trace": (4, 4), "mean": (2, 9), "var": (3, 4), "max": (6, 7), "min": (4, 5), "sigmoid": (4, 7), "softmax": (10, 5), "pad": (6, 3), # "avg_pool2d": (1, 3, 21, 21), # TODO: Enable once avg_pool2d is # fixed in gradients.py. "max_pool2d": (1, 3, 21, 21), "conv2d": (1, 4, 21, 21), "binary_cross_entropy": (8, ), "cross_entropy": (8, 4), } additional_args = { "transpose": [2, 0], "flip": [(1, 3, 2)], "view": [(4, 12)], "reshape": [(4, 12)], "narrow": [1, 2, 3], "gather": [1, torch.tensor([[0, 0], [1, 0]])], "scatter": [ 0, torch.tensor([[0, 1, 2, 0, 0], [2, 0, 0, 1, 2]]), get_random_test_tensor(size=(2, 5), is_float=True), ], "roll": [(2, -1), (0, 1)], "squeeze": [1], "unsqueeze": [1], "__getitem__": [1], "div": [4.0], "pow": [2.0], "cumsum": [1], "softmax": [1], "pad": [(1, 2, 3, 4)], "avg_pool2d": [5], "max_pool2d": [3], "conv2d": [get_random_test_tensor(size=(2, 4, 3, 3), is_float=True)], "take": [torch.tensor([0, 5, 10])], "binary_cross_entropy": [ get_random_test_tensor(size=(8, ), is_float=True).gt(0.0).float() ], "cross_entropy": [ onehot(get_random_test_tensor(size=(8, ), max_value=3).abs(), num_targets=4) ], } binary_functions = ["add", "sub", "mul", "dot", "ger", "matmul"] positive_only = ["pow", "sqrt", "log", "binary_cross_entropy"] # loop over all autograd functions: for func_name in input_size.keys(): # generate inputs: inputs = [ get_random_test_tensor(size=input_size[func_name], max_value=1.0, is_float=True) for _ in range(2 if func_name in binary_functions else 1) ] if func_name in positive_only: # some functions do not take negative values inputs = [input.abs().add_(0.001) for input in inputs] for input in inputs: input.requires_grad = True encr_inputs = [crypten.cryptensor(input) for input in inputs] number_of_inputs = len(inputs) # add additional arguments, encrypting only tensors (if found): if func_name in additional_args: inputs += additional_args[func_name] encr_inputs += additional_args[func_name] if func_name == "take": encr_inputs += [None] elif func_name not in ["gather", "scatter"]: encr_inputs = [ crypten.cryptensor(t) if torch.is_tensor(t) else t for t in encr_inputs ] # cross_entropy uses one-hot targets in crypten but not in PyTorch: if func_name == "cross_entropy": inputs[1] = inputs[1].argmax(1) # AutogradFunction.forward() does not accept unpacked inputs: if len(encr_inputs) == 1: encr_inputs = encr_inputs[0] # test forward function: if hasattr(inputs[0], func_name): # torch.function() reference = getattr(inputs[0], func_name)(*inputs[1:]) elif hasattr(F, func_name): # torch.nn.functional.function() reference = getattr(F, func_name)(*inputs) elif func_name == "square": reference = inputs[0].pow(2.0) else: raise ValueError("unknown PyTorch function: %s" % func_name) ctx = AutogradContext() grad_fn = gradients.get_grad_fn(func_name) encr_output = grad_fn.forward(ctx, encr_inputs) self._check(encr_output, reference, "%s forward failed" % func_name) if func_name == "view": ctx = AutogradContext() # check view() with a list of int to represent size. # encr_inputs[0]: input # encr_inputs[1]: tuple as torch.Size, to be unpacked. view_input, sizes = encr_inputs encr_output = grad_fn.forward(ctx, [view_input] + [size for size in sizes]) self._check(encr_output, reference, "%s forward failed" % func_name) # run backward functions: grad_output = get_random_test_tensor(max_value=2, size=reference.size(), is_float=True) encr_grad_output = encr_output.new(grad_output) reference.backward(grad_output) encr_grad = grad_fn.backward(ctx, encr_grad_output) # test result of running backward function: if not isinstance(encr_grad, (list, tuple)): encr_grad = (encr_grad, ) for idx in range(number_of_inputs): self._check(encr_grad[idx], inputs[idx].grad, "%s backward failed" % func_name)