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)
