def test_pos_pow(self):
"""Test gradient crypten pos_pow"""
for power in [3, -2, 1.75]:
# ensure base is positive for pos_pow
tensor = get_random_test_tensor(is_float=True, max_value=2) + 4
tensor.requires_grad = True
tensor_encr = AutogradCrypTensor(
crypten.cryptensor(tensor), requires_grad=True
)
reference = tensor.pow(power)
out_encr = tensor_encr.pos_pow(power)
self._check(
out_encr, reference, f"pos_pow forward failed with power {power}"
)
grad_out = get_random_test_tensor(is_float=True)
grad_out_encr = crypten.cryptensor(grad_out)
reference.backward(grad_out)
out_encr.backward(grad_out_encr)
self._check(
tensor_encr.grad,
tensor.grad,
f"pos_pow backward failed with power {power}",
)
def train_encrypted(
x_encrypted,
y_encrypted,
encrypted_model,
num_epochs,
learning_rate,
batch_size,
print_freq,
):
rank = comm.get().get_rank()
loss = crypten.nn.MSELoss()
num_samples = x_encrypted.size(0)
label_eye = torch.eye(2)
for epoch in range(num_epochs):
last_progress_logged = 0
# only print from rank 0 to avoid duplicates for readability
if rank == 0:
print(f"Epoch {epoch} in progress:")
for j in range(0, num_samples, batch_size):
# define the start and end of the training mini-batch
start, end = j, min(j + batch_size, num_samples)
# construct AutogradCrypTensors out of training examples
x_train = AutogradCrypTensor(x_encrypted[start:end])
y_one_hot = label_eye[y_encrypted[start:end]]
y_train = AutogradCrypTensor(crypten.cryptensor(y_one_hot))
# perform forward pass:
output = encrypted_model(x_train)
loss_value = loss(output, y_train)
# backprop
encrypted_model.zero_grad()
loss_value.backward()
encrypted_model.update_parameters(learning_rate)
# log progress
if j + batch_size - last_progress_logged >= print_freq:
last_progress_logged += print_freq
print(f"Loss {loss_value.get_plain_text().item():.4f}")
# compute accuracy every epoch
pred = output.get_plain_text().argmax(1)
correct = pred.eq(y_encrypted[start:end])
correct_count = correct.sum(0, keepdim=True).float()
accuracy = correct_count.mul_(100.0 / output.size(0))
loss_plaintext = loss_value.get_plain_text().item()
print(f"Epoch {epoch} completed: "
f"Loss {loss_plaintext:.4f} Accuracy {accuracy.item():.2f}")
def _check_forward_backward(
self, func_name, input_tensor, *args, torch_func_name=None, msg=None, **kwargs
):
"""Checks forward and backward against PyTorch
Args:
func_name (str): PyTorch/CrypTen function name
input_tensor (torch.tensor): primary input
args (list): contains arguments for function
msg (str): additional message for mismatch
kwargs (list): keyword arguments for function
"""
if msg is None:
msg = f"{func_name} grad_fn incorrect"
input = input_tensor.clone()
input.requires_grad = True
input_encr = AutogradCrypTensor(crypten.cryptensor(input), requires_grad=True)
for private in [False, True]:
input.grad = None
input_encr.grad = None
args = self._set_grad_to_zero(args)
args_encr = self._set_grad_to_zero(list(args), make_private=private)
# obtain torch function
if torch_func_name is not None:
torch_func = self._get_torch_func(torch_func_name)
else:
torch_func = self._get_torch_func(func_name)
reference = torch_func(input, *args, **kwargs)
encrypted_out = getattr(input_encr, func_name)(*args_encr, **kwargs)
# extract argmax output for max / min with keepdim=False
if isinstance(encrypted_out, (list, tuple)):
reference = reference[0]
encrypted_out = encrypted_out[0]
self._check(encrypted_out, reference, msg + " in forward")
# check backward pass
grad_output = get_random_test_tensor(
max_value=2, size=reference.size(), is_float=True
)
grad_output_encr = crypten.cryptensor(grad_output)
reference.backward(grad_output)
encrypted_out.backward(grad_output_encr)
self._check(input_encr.grad, input.grad, msg + " in backward")
for i, arg_encr in enumerate(args_encr):
if crypten.is_encrypted_tensor(arg_encr):
self._check(arg_encr.grad, args[i].grad, msg + " in backward args")
def test_training(self):
"""
Tests training of simple model in crypten.nn.
"""
# create MLP with one hidden layer:
learning_rate = 0.1
batch_size, num_inputs, num_intermediate, num_outputs = 8, 10, 5, 1
model = crypten.nn.Sequential([
crypten.nn.Linear(num_inputs, num_intermediate),
crypten.nn.ReLU(),
crypten.nn.Linear(num_intermediate, num_outputs),
])
model.train()
model.encrypt()
loss = crypten.nn.MSELoss()
# perform training iterations:
for _ in range(10):
for wrap in [True, False]:
# get training sample:
input = get_random_test_tensor(size=(batch_size, num_inputs),
is_float=True)
target = input.mean(dim=1, keepdim=True)
# encrypt training sample:
input = crypten.cryptensor(input)
target = crypten.cryptensor(target)
if wrap:
input = AutogradCrypTensor(input)
target = AutogradCrypTensor(target)
# perform forward pass:
output = model(input)
loss_value = loss(output, target)
# set gradients to "zero" (setting to None is more efficient):
model.zero_grad()
for param in model.parameters():
self.assertIsNone(param.grad,
"zero_grad did not reset gradients")
# perform backward pass:
loss_value.backward()
# perform parameter update:
reference = {}
reference = self._compute_reference_parameters(
"", reference, model, learning_rate)
model.update_parameters(learning_rate)
self._check_reference_parameters("", reference, model)
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_graph(self):
"""
Tests crypten.nn.Graph module.
"""
for wrap in [True, False]:
# define test case:
input_size = (3, 10)
input = get_random_test_tensor(size=input_size, is_float=True)
encr_input = crypten.cryptensor(input)
if wrap:
encr_input = AutogradCrypTensor(encr_input)
# test residual block with subsequent linear layer:
graph = crypten.nn.Graph("input", "output")
linear1 = get_random_linear(input_size[1], input_size[1])
linear2 = get_random_linear(input_size[1], input_size[1])
graph.add_module("linear", crypten.nn.from_pytorch(linear1, input),
["input"])
graph.add_module("residual", crypten.nn.Add(), ["input", "linear"])
graph.add_module("output", crypten.nn.from_pytorch(linear2, input),
["residual"])
graph.encrypt()
# check container:
self.assertTrue(graph.encrypted, "nn.Graph not encrypted")
for module in graph.modules():
self.assertTrue(module.encrypted, "module not encrypted")
assert (sum(1 for _ in graph.modules()) == 3
), "nn.Graph contains incorrect number of modules"
# compare output to reference:
encr_output = graph(encr_input)
reference = linear2(linear1(input) + input)
self._check(encr_output, reference, "nn.Graph forward failed")
def test_autograd(self):
"""Tests autograd graph construction and backprop."""
# define test cases:
tests = [
(1, ["relu", "neg", "relu", "sum"]),
(2, ["t", "neg", "add", "sum"]),
(2, ["relu", "mul", "t", "sum"]),
]
binary_functions = ["add", "sub", "mul", "dot", "matmul"]
# PyTorch test case:
for test in tests:
# get test case:
number_of_inputs, ops = test
inputs = [
get_random_test_tensor(size=(12, 5), is_float=True)
for _ in range(number_of_inputs)
]
encr_inputs = [crypten.cryptensor(input) for input in inputs]
# get autograd variables:
for input in inputs:
input.requires_grad = True
encr_inputs = [AutogradCrypTensor(encr_input) for encr_input in encr_inputs]
# perform forward pass, logging all intermediate outputs:
outputs, encr_outputs = [inputs], [encr_inputs]
for op in ops:
# get inputs for current operation:
input, output = outputs[-1], []
encr_input, encr_output = encr_outputs[-1], []
# apply current operation:
if op in binary_functions: # combine outputs via operation
output.append(getattr(input[0], op)(input[1]))
encr_output.append(getattr(encr_input[0], op)(encr_input[1]))
else:
for idx in range(len(input)):
output.append(getattr(input[idx], op)())
encr_output.append(getattr(encr_input[idx], op)())
# keep references to outputs of operation:
outputs.append(output)
encr_outputs.append(encr_output)
# check output of forward pass:
output, encr_output = outputs[-1][0], encr_outputs[-1][0]
self._check(encr_output._tensor, output, "forward failed")
self.assertTrue(encr_output.requires_grad, "requires_grad incorrect")
# perform backward pass:
output.backward()
encr_output.backward()
# test result of running backward function:
for idx in range(number_of_inputs):
self._check(encr_inputs[idx].grad, inputs[idx].grad, "backward failed")
def _to_autograd(args):
"""
Recursively converts inputs to AutogradCrypTensors.
"""
# convert tuples to lists to allow changes:
convert_to_tuple = False
if isinstance(args, tuple):
args = list(args)
convert_to_tuple = True
# wrap all input tensors in AutogradCrypTensor:
for idx in range(len(args)):
if isinstance(args[idx],
(list, tuple)): # input may be list of tensors
args[idx] = _to_autograd(args[idx])
elif isinstance(args[idx], AutogradCrypTensor) or args[idx] is None:
pass
elif isinstance(args[idx], crypten.CrypTensor):
args[idx] = AutogradCrypTensor(args[idx])
else:
raise ValueError(
"Cannot convert type {} to AutogradCrypTensor.".format(
type(args[idx])))
# return:
if convert_to_tuple:
args = tuple(args)
return args
def test_dropout(self):
"""Tests forward and backward passes for dropout"""
# Create a separate test for dropout since it cannot use the
# regular forward function
all_prob_values = [x * 0.2 for x in range(0, 5)]
for dropout_fn in [
"dropout", "_feature_dropout", "dropout2d", "dropout3d"
]:
for prob in all_prob_values:
for size in [(5, 10), (5, 10, 15), (5, 10, 15, 20)]:
for use_zeros in [False, True]:
tensor = get_random_test_tensor(size=size,
ex_zero=True,
min_value=1.0,
is_float=True)
if use_zeros:
# turn the first row to all zeros
index = [1] + [
slice(0, tensor.size(i))
for i in range(1, tensor.dim())
]
tensor[index] = 0.0
encr_tensor = AutogradCrypTensor(
crypten.cryptensor(tensor), requires_grad=True)
encr_tensor_out = getattr(encr_tensor,
dropout_fn)(p=prob)
dropout_tensor = encr_tensor_out.get_plain_text()
# Check the scaling for non-zero elements
scaled_tensor = tensor / (1 - prob)
reference = dropout_tensor.where(
dropout_tensor == 0, scaled_tensor)
self._check(
encr_tensor_out,
reference,
"dropout failed with size {}, use_zeros {}, and "
"probability {}".format(size, use_zeros, prob),
)
# Check backward pass
grad_output = get_random_test_tensor(max_value=2,
size=reference.size(),
is_float=True)
grad_output_encr = crypten.cryptensor(grad_output)
# Do not check backward if pytorch backward fails
try:
reference.backward(grad_output)
except RuntimeError:
logging.info("skipped")
continue
encr_tensor_out.backward(grad_output_encr)
self._check(
encr_tensor.grad,
input.grad,
"dropout failed in backward with size {}, use_zeros {} and "
"probability {}".format(size, use_zeros, prob),
)
def _set_grad_to_zero(self, args, make_private=False):
"""Sets gradients for args to zero
Args:
args (list of torch.tensors): contains arguments
make_private (bool): encrypt args using AutogradCrypTensor
"""
args_zero_grad = []
for arg in args:
if is_float_tensor(arg) and make_private:
arg = AutogradCrypTensor(crypten.cryptensor(arg), requires_grad=True)
elif is_float_tensor(arg):
arg.requires_grad = True
arg.grad = None
args_zero_grad.append(arg)
return args_zero_grad
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 register_parameter(self, name, param, requires_grad=True):
"""Register parameter in the module."""
if name in self._parameters or hasattr(self, name):
raise ValueError("Parameter or field %s already exists." % name)
if torch.is_tensor(param): # unencrypted model
param.requires_grad = requires_grad
self._parameters[name] = param
else: # encryped model
self._parameters[name] = AutogradCrypTensor(
param, requires_grad=requires_grad)
setattr(self, name, param)
def test_batchnorm_module(self):
"""Test module correctly sets and updates running stats"""
batchnorm_fn_and_size = (
("BatchNorm1d", (500, 10, 3)),
("BatchNorm2d", (600, 7, 4, 20)),
("BatchNorm3d", (800, 5, 4, 8, 15)),
)
for batchnorm_fn, size in batchnorm_fn_and_size:
for is_trainning in (True, False):
tensor = get_random_test_tensor(size=size, is_float=True)
tensor.requires_grad = True
encrypted_input = AutogradCrypTensor(
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)
# check running stats initial
enc_model = getattr(crypten.nn.module,
batchnorm_fn)(C).encrypt()
plain_model = getattr(torch.nn.modules, batchnorm_fn)(C)
stats = ["running_var", "running_mean"]
for stat in stats:
self._check(
enc_model._buffers[stat],
plain_model._buffers[stat],
f"{stat} initial module value incorrect",
)
# set trainning mode
plain_model.training = is_trainning
enc_model.training = is_trainning
# check running_stats update
enc_model.forward(encrypted_input)
plain_model.forward(tensor)
for stat in stats:
self._check(
enc_model._buffers[stat],
plain_model._buffers[stat],
f"{stat} momentum update in module incorrect",
)
def test_autograd_repetition(self):
"""Tests running autograd on the same input repeatedly."""
# create test case:
input = get_random_test_tensor(size=(12, 5), is_float=True)
input.requires_grad = True
encr_input = AutogradCrypTensor(crypten.cryptensor(input))
# re-use the same input multiple times:
for _ in range(7):
# perform forward pass:
output = input.exp().sum()
encr_output = encr_input.exp().sum()
self._check(encr_output._tensor, output, "forward failed")
self.assertTrue(encr_output.requires_grad, "requires_grad incorrect")
# perform backward computation:
output.backward()
encr_output.backward()
self._check(encr_input.grad, input.grad, "backward failed")
def test_detach(self):
"""Tests that detach() works as expected."""
for func_name in ["detach", "detach_"]:
# get test case:
input_size = (12, 5)
input1 = get_random_test_tensor(size=input_size, is_float=True)
input2 = get_random_test_tensor(size=input_size, is_float=True)
input1 = AutogradCrypTensor(crypten.cryptensor(input1))
input2 = AutogradCrypTensor(crypten.cryptensor(input2))
# perform forward computation with detach in the middle:
intermediate = input1.add(1.0)
intermediate = getattr(intermediate, func_name)()
output = intermediate.add(input2).sum()
# perform backward:
output.backward()
msg = "detach() function does not behave as expected"
self.assertIsNone(output.grad, msg)
self.assertIsNone(intermediate.grad, msg)
self.assertIsNone(input1.grad, msg)
self.assertIsNotNone(input2.grad, msg)
def test_square(self):
"""Tests square function gradient.
Note: torch pow(2) is used to verify gradient,
since PyTorch does not implement square().
"""
for size in SIZES:
tensor = get_random_test_tensor(size=size, is_float=True)
tensor.requires_grad = True
tensor_encr = AutogradCrypTensor(crypten.cryptensor(tensor),
requires_grad=True)
out = tensor.pow(2)
out_encr = tensor_encr.square()
self._check(out_encr, out,
f"square forward failed with size {size}")
grad_output = get_random_test_tensor(size=out.shape, is_float=True)
out.backward(grad_output)
out_encr.backward(crypten.cryptensor(grad_output))
self._check(
tensor_encr.grad,
tensor.grad,
f"square backward failed with size {size}",
)
def encrypt(self, mode=True, src=0):
"""Encrypts the model."""
if mode != self.encrypted:
# encrypt / decrypt parameters:
self.encrypted = mode
for name, param in self.named_parameters(recurse=False):
requires_grad = param.requires_grad
if mode: # encrypt parameter
self.set_parameter(
name,
AutogradCrypTensor(
crypten.cryptensor(param, **{"src": src}),
requires_grad=requires_grad,
),
)
else: # decrypt parameter
self.set_parameter(name, param.get_plain_text())
self._parameters[name].requires_grad = requires_grad
# encrypt / decrypt buffers:
for name, buffer in self.named_buffers(recurse=False):
if mode: # encrypt buffer
self.set_buffer(
name,
AutogradCrypTensor(
crypten.cryptensor(buffer, **{"src": src}),
requires_grad=False,
),
)
else: # decrypt buffer
self.set_buffer(name, buffer.get_plain_text())
# apply encryption recursively:
return self._apply(lambda m: m.encrypt(mode=mode, src=src))
return self
def test_non_differentiable_marking(self):
"""Tests whether marking of non-differentiability works correctly."""
# generate random inputs:
inputs = [get_random_test_tensor(is_float=True) for _ in range(5)]
inputs = [crypten.cryptensor(input) for input in inputs]
ctx = AutogradContext()
# repeat test multiple times:
for _ in range(10):
# mark non-differentiable inputs as such:
differentiable = [
random.random() > 0.5 for _ in range(len(inputs))
]
for idx, diff in enumerate(differentiable):
if not diff:
ctx.mark_non_differentiable(inputs[idx])
# check that inputs were correctly marked:
for idx, input in enumerate(inputs):
self.assertEqual(
ctx.is_differentiable(input),
differentiable[idx],
"marking of differentiability failed",
)
ctx.reset()
# test behavior of AutogradCrypTensor:
input = AutogradCrypTensor(inputs[0])
reference = [True, True, False]
for func_name in ["min", "max"]:
outputs = [None] * 3
outputs[0] = getattr(input, func_name)()
outputs[1], outputs[2] = getattr(input, func_name)(dim=0)
for idx, output in enumerate(outputs):
self.assertEqual(
output.requires_grad,
reference[idx],
"value of requires_grad is incorrect",
)
# behavior of max_pool2d in which indices are returned:
input = get_random_test_tensor(size=(1, 3, 8, 8), is_float=True)
input = AutogradCrypTensor(crypten.cryptensor(input))
reference = [True, True, False]
outputs = [None] * 3
outputs[0] = input.max_pool2d(2, return_indices=False)
outputs[1], outputs[2] = input.max_pool2d(2, return_indices=True)
for idx, output in enumerate(outputs):
self.assertEqual(
output.requires_grad,
reference[idx],
"value of requires_grad is incorrect",
)
def test_sequential(self):
"""
Tests crypten.nn.Sequential module.
"""
# try networks of different depth:
for num_layers in range(1, 6):
for wrap in [True, False]:
# construct sequential container:
input_size = (3, 10)
output_size = (input_size[0], input_size[1] - num_layers)
layer_idx = range(input_size[1], output_size[1], -1)
module_list = [
crypten.nn.Linear(num_feat, num_feat - 1)
for num_feat in layer_idx
]
sequential = crypten.nn.Sequential(module_list)
sequential.encrypt()
# check container:
self.assertTrue(sequential.encrypted,
"nn.Sequential not encrypted")
for module in sequential.modules():
self.assertTrue(module.encrypted, "module not encrypted")
assert sum(1 for _ in sequential.modules()) == len(
module_list
), "nn.Sequential contains incorrect number of modules"
# construct test input and run through sequential container:
input = get_random_test_tensor(size=input_size, is_float=True)
encr_input = crypten.cryptensor(input)
if wrap:
encr_input = AutogradCrypTensor(encr_input)
encr_output = sequential(encr_input)
# compute reference output:
encr_reference = encr_input
for module in sequential.modules():
encr_reference = module(encr_reference)
reference = encr_reference.get_plain_text()
# compare output to reference:
self._check(encr_output, reference,
"nn.Sequential forward failed")
def test_cat_stack(self):
for func in ["cat", "stack"]:
for dimensions in range(1, 5):
size = [5] * dimensions
for num_tensors in range(1, 5):
for dim in range(dimensions):
tensors = [
get_random_test_tensor(size=size, is_float=True)
for _ in range(num_tensors)
]
encrypted_tensors = [
AutogradCrypTensor(crypten.cryptensor(t))
for t in tensors
]
for i in range(len(tensors)):
tensors[i].grad = None
tensors[i].requires_grad = True
encrypted_tensors[i].grad = None
encrypted_tensors[i].requires_grad = True
# Forward
reference = getattr(torch, func)(tensors, dim=dim)
encrypted_out = getattr(crypten,
func)(encrypted_tensors,
dim=dim)
self._check(encrypted_out, reference,
f"{func} forward failed")
# Backward
grad_output = get_random_test_tensor(
size=reference.size(), is_float=True)
encrypted_grad_output = crypten.cryptensor(grad_output)
reference.backward(grad_output)
encrypted_out.backward(encrypted_grad_output)
for i in range(len(tensors)):
self._check(
encrypted_tensors[i].grad,
tensors[i].grad,
f"{func} backward failed",
)
def test_gather_scatter(self):
"""Tests gather and scatter gradients"""
sizes = [(2, 2), (3, 5), (3, 5, 10)]
indices = [[0, 1, 0, 0], [0, 1, 0, 0, 1] * 3, [0, 0, 1] * 50]
dims = [0, 1]
funcs = ["scatter", "gather"]
for dim, func in itertools.product(dims, funcs):
for size, index in zip(sizes, indices):
tensor = get_random_test_tensor(size=size, is_float=True)
index = torch.tensor(index).reshape(tensor.shape)
tensor.requires_grad = True
tensor_encr = AutogradCrypTensor(crypten.cryptensor(tensor),
requires_grad=True)
if func == "gather":
reference = getattr(tensor, func)(dim, index)
out_encr = getattr(tensor_encr, func)(dim, index)
else:
src = get_random_test_tensor(size=index.shape,
is_float=True)
reference = getattr(tensor, func)(dim, index, src)
out_encr = getattr(tensor_encr, func)(dim, index, src)
self._check(out_encr, reference,
f"{func} forward failed with index {index}")
grad_out = get_random_test_tensor(size=reference.shape,
is_float=True)
grad_out_encr = crypten.cryptensor(grad_out)
reference.backward(grad_out)
out_encr.backward(grad_out_encr)
self._check(
tensor_encr.grad,
tensor.grad,
f"{func} backward failed with index {index}",
)
def _check_forward_backward(self,
fn_name,
input_tensor,
*args,
msg=None,
**kwargs):
if msg is None:
msg = f"{fn_name} grad_fn incorrect"
for requires_grad in [True]:
# Setup input
input = input_tensor.clone()
input.requires_grad = requires_grad
input_encr = AutogradCrypTensor(crypten.cryptensor(input),
requires_grad=requires_grad)
for private in [False, True]:
input.grad = None
input_encr.grad = None
# Setup args
args_encr = list(args)
for i, arg in enumerate(args):
if private and is_float_tensor(arg):
args_encr[i] = AutogradCrypTensor(
crypten.cryptensor(arg),
requires_grad=requires_grad)
args_encr[i].grad = None # zero grad
if is_float_tensor(arg):
args[i].requires_grad = requires_grad
args[i].grad = None # zero grad
# Check forward pass
if hasattr(input, fn_name):
reference = getattr(input, fn_name)(*args, **kwargs)
elif hasattr(F, fn_name):
reference = getattr(F, fn_name)(input, *args, **kwargs)
elif fn_name == "square":
reference = input.pow(2)
else:
raise ValueError("unknown PyTorch function: %s" % fn_name)
encrypted_out = getattr(input_encr, fn_name)(*args_encr,
**kwargs)
# Remove argmax output from max / min
if isinstance(encrypted_out, (list, tuple)):
reference = reference[0]
encrypted_out = encrypted_out[0]
self._check(encrypted_out, reference, msg + " in forward")
# Check backward pass
grad_output = get_random_test_tensor(max_value=2,
size=reference.size(),
is_float=True)
grad_output_encr = crypten.cryptensor(grad_output)
# Do not check backward if pytorch backward fails
try:
reference.backward(grad_output)
except RuntimeError:
logging.info("skipped")
continue
encrypted_out.backward(grad_output_encr)
self._check(input_encr.grad, input.grad, msg + " in backward")
for i, arg_encr in enumerate(args_encr):
if crypten.is_encrypted_tensor(arg_encr):
self._check(arg_encr.grad, args[i].grad,
msg + " in backward args")
def test_autograd_accumulation(self):
"""Tests accumulation in autograd."""
# define test cases that have nodes with multiple parents:
def test_case1(input, encr_input):
output = input.add(1.0).add(input.exp()).sum()
encr_output = encr_input.add(1.0).add(encr_input.exp()).sum()
return output, encr_output
def test_case2(input, encr_input):
intermediate = input.pow(2.0) # PyTorch
output = intermediate.add(1.0).add(intermediate.mul(2.0)).sum()
encr_intermediate = encr_input.square() # CrypTen
encr_output = (
encr_intermediate.add(1.0).add(encr_intermediate.mul(2.0)).sum()
)
return output, encr_output
def test_case3(input, encr_input):
intermediate1 = input.pow(2.0) # PyTorch
intermediate2 = intermediate1.add(1.0).add(intermediate1.mul(2.0))
output = intermediate2.pow(2.0).sum()
encr_intermediate1 = encr_input.square() # CrypTen
encr_intermediate2 = encr_intermediate1.add(1.0).add(
encr_intermediate1.mul(2.0)
)
encr_output = encr_intermediate2.square().sum()
return output, encr_output
# loop over test cases:
for idx, test_case in enumerate([test_case1, test_case2, test_case2]):
# get input tensors:
input = get_random_test_tensor(size=(12, 5), is_float=True)
input.requires_grad = True
encr_input = AutogradCrypTensor(crypten.cryptensor(input))
# perform multiple forward computations on input that get combined:
output, encr_output = test_case(input, encr_input)
self._check(
encr_output._tensor, output, "forward for test case %d failed" % idx
)
self.assertTrue(
encr_output.requires_grad,
"requires_grad incorrect for test case %d" % idx,
)
# perform backward computation:
output.backward()
encr_output.backward()
self._check(
encr_input.grad, input.grad, "backward for test case %d failed" % idx
)
# test cases in which tensor gets combined with itself:
for func_name in ["sub", "add", "mul"]:
# get input tensors:
input = get_random_test_tensor(size=(12, 5), is_float=True)
input.requires_grad = True
encr_input = AutogradCrypTensor(crypten.cryptensor(input))
# perform forward-backward pass:
output = getattr(input, func_name)(input).sum()
encr_output = getattr(encr_input, func_name)(encr_input).sum()
self._check(encr_output._tensor, output, "forward failed")
self.assertTrue(encr_output.requires_grad, "requires_grad incorrect")
output.backward()
encr_output.backward()
self._check(encr_input.grad, input.grad, "%s backward failed" % func_name)
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_custom_module_training(self):
"""Tests training CrypTen models created directly using the crypten.nn.Module"""
class ExampleNet(crypten.nn.Module):
def __init__(self):
super(ExampleNet, self).__init__()
self.fc1 = crypten.nn.Linear(20, 5)
self.fc2 = crypten.nn.Linear(5, 2)
def forward(self, x):
out = self.fc1(x)
out = self.fc2(out)
return out
model = ExampleNet()
batch_size = 5
x_orig = get_random_test_tensor(size=(batch_size, 20), 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)
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_pytorch_modules(self):
"""
Tests all non-container Modules in crypten.nn that have equivalent
modules in PyTorch.
"""
# input arguments for modules and input sizes:
module_args = {
"AdaptiveAvgPool2d": (2, ),
"AvgPool2d": (2, ),
# "BatchNorm1d": (400,), # FIXME: Unit tests claim gradients are incorrect.
# "BatchNorm2d": (3,),
# "BatchNorm3d": (6,),
"ConstantPad1d": (3, 1.0),
"ConstantPad2d": (2, 2.0),
"ConstantPad3d": (1, 0.0),
"Conv2d": (3, 6, 5),
"Linear": (400, 120),
"MaxPool2d": (2, ),
"ReLU": (),
"Softmax": (0, ),
"LogSoftmax": (0, ),
}
input_sizes = {
"AdaptiveAvgPool2d": (1, 3, 32, 32),
"AvgPool2d": (1, 3, 32, 32),
"BatchNorm1d": (8, 400),
"BatchNorm2d": (8, 3, 32, 32),
"BatchNorm3d": (8, 6, 32, 32, 4),
"ConstantPad1d": (9, ),
"ConstantPad2d": (3, 6),
"ConstantPad3d": (4, 2, 7),
"Conv2d": (1, 3, 32, 32),
"Linear": (1, 400),
"MaxPool2d": (1, 2, 32, 32),
"ReLU": (1, 3, 32, 32),
"Softmax": (5, 5, 5),
"LogSoftmax": (5, 5, 5),
}
# loop over all modules:
for module_name in module_args.keys():
for wrap in [True, False]:
# generate inputs:
input = get_random_test_tensor(size=input_sizes[module_name],
is_float=True)
input.requires_grad = True
encr_input = crypten.cryptensor(input)
if wrap:
encr_input = AutogradCrypTensor(encr_input)
# create PyTorch module:
module = getattr(torch.nn,
module_name)(*module_args[module_name])
module.train()
# create encrypted CrypTen module:
encr_module = crypten.nn.from_pytorch(module, input)
# check that module properly encrypts / decrypts and
# check that encrypting with current mode properly performs no-op
for encrypted in [False, True, True, False, True]:
encr_module.encrypt(mode=encrypted)
if encrypted:
self.assertTrue(encr_module.encrypted,
"module not encrypted")
else:
self.assertFalse(encr_module.encrypted,
"module encrypted")
for key in ["weight", "bias"]:
if hasattr(module, key): # if PyTorch model has key
encr_param = None
# find that key in the crypten.nn.Graph:
if isinstance(encr_module, crypten.nn.Graph):
for encr_node in encr_module.modules():
if hasattr(encr_node, key):
encr_param = getattr(encr_node, key)
break
# or get it from the crypten Module directly:
else:
encr_param = getattr(encr_module, key)
# compare with reference:
# NOTE: Because some parameters are initialized randomly
# with different values on each process, we only want to
# check that they are consistent with source parameter value
reference = getattr(module, key)
src_reference = comm.get().broadcast(reference,
src=0)
msg = "parameter %s in %s incorrect" % (
key, module_name)
if not encrypted:
encr_param = crypten.cryptensor(encr_param)
self._check(encr_param, src_reference, msg)
# compare model outputs:
self.assertTrue(encr_module.training,
"training value incorrect")
reference = module(input)
encr_output = encr_module(encr_input)
self._check(encr_output, reference,
"%s forward failed" % module_name)
# test backward pass:
reference.backward(torch.ones(reference.size()))
encr_output.backward()
if wrap: # you cannot get input gradients on MPCTensor inputs
self._check(
encr_input.grad,
input.grad,
"%s backward on input failed" % module_name,
)
else:
self.assertFalse(hasattr(encr_input, "grad"))
for name, param in module.named_parameters():
encr_param = getattr(encr_module, name)
self._check(
encr_param.grad,
param.grad,
"%s backward on %s failed" % (module_name, name),
)
def test_dropout_module(self):
"""Tests the dropout module"""
input_size = [3, 3, 3]
prob_list = [0.2 * x for x in range(1, 5)]
for module_name in ["Dropout", "Dropout2d", "Dropout3d"]:
for prob in prob_list:
for wrap in [True, False]:
# generate inputs:
input = get_random_test_tensor(size=input_size,
is_float=True,
ex_zero=True)
input.requires_grad = True
encr_input = crypten.cryptensor(input)
if wrap:
encr_input = AutogradCrypTensor(encr_input)
# create PyTorch module:
for inplace in [False, True]:
module = getattr(torch.nn,
module_name)(prob, inplace=inplace)
module.train()
# create encrypted CrypTen module:
encr_module = crypten.nn.from_pytorch(module, input)
# check that module properly encrypts / decrypts and
# check that encrypting with current mode properly
# performs no-op
for encrypted in [False, True, True, False, True]:
encr_module.encrypt(mode=encrypted)
if encrypted:
self.assertTrue(encr_module.encrypted,
"module not encrypted")
else:
self.assertFalse(encr_module.encrypted,
"module encrypted")
# compare model outputs:
# compare the zero and non-zero entries of the encrypted tensor
# with a directly constructed plaintext tensor, since we cannot
# ensure that the randomization produces the same output
# for both encrypted and plaintext tensors
self.assertTrue(encr_module.training,
"training value incorrect")
encr_output = encr_module(encr_input)
plaintext_output = encr_output.get_plain_text()
scaled_tensor = input / (1 - prob)
reference = plaintext_output.where(
plaintext_output == 0, scaled_tensor)
self._check(encr_output, reference,
"Dropout forward failed")
if inplace:
self._check(encr_input, reference,
"In-place dropout failed")
else:
self._check(encr_input, input,
"Out-of-place dropout failed")
# check backward
# compare the zero and non-zero entries of the grad in
# the encrypted tensor with a directly constructed plaintext
# tensor: we do this because we cannot ensure that the
# randomization produces the same output for the input
# encrypted and plaintext tensors and so we cannot ensure
# that the grad in the input tensor is populated identically
all_ones = torch.ones(reference.size())
ref_grad = plaintext_output.where(
plaintext_output == 0, all_ones)
ref_grad_input = ref_grad / (1 - prob)
encr_output.backward()
if (
wrap
): # you cannot get input gradients on MPCTensor inputs
self._check(
encr_input.grad,
ref_grad_input,
"dropout backward on input failed",
)
# check testing mode for Dropout module
encr_module.train(mode=False)
encr_output = encr_module(encr_input)
result = encr_input.eq(encr_output)
result_plaintext = result.get_plain_text().bool()
self.assertTrue(result_plaintext.all(),
"dropout failed in test mode")
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 = crypten.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,
)
