def test_cross_entropy(self):
"""Tests cross_entropy and binary_cross_entropy"""
sizes = [(3, 2), (8, 4), (5, 10)]
losses = ["binary_cross_entropy", "cross_entropy"]
for size, loss in itertools.product(sizes, losses):
batch_size, num_targets = size
if loss == "binary_cross_entropy":
tensor = get_random_test_tensor(
size=(batch_size,), max_value=1.0, is_float=True
)
tensor = tensor.abs().add_(0.001)
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)
self._check(out_encr, reference, f"{loss} forward failed")
# backward
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"{loss} backward failed with")
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_sum(self):
"""Tests sum reduction on encrypted tensor."""
tensor = get_random_test_tensor(size=(5, 100, 100), is_float=True)
encrypted = ArithmeticSharedTensor(tensor)
self._check(encrypted.sum(), tensor.sum(), "sum failed")
for dim in [0, 1, 2]:
reference = tensor.sum(dim)
with self.benchmark(type="sum", dim=dim) as bench:
for _ in bench.iters:
encrypted_out = encrypted.sum(dim)
self._check(encrypted_out, reference, "sum failed")
def test_gather_random(self):
sizes = [(), (1, ), (5, ), (5, 5), (5, 5, 5), (1000, )]
for rank in range(self.world_size):
for size in sizes:
tensor = get_random_test_tensor(size=size)
result = comm.get().gather(tensor, rank)
if rank == self.rank:
self.assertTrue(isinstance(result, list))
for res in result:
self.assertTrue((res == tensor).all())
else:
self.assertIsNone(result[0])
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 compute_gradients 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 compute_gradients:
input.requires_grad = True
target.requires_grad = True
# 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_torch_comparators(self):
"""Test torch comparators on CUDALongTensor"""
for comp in ["gt", "ge", "lt", "le", "eq", "ne"]:
tensor = get_random_test_tensor(is_float=False)
tensor2 = get_random_test_tensor(is_float=False)
t_cuda = CUDALongTensor(tensor)
t2_cuda = CUDALongTensor(tensor2)
reference = getattr(torch, comp)(tensor, tensor2).long()
result1 = getattr(t_cuda, comp)(t2_cuda)
result2 = getattr(torch, comp)(t_cuda, t2_cuda)
self.assertTrue(
type(result1) == CUDALongTensor, "result should be a CUDALongTensor"
)
self.assertTrue(
type(result2) == CUDALongTensor, "result should be a CUDALongTensor"
)
self._check_int(result1.cpu(), reference, "%s comparator failed" % comp)
self._check_int(result2.cpu(), reference, "%s comparator failed" % comp)
def test_torch_broadcast_tensor(self):
"""Test torch.broadcast_tensor on CUDALongTensor"""
x = get_random_test_tensor(size=(1, 5), is_float=False)
y = get_random_test_tensor(size=(5, 1), is_float=False)
x_cuda = CUDALongTensor(x)
y_cuda = CUDALongTensor(y)
a, b = torch.broadcast_tensors(x, y)
a_cuda, b_cuda = torch.broadcast_tensors(x_cuda, y_cuda)
self.assertTrue(
type(a_cuda) == CUDALongTensor,
"result should be a CUDALongTensor")
self.assertTrue(
type(b_cuda) == CUDALongTensor,
"result should be a CUDALongTensor")
self._check_int(a, a_cuda.cpu(),
"torch.broadcast_tensor failed for CUDALongTensor")
self._check_int(b, b_cuda.cpu(),
"torch.broadcast_tensor failed for CUDALongTensor")
def test_sum(self):
"""Tests sum using binary shares"""
tensor = get_random_test_tensor(size=(5, 5, 5), is_float=False)
encrypted = BinarySharedTensor(tensor)
self._check(encrypted.sum(), tensor.sum(), "sum failed")
for dim in [0, 1, 2]:
reference = tensor.sum(dim)
with self.benchmark(type="sum", dim=dim) as bench:
for _ in bench.iters:
encrypted_out = encrypted.sum(dim)
self._check(encrypted_out, reference, "sum failed")
def test_squeeze_unsqueeze(self):
"""Test addition and removal of tensor dimensions"""
for size in SIZES:
tensor = get_random_test_tensor(size=size, is_float=True)
self._check_forward_backward("squeeze", tensor)
for dim in range(tensor.dim()):
self._check_forward_backward("squeeze", tensor, dim)
self._check_forward_backward("unsqueeze", tensor, dim)
# Check unsqueeze on last dimension
self._check_forward_backward("unsqueeze", tensor, tensor.dim())
def test_where(self):
"""Test that crypten.where properly conditions"""
sizes = [(10,), (5, 10), (1, 5, 10)]
y_types = [lambda x: x, crypten.cryptensor]
for size, y_type in itertools.product(sizes, y_types):
tensor1 = get_random_test_tensor(size=size, is_float=True)
encrypted_tensor1 = crypten.cryptensor(tensor1)
tensor2 = get_random_test_tensor(size=size, is_float=True)
encrypted_tensor2 = y_type(tensor2)
condition_tensor = (
get_random_test_tensor(max_value=1, size=size, is_float=False) + 1
)
condition_encrypted = crypten.cryptensor(condition_tensor)
condition_bool = condition_tensor.bool()
reference_out = torch.where(condition_bool, tensor1, tensor2)
encrypted_out = crypten.where(
condition_bool, encrypted_tensor1, encrypted_tensor2
)
y_is_private = crypten.is_encrypted_tensor(tensor2)
self._check(
encrypted_out,
reference_out,
f"{'private' if y_is_private else 'public'} y "
"where failed with public condition",
)
encrypted_out = encrypted_tensor1.where(
condition_encrypted, encrypted_tensor2
)
self._check(
encrypted_out,
reference_out,
f"{'private' if y_is_private else 'public'} y "
"where failed with private condition",
)
def test_where(self):
"""Tests where() conditional element selection"""
sizes = [(10,), (5, 10), (1, 5, 10)]
y_types = [lambda x: x, BinarySharedTensor]
for size, y_type in itertools.product(sizes, y_types):
tensor1 = get_random_test_tensor(size=size, is_float=False)
encrypted_tensor1 = BinarySharedTensor(tensor1)
tensor2 = get_random_test_tensor(size=size, is_float=False)
encrypted_tensor2 = y_type(tensor2)
condition_tensor = (
get_random_test_tensor(max_value=1, size=[1], is_float=False) + 1
)
condition_encrypted = BinarySharedTensor(condition_tensor)
condition_bool = condition_tensor.bool()
reference_out = tensor1 * condition_tensor + tensor2 * (
1 - condition_tensor
)
encrypted_out = encrypted_tensor1.where(condition_bool, encrypted_tensor2)
y_is_private = y_type == BinarySharedTensor
self._check(
encrypted_out,
reference_out,
f"{'private' if y_is_private else 'public'} y "
"where failed with public condition",
)
encrypted_out = encrypted_tensor1.where(
condition_encrypted, encrypted_tensor2
)
self._check(
encrypted_out,
reference_out,
f"{'private' if y_is_private else 'public'} y "
"where failed with private condition",
)
def _conv2d(self, image_size, in_channels):
"""Test convolution of encrypted tensor with public/private tensors."""
nbatches = [1, 3]
kernel_sizes = [(1, 1), (2, 2), (2, 3)]
ochannels = [1, 3]
paddings = [0, 1, (0, 1)]
strides = [1, 2, (1, 2)]
dilations = [1, 2, (1, 2)]
groupings = [1, 2]
for (
batches,
kernel_size,
out_channels,
padding,
stride,
dilation,
groups,
) in itertools.product(
nbatches, kernel_sizes, ochannels, paddings, strides, dilations, groupings
):
# TODO: Fix conv2d gradient in this case:
if in_channels > 1 and groups > 1:
continue
size = (batches, in_channels * groups, *image_size)
image = get_random_test_tensor(size=size, is_float=True)
kernel_size = (out_channels * groups, in_channels, *kernel_size)
kernel = get_random_test_tensor(size=kernel_size, is_float=True)
self._check_forward_backward(
"conv2d",
image,
kernel,
stride=stride,
padding=padding,
dilation=dilation,
groups=groups,
)
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 test_torch_scatter(self):
""" Test scatter/scatter_add function of CUDALongTensor
This test will be skipped for now since torch.scatter provides
inconsistent result given the same input on CUDA. This is likely
due to a potential bug on pytorch's implementation of scatter
"""
funcs = ["scatter", "scatter_add"]
sizes = [(5, 5), (5, 5, 5), (5, 5, 5, 5)]
for func in funcs:
for size in sizes:
for dim in range(len(size)):
tensor1 = get_random_test_tensor(size=size, is_float=False)
tensor2 = get_random_test_tensor(size=size, is_float=False)
index = get_random_test_tensor(size=size, is_float=False)
index = index.abs().clamp(0, 4)
t1_cuda = CUDALongTensor(tensor1)
t2_cuda = CUDALongTensor(tensor2)
idx_cuda = CUDALongTensor(index)
reference = getattr(torch, func)(tensor1, dim, index,
tensor2)
result = getattr(torch, func)(t1_cuda, dim, idx_cuda,
t2_cuda)
result2 = getattr(t1_cuda, func)(dim, idx_cuda, t2_cuda)
self.assertTrue(
type(result) == CUDALongTensor,
"result should be a CUDALongTensor",
)
self.assertTrue(
type(result2) == CUDALongTensor,
"result should be a CUDALongTensor",
)
self._check_int(result.cpu(), reference,
"{} failed".format(func))
self._check_int(result2.cpu(), reference,
"{} failed".format(func))
def test_index_add(self):
"""Test index_add function of encrypted tensor"""
index_add_functions = ["index_add", "index_add_"]
tensor_size1 = [5, 5, 5, 5]
index = torch.tensor([1, 2, 3, 4, 4, 2, 1, 3], dtype=torch.long)
for dimension in range(0, 4):
tensor_size2 = [5, 5, 5, 5]
tensor_size2[dimension] = index.size(0)
for func in index_add_functions:
for tensor_type in [lambda x: x, ArithmeticSharedTensor]:
tensor1 = get_random_test_tensor(size=tensor_size1, is_float=True)
tensor2 = get_random_test_tensor(size=tensor_size2, is_float=True)
encrypted = ArithmeticSharedTensor(tensor1)
encrypted2 = tensor_type(tensor2)
reference = getattr(tensor1, func)(dimension, index, tensor2)
encrypted_out = getattr(encrypted, func)(
dimension, index, encrypted2
)
private = tensor_type == ArithmeticSharedTensor
self._check(
encrypted_out,
reference,
"%s %s failed" % ("private" if private else "public", func),
)
if func.endswith("_"):
# Check in-place index_add worked
self._check(
encrypted,
reference,
"%s %s failed" % ("private" if private else "public", func),
)
else:
# Check original is not modified
self._check(
encrypted,
tensor1,
"%s %s failed" % ("private" if private else "public", func),
)
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 = 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 test_prod(self):
"""Tests prod reduction on encrypted tensor."""
# Increaing size to reduce relative error due to quantization
tensor = get_random_test_tensor(size=(5, 5, 5), is_float=False)
encrypted = ArithmeticSharedTensor(tensor)
self._check(encrypted.prod(), tensor.prod().float(), "prod failed")
for dim in [0, 1, 2]:
reference = tensor.prod(dim).float()
with self.benchmark(type="prod", dim=dim) as bench:
for _ in bench.iters:
encrypted_out = encrypted.prod(dim)
self._check(encrypted_out, reference, "prod failed")
def test_view_reshape(self):
"""Tests view and reshape gradients"""
size_to_views = {
(10, ): [(5, 2), (1, 10)],
(10, 5): [(50), (2, 5, 5)],
(5, 10, 8): [(400), (50, 8), (5, 5, 2, 8)],
}
for size in size_to_views:
for view in size_to_views[size]:
tensor = get_random_test_tensor(size=size, is_float=True)
self._check_forward_backward("view", tensor, view)
self._check_forward_backward("reshape", tensor, view)
def test_torch_arithmetic(self):
"""Test torch arithmetic on CUDALongTensor"""
funcs = ["add", "sub", "mul", "div"]
a = get_random_test_tensor(is_float=False)
b = get_random_test_tensor(min_value=1, is_float=False)
a_cuda = CUDALongTensor(a)
b_cuda = CUDALongTensor(b)
for op in funcs:
reference = getattr(torch, op)(a, b)
result = getattr(torch, op)(a_cuda, b_cuda)
result2 = getattr(a_cuda, op)(b_cuda)
self.assertTrue(type(result), CUDALongTensor)
self._check_int(reference, result.cpu(),
"torch.{} failed for CUDALongTensor".format(op))
self._check_int(
reference,
result2.cpu(),
"torch.{} failed for CUDALongTensor".format(op),
)
def test_torch_gather(self):
"""Test torch.gather on CUDALongTensor"""
sizes = [(5, 5), (5, 5, 5), (5, 5, 5, 5)]
for size in sizes:
for dim in range(len(size)):
tensor = get_random_test_tensor(size=size, is_float=False)
index = get_random_test_tensor(size=size, is_float=False)
index = index.abs().clamp(0, 4)
t_cuda = CUDALongTensor(tensor)
idx_cuda = CUDALongTensor(index)
reference = tensor.gather(dim, index)
result = t_cuda.gather(dim, idx_cuda)
result2 = torch.gather(t_cuda, dim, idx_cuda)
self._check_int(
result.cpu(), reference, f"gather failed with size {size}"
)
self._check_int(
result2.cpu(), reference, f"gather failed with size {size}"
)
def test_scatter(self):
"""Test scatter/scatter_add function of encrypted tensor"""
funcs = ["scatter", "scatter_", "scatter_add", "scatter_add_"]
sizes = [(5, 5), (5, 5, 5), (5, 5, 5, 5)]
for func in funcs:
for size in sizes:
for tensor_type in [lambda x: x, ArithmeticSharedTensor]:
for dim in range(len(size)):
tensor1 = get_random_test_tensor(size=size, is_float=True)
tensor2 = get_random_test_tensor(size=size, is_float=True)
index = get_random_test_tensor(size=size, is_float=False)
index = index.abs().clamp(0, 4)
encrypted = ArithmeticSharedTensor(tensor1)
encrypted2 = tensor_type(tensor2)
reference = getattr(tensor1, func)(dim, index, tensor2)
encrypted_out = getattr(encrypted, func)(dim, index, encrypted2)
private = tensor_type == ArithmeticSharedTensor
self._check(
encrypted_out,
reference,
"%s %s failed" % ("private" if private else "public", func),
)
if func.endswith("_"):
# Check in-place scatter/scatter-add modified input
self._check(
encrypted,
reference,
"%s %s failed to modify input"
% ("private" if private else "public", func),
)
else:
# Check original is not modified
self._check(
encrypted,
tensor1,
"%s %s unintendedly modified input"
% ("private" if private else "public", func),
)
def _conv1d(self, signal_size, in_channels):
"""Test convolution of encrypted tensor with public/private tensors."""
nbatches = [1, 3]
kernel_sizes = [1, 2, 3]
ochannels = [1, 3, 6]
paddings = [0, 1]
strides = [1, 2]
for func_name in ["conv1d", "conv_transpose1d"]:
for kernel_type in [lambda x: x, ArithmeticSharedTensor]:
for (
batches,
kernel_size,
out_channels,
padding,
stride,
) in itertools.product(
nbatches, kernel_sizes, ochannels, paddings, strides
):
input_size = (batches, in_channels, signal_size)
signal = get_random_test_tensor(size=input_size, is_float=True)
if func_name == "conv1d":
k_size = (out_channels, in_channels, kernel_size)
else:
k_size = (in_channels, out_channels, kernel_size)
kernel = get_random_test_tensor(size=k_size, is_float=True)
reference = getattr(F, func_name)(
signal, kernel, padding=padding, stride=stride
)
encrypted_signal = ArithmeticSharedTensor(signal)
encrypted_kernel = kernel_type(kernel)
encrypted_conv = getattr(encrypted_signal, func_name)(
encrypted_kernel, padding=padding, stride=stride
)
self._check(encrypted_conv, reference, f"{func_name} failed")
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_get_set(self):
for tensor_type in [lambda x: x, BinarySharedTensor]:
for size in range(1, 5):
# Test __getitem__
tensor = get_random_test_tensor(size=(size, size), is_float=False)
reference = tensor[:, 0]
encrypted_tensor = BinarySharedTensor(tensor)
encrypted_out = encrypted_tensor[:, 0]
self._check(encrypted_out, reference, "getitem failed")
reference = tensor[0, :]
encrypted_out = encrypted_tensor[0, :]
self._check(encrypted_out, reference, "getitem failed")
# Test __setitem__
tensor2 = get_random_test_tensor(size=(size,), is_float=False)
reference = tensor.clone()
reference[:, 0] = tensor2
encrypted_out = BinarySharedTensor(tensor)
encrypted2 = tensor_type(tensor2)
encrypted_out[:, 0] = encrypted2
self._check(
encrypted_out, reference, "%s setitem failed" % type(encrypted2)
)
reference = tensor.clone()
reference[0, :] = tensor2
encrypted_out = BinarySharedTensor(tensor)
encrypted2 = tensor_type(tensor2)
encrypted_out[0, :] = encrypted2
self._check(
encrypted_out, reference, "%s setitem failed" % type(encrypted2)
)
def _check_forward_backward(
self, func_name, input_tensor, *args, 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)
# check forward pass
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 _conv1d(self, signal_size, in_channels):
"""Test convolution of encrypted tensor with public/private tensors."""
nbatches = [1, 3]
nout_channels = [1, 5]
kernel_sizes = [1, 2, 3]
paddings = [0, 1]
strides = [1, 2]
for batches in nbatches:
size = (batches, in_channels, signal_size)
signal = get_random_test_tensor(size=size, is_float=True)
for kernel_size, out_channels in itertools.product(
kernel_sizes, nout_channels
):
kernel_size = (out_channels, in_channels, kernel_size)
kernel = get_random_test_tensor(size=kernel_size, is_float=True)
for padding in paddings:
for stride in strides:
self._check_forward_backward(
"conv1d", signal, kernel, stride=stride, padding=padding
)
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 = 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 test_patched_matmul(self):
"""Test torch.matmul on CUDALongTensor"""
input_sizes = [
(5,),
(5, 5),
(5,),
(5, 5),
(5, 5, 5),
(5,),
(5, 5, 5, 5),
(5, 5),
]
other_sizes = [
(5,),
(5, 5),
(5, 5),
(5,),
(5,),
(5, 5, 5),
(5, 5),
(5, 5, 5, 5),
]
for x_size, y_size in zip(input_sizes, other_sizes):
x = get_random_test_tensor(size=x_size, max_value=2 ** 62, is_float=False)
x_cuda = CUDALongTensor(x)
y = get_random_test_tensor(size=y_size, max_value=2 ** 62, is_float=False)
y_cuda = CUDALongTensor(y)
z = torch.matmul(x_cuda, y_cuda)
self.assertTrue(
type(z) == CUDALongTensor, "result should be a CUDALongTensor"
)
reference = torch.matmul(x, y)
self._check_int(z.cpu(), reference, "matmul failed for cuda_patches")
def test_save_load(self):
"""Test that crypten.save and crypten.load properly save and load
shares of cryptensors"""
import io
import pickle
def custom_load_function(f):
obj = pickle.load(f)
return obj
def custom_save_function(obj, f):
pickle.dump(obj, f)
all_save_fns = [torch.save, custom_save_function]
all_load_fns = [torch.load, custom_load_function]
tensor = get_random_test_tensor()
cryptensor1 = crypten.cryptensor(tensor)
for i, save_closure in enumerate(all_save_fns):
load_closure = all_load_fns[i]
f = [
io.BytesIO() for i in range(crypten.communicator.get().get_world_size())
]
crypten.save(cryptensor1, f[self.rank], save_closure=save_closure)
f[self.rank].seek(0)
cryptensor2 = crypten.load(f[self.rank], load_closure=load_closure)
# test whether share matches
self.assertTrue(cryptensor1.share.allclose(cryptensor2.share))
# test whether tensor matches
self.assertTrue(
cryptensor1.get_plain_text().allclose(cryptensor2.get_plain_text())
)
attributes = [
a
for a in dir(cryptensor1)
if not a.startswith("__")
and not callable(getattr(cryptensor1, a))
and a not in ["share", "_tensor", "ctx"]
]
for a in attributes:
attr1, attr2 = getattr(cryptensor1, a), getattr(cryptensor2, a)
if a == "encoder":
self.assertTrue(attr1._scale == attr2._scale)
self.assertTrue(attr1._precision_bits == attr2._precision_bits)
elif torch.is_tensor(attr1):
self.assertTrue(attr1.eq(attr2).all())
else:
self.assertTrue(attr1 == attr2)
def _check_max_pool2d_forward_backward(self,
image,
kernel_size,
padding,
stride,
tol=0.1):
"""Checks forward and backward are for max pool 2d.
Verifies gradients by checking sum of non-matching elements to account for
differences in tie resolution in max between PyTorch and CrypTen:
PyTorch returns smallest index for max entries,
whereas CrypTen returns a random index.
Args:
image (torch.tensor): input
kernel_size (tuple of ints): size of the window over which to compute max
padding (int or tuple of ints): implicit zero padding to added on both sides
stride (int or tuple of ints): the stride of the window
"""
# check forward
image = image.clone()
image.requires_grad = True
image_enc = crypten.cryptensor(image, requires_grad=True)
out = torch.nn.functional.max_pool2d(image,
kernel_size,
padding=padding,
stride=stride)
out_enc = image_enc.max_pool2d(kernel_size,
padding=padding,
stride=stride)
self._check(out_enc, out, "max_pool2d forward incorrect")
# check backward
grad_output = get_random_test_tensor(size=out.size(), is_float=True)
grad_output_enc = crypten.cryptensor(grad_output)
out.backward(grad_output)
out_enc.backward(grad_output_enc)
# check sum of non-matching gradient entries
crypten_grad = image_enc.grad.get_plain_text()
non_matching_indices = (image.grad - crypten_grad).abs() > tol
sum_is_close = (crypten_grad[non_matching_indices].sum() -
image.grad[non_matching_indices].sum()) < tol
if not sum_is_close:
msg = "max_pool2d backward failed"
logging.info(msg)
logging.info(f"Result: crypten image gradient {crypten_grad}")
logging.info(f"Result - Reference {image.grad - crypten_grad}")
self.assertTrue(sum_is_close, msg=msg)
