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_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 autograd in CrypTensor:
input = inputs[0]
input.requires_grad = True
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)(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 = crypten.cryptensor(input, requires_grad=True)
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_batchnorm(self):
"""
Tests batchnorm forward and backward steps with training on / off.
"""
tolerance = 0.1
sizes = [(8, 5), (16, 3), (32, 5), (8, 6, 4), (8, 4, 3, 5)]
for size in sizes:
for is_training in (False, True):
# sample input data, weight, and bias:
tensor = get_random_test_tensor(size=size, is_float=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 over which means and variances are computed:
stats_dimensions = list(range(tensor.dim()))
stats_dimensions.pop(1)
# dummy running mean and variance:
running_mean = tensor.mean(stats_dimensions).detach()
running_var = tensor.var(stats_dimensions).detach()
enc_running_mean = crypten.cryptensor(running_mean)
enc_running_var = crypten.cryptensor(running_var)
# compute reference output:
tensor.requires_grad = True
reference = torch.nn.functional.batch_norm(
tensor,
running_mean,
running_var,
weight=weight,
bias=bias,
training=is_training,
)
# compute CrypTen output:
encrypted_input.requires_grad = True
ctx = AutogradContext()
batch_norm_fn, _ = crypten.gradients.get_grad_fn("batchnorm")
with crypten.no_grad():
encrypted_out = batch_norm_fn.forward(
ctx,
encrypted_input,
weight,
bias,
training=is_training,
running_mean=enc_running_mean,
running_var=enc_running_var,
)
# check forward
self._check(
encrypted_out,
reference,
"batchnorm forward failed with training "
f"{is_training} on {tensor.dim()}-D",
tolerance=tolerance,
)
# check backward (input, weight, and bias gradients):
reference.backward(reference)
with crypten.no_grad():
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_training} on {tensor.dim()}-D",
tolerance=tolerance,
)
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.requires_grad = True
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,
)