def test_cross_entropy(self): """Tests cross_entropy and binary_cross_entropy""" sizes = [(3, 2), (8, 4), (5, 10)] losses = [ "binary_cross_entropy", "binary_cross_entropy_with_logits", "cross_entropy", ] for size, loss in itertools.product(sizes, losses): for skip_forward in [False, True]: batch_size, num_targets = size if loss in [ "binary_cross_entropy", "binary_cross_entropy_with_logits" ]: if loss == "binary_cross_entropy": tensor = get_random_test_tensor(size=(batch_size, ), max_value=0.998, is_float=True) tensor = tensor.abs().add_(0.001) else: tensor = get_random_test_tensor(size=(batch_size, ), is_float=True) target = get_random_test_tensor(size=(batch_size, ), is_float=True) target = target.gt(0.0).float() target_encr = crypten.cryptensor(target) else: tensor = get_random_test_tensor(size=size, is_float=True) target = get_random_test_tensor(size=(batch_size, ), max_value=num_targets - 1) target = onehot(target.abs(), num_targets=num_targets) target_encr = crypten.cryptensor(target) # CrypTen, unlike PyTorch, uses one-hot targets target = target.argmax(1) # forward tensor.requires_grad = True tensor_encr = AutogradCrypTensor(crypten.cryptensor(tensor), requires_grad=True) reference = getattr(torch.nn.functional, loss)(tensor, target) out_encr = getattr(tensor_encr, loss)(target_encr, skip_forward=skip_forward) if not skip_forward: self._check(out_encr, reference, f"{loss} forward failed") # backward reference.backward() out_encr.backward() self._check(tensor_encr.grad, tensor.grad, f"{loss} backward failed with")
def test_losses(self): """ Tests all Losses implemented in crypten.nn. """ # create test tensor: input = get_random_test_tensor(max_value=0.999, is_float=True).abs() + 0.001 target = get_random_test_tensor(max_value=0.999, is_float=True).abs() + 0.001 encrypted_input = crypten.cryptensor(input) encrypted_target = crypten.cryptensor(target) # test forward() function of all simple losses: for loss_name in ["BCELoss", "L1Loss", "MSELoss"]: enc_loss_object = getattr(torch.nn, loss_name)() self.assertEqual(enc_loss_object.reduction, "mean", "Reduction used is not 'mean'") loss = getattr(torch.nn, loss_name)()(input, target) encrypted_loss = getattr(crypten.nn, loss_name)()(encrypted_input, encrypted_target) self._check(encrypted_loss, loss, "%s failed" % loss_name) encrypted_loss = getattr(crypten.nn, loss_name)()( AutogradCrypTensor(encrypted_input), AutogradCrypTensor(encrypted_target), ) self._check(encrypted_loss, loss, "%s failed" % loss_name) # test forward() function of cross-entropy loss: batch_size, num_targets = 16, 5 input = get_random_test_tensor(size=(batch_size, num_targets), is_float=True) target = get_random_test_tensor(size=(batch_size, ), max_value=num_targets - 1).abs() encrypted_input = crypten.cryptensor(input) encrypted_target = crypten.cryptensor( onehot(target, num_targets=num_targets)) enc_loss_object = crypten.nn.CrossEntropyLoss() self.assertEqual(enc_loss_object.reduction, "mean", "Reduction used is not 'mean'") loss = torch.nn.CrossEntropyLoss()(input, target) encrypted_loss = crypten.nn.CrossEntropyLoss()(encrypted_input, encrypted_target) self._check(encrypted_loss, loss, "cross-entropy loss failed") encrypted_loss = crypten.nn.CrossEntropyLoss()( AutogradCrypTensor(encrypted_input), AutogradCrypTensor(encrypted_target)) self._check(encrypted_loss, loss, "cross-entropy loss failed")
def test_from_pytorch_training_classification(self): """Tests from_pytorch CrypTen training for classification models""" import torch.nn as nn import torch.nn.functional as F class CNN(nn.Module): def __init__(self): super(CNN, self).__init__() self.conv1 = nn.Conv2d(1, 16, kernel_size=5, padding=1) self.fc1 = nn.Linear(16 * 13 * 13, 100) self.fc2 = nn.Linear(100, 2) def forward(self, x): out = self.conv1(x) out = F.relu(out) out = F.max_pool2d(out, 2) out = out.view(-1, 16 * 13 * 13) out = self.fc1(out) out = F.relu(out) out = self.fc2(out) out = F.softmax(out, dim=1) return out model_plaintext = CNN() batch_size = 5 x_orig = get_random_test_tensor(size=(batch_size, 1, 28, 28), is_float=True) y_orig = (get_random_test_tensor(size=(batch_size, 1), is_float=True).gt(0).long()) y_one_hot = onehot(y_orig, num_targets=2) # encrypt training sample: x_train = crypten.cryptensor(x_orig, requires_grad=True) y_train = crypten.cryptensor(y_one_hot) dummy_input = torch.empty((1, 1, 28, 28)) for loss_name in ["BCELoss", "CrossEntropyLoss"]: # create encrypted model model = crypten.nn.from_pytorch(model_plaintext, dummy_input) model.train() model.encrypt() self._check_training(model, x_train, y_train, loss_name) self._check_model_export(model, x_train)
def test_custom_module_training(self): """Tests training CrypTen models created directly using the crypten.nn.Module""" BATCH_SIZE = 32 NUM_FEATURES = 3 class ExampleNet(crypten.nn.Module): def __init__(self): super(ExampleNet, self).__init__() self.fc1 = crypten.nn.Linear(NUM_FEATURES, BATCH_SIZE) self.fc2 = crypten.nn.Linear(BATCH_SIZE, 2) def forward(self, x): out = self.fc1(x) out = self.fc2(out) return out model = ExampleNet() x_orig = get_random_test_tensor(size=(BATCH_SIZE, NUM_FEATURES), is_float=True) # y is a linear combo of x to ensure network can easily learn pattern y_orig = (2 * x_orig.mean(dim=1)).gt(0).long() y_one_hot = onehot(y_orig, num_targets=2) # encrypt training sample: x_train = crypten.cryptensor(x_orig, requires_grad=True) y_train = crypten.cryptensor(y_one_hot) for loss_name in ["BCELoss", "CrossEntropyLoss", "MSELoss"]: # create loss function loss = getattr(crypten.nn, loss_name)() # create encrypted model model.train() model.encrypt() num_epochs = 3 learning_rate = 0.001 for i in range(num_epochs): output = model(x_train) if loss_name == "MSELoss": output_norm = output else: output_norm = output.softmax(1) loss_value = loss(output_norm, y_train) # set gradients to "zero" model.zero_grad() for param in model.parameters(): self.assertIsNone(param.grad, "zero_grad did not reset gradients") # perform backward pass: loss_value.backward() for param in model.parameters(): if param.requires_grad: self.assertIsNotNone( param.grad, "required parameter gradient not created") # update parameters orig_parameters, upd_parameters = {}, {} orig_parameters = self._compute_reference_parameters( "", orig_parameters, model, 0) model.update_parameters(learning_rate) upd_parameters = self._compute_reference_parameters( "", upd_parameters, model, learning_rate) parameter_changed = False for name, value in orig_parameters.items(): if param.requires_grad and param.grad is not None: unchanged = torch.allclose(upd_parameters[name], value) if unchanged is False: parameter_changed = True self.assertTrue( parameter_changed, "no parameter changed in training step") # record initial and current loss if i == 0: orig_loss = loss_value.get_plain_text() curr_loss = loss_value.get_plain_text() # check that the loss has decreased after training self.assertTrue( curr_loss.item() < orig_loss.item(), "loss has not decreased after training", )
def test_losses(self): """ Tests all Losses implemented in crypten.nn. """ # create test tensor: input = get_random_test_tensor(max_value=0.999, is_float=True).abs() + 0.001 target = get_random_test_tensor(max_value=0.999, is_float=True).abs() + 0.001 encrypted_input = crypten.cryptensor(input) encrypted_target = crypten.cryptensor(target) # test forward() function of all simple losses: for loss_name in ["BCELoss", "BCEWithLogitsLoss", "L1Loss", "MSELoss"]: for skip_forward in [False, True]: enc_loss_object = getattr(torch.nn, loss_name)() self.assertEqual(enc_loss_object.reduction, "mean", "Reduction used is not 'mean'") input.requires_grad = True input.grad = None loss = getattr(torch.nn, loss_name)()(input, target) if not skip_forward: encrypted_loss = getattr(crypten.nn, loss_name)()(encrypted_input, encrypted_target) self._check(encrypted_loss, loss, "%s failed" % loss_name) encrypted_input.requires_grad = True encrypted_input.grad = None encrypted_loss = getattr(crypten.nn, loss_name)(skip_forward=skip_forward)( encrypted_input, encrypted_target) if not skip_forward: self._check(encrypted_loss, loss, "%s failed" % loss_name) # Check backward loss.backward() encrypted_loss.backward() self._check(encrypted_input.grad, input.grad, "%s grad failed" % loss_name) # test forward() function of cross-entropy loss: batch_size, num_targets = 16, 5 input = get_random_test_tensor(size=(batch_size, num_targets), is_float=True) target = get_random_test_tensor(size=(batch_size, ), max_value=num_targets - 1).abs() encrypted_input = crypten.cryptensor(input) encrypted_target = crypten.cryptensor( onehot(target, num_targets=num_targets)) enc_loss_object = crypten.nn.CrossEntropyLoss() self.assertEqual(enc_loss_object.reduction, "mean", "Reduction used is not 'mean'") loss = torch.nn.CrossEntropyLoss()(input, target) encrypted_loss = crypten.nn.CrossEntropyLoss()(encrypted_input, encrypted_target) self._check(encrypted_loss, loss, "cross-entropy loss failed") encrypted_input.requires_grad = True encrypted_target.requires_grad = True encrypted_loss = crypten.nn.CrossEntropyLoss()(encrypted_input, encrypted_target) self._check(encrypted_loss, loss, "cross-entropy loss failed")
def test_from_pytorch_training(self): """Tests the from_pytorch code path for training CrypTen models""" import torch.nn as nn import torch.nn.functional as F class ExampleNet(nn.Module): def __init__(self): super(ExampleNet, self).__init__() self.conv1 = nn.Conv2d(1, 16, kernel_size=5, padding=1) self.fc1 = nn.Linear(16 * 13 * 13, 100) self.fc2 = nn.Linear(100, 2) def forward(self, x): out = self.conv1(x) out = F.relu(out) out = F.max_pool2d(out, 2) out = out.view(out.size(0), -1) out = self.fc1(out) out = F.relu(out) out = self.fc2(out) return out model_plaintext = ExampleNet() batch_size = 5 x_orig = get_random_test_tensor(size=(batch_size, 1, 28, 28), is_float=True) y_orig = (get_random_test_tensor(size=(batch_size, 1), is_float=True).gt(0).long()) y_one_hot = onehot(y_orig, num_targets=2) # encrypt training sample: x_train = AutogradCrypTensor(crypten.cryptensor(x_orig)) y_train = crypten.cryptensor(y_one_hot) dummy_input = torch.empty((1, 1, 28, 28)) for loss_name in ["BCELoss", "CrossEntropyLoss", "MSELoss"]: # create loss function loss = getattr(crypten.nn, loss_name)() # create encrypted model model = crypten.nn.from_pytorch(model_plaintext, dummy_input) model.train() model.encrypt() num_epochs = 3 learning_rate = 0.001 for i in range(num_epochs): output = model(x_train) if loss_name == "MSELoss": output_norm = output else: output_norm = output.softmax(1) loss_value = loss(output_norm, y_train) # set gradients to "zero" model.zero_grad() for param in model.parameters(): self.assertIsNone(param.grad, "zero_grad did not reset gradients") # perform backward pass: loss_value.backward() for param in model.parameters(): if param.requires_grad: self.assertIsNotNone( param.grad, "required parameter gradient not created") # update parameters orig_parameters, upd_parameters = {}, {} orig_parameters = self._compute_reference_parameters( "", orig_parameters, model, 0) model.update_parameters(learning_rate) upd_parameters = self._compute_reference_parameters( "", upd_parameters, model, learning_rate) # FIX check that any parameter with a non-zero gradient has changed?? parameter_changed = False for name, value in orig_parameters.items(): if param.requires_grad and param.grad is not None: unchanged = torch.allclose(upd_parameters[name], value) if unchanged is False: parameter_changed = True self.assertTrue( parameter_changed, "no parameter changed in training step") # record initial and current loss if i == 0: orig_loss = loss_value.get_plain_text() curr_loss = loss_value.get_plain_text() # check that the loss has decreased after training self.assertTrue( curr_loss.item() < orig_loss.item(), "loss has not decreased after training", )
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)