def compare(self, debug=False):
rank = comm.get().get_rank()
#Receive secret share from previous function.
with open('dist_rank_{}.pickle'.format(rank), 'rb') as handle:
dist_dict = pickle.load(handle)
dist_enc = dist_dict["distance_share_list_rank{}".format(rank)]
n_dust = len(dist_enc)
dist_enc_new = crypten.cryptensor(torch.ones(n_dust, self.n_point))
for i in range(n_dust):
for j in range(self.n_point):
dist_enc_new[i, j] = dist_enc[i][j]
with open('data_rank_{}.pickle'.format(rank), 'rb') as handle:
ss_dict = pickle.load(handle)
point_enc = ss_dict["point_share_list_rank{}".format(rank)]
point_enc_new = crypten.cryptensor(torch.ones(self.n_point, 2))
for i in range(self.n_point):
point_enc_new[i, :] = point_enc[i][:]
with open('centroid_rank_{}.pickle'.format(rank), 'rb') as handle:
ss_dict = pickle.load(handle)
dust_enc = ss_dict["centroid_share_list_rank{}".format(rank)]
dust_enc_new = crypten.cryptensor(torch.ones(n_dust, 2))
for i in range(n_dust):
dust_enc_new[i, :] = dust_enc[i][:]
#temp is the radius
distance_bool_list = []
changed = False
total_udpate_volume = 0
for i in range(n_dust):
templist = []
temprad = torch.ones(len(dist_enc[i])) * self.radius
#create shared radius ([r,r,r,r....])
radius_enc = crypten.cryptensor(temprad,
ptype=crypten.ptype.arithmetic)
#calculates if point distance is le radius
temp_bool = crypten.cryptensor(temprad,
ptype=crypten.ptype.arithmetic)
temp_bool = dist_enc_new[i] <= radius_enc
#multiply point value with 0/1 matrix,
#result should be the updated centroid location
temp_pts = crypten.cryptensor(torch.ones(point_enc_new.shape))
for j in range(self.n_point):
temp_pts[j, :] = point_enc_new[j, :] * temp_bool[j]
#sum them up, divide by sum of 0/1 matrix
updated_centroid = (temp_pts.sum(0)) / temp_bool.sum()
"""if debug:
a = updated_centroid.get_plain_text()
b = dust_enc_new[i,:].get_plain_text()
if rank==0:
print("a and b:")
print(a)
print(b)"""
#get plain text of ne, if is 1(not equal) then set changed
# udpate_volume = (updated_centroid - dust_enc_new[i,:]).get_plain_text().sum().item()
# print((updated_centroid - dust_enc_new[i,:]).get_plain_text().sum().item())
total_udpate_volume += (
updated_centroid -
dust_enc_new[i, :]).get_plain_text().sum().item()
# changetemp = (updated_centroid!=dust_enc_new[i,:]).get_plain_text().sum().item()
# if (changetemp):
# #if changed, set flag and change dust
# changed = True
# if debug:
# a = updated_centroid.get_plain_text()
# b = dust_enc_new[i,:].get_plain_text()
# if rank==0:
# print("a and b:")
# print(a)
# print(b)
dust_enc[i] = updated_centroid
# print(np.abs(total_udpate_volume))
if np.abs(total_udpate_volume) > 5e-5:
changed = True
#if changed, then update the data file with new dust
if changed:
ss_dict["centroid_share_list_rank{}".format(rank)] = dust_enc
with open('centroid_rank_{}.pickle'.format(rank), 'wb') as handle:
pickle.dump(ss_dict, handle, protocol=pickle.HIGHEST_PROTOCOL)
with open('result.pickle', 'wb') as handle:
pickle.dump(True, handle, protocol=pickle.HIGHEST_PROTOCOL)
else:
with open('result.pickle', 'wb') as handle:
pickle.dump(False, handle, protocol=pickle.HIGHEST_PROTOCOL)
def test_non_pytorch_modules(self):
"""
Tests all non-container Modules in crypten.nn that do not have
equivalent modules in PyTorch.
"""
# input arguments for modules and input sizes:
no_input_modules = ["Constant"]
binary_modules = ["Add", "Sub", "Concat", "MatMul"]
ex_zero_modules = []
module_args = {
"Add": (),
"Concat": (0, ),
"Constant": (1.2, ),
"Exp": (),
"Gather": (0, ),
"MatMul": (),
"Mean": ([0], True),
"Reshape": (),
"Shape": (),
"Sub": (),
"Sum": ([0], True),
"Squeeze": (0, ),
"Transpose": ([1, 3, 0, 2], ),
"Unsqueeze": (0, ),
}
module_lambdas = {
"Add":
lambda x: x[0] + x[1],
"Concat":
lambda x: torch.cat((x[0], x[1])),
"Constant":
lambda _: torch.tensor(module_args["Constant"][0]),
"Exp":
lambda x: torch.exp(x),
"Gather":
lambda x: torch.from_numpy(x[0].numpy().take(
x[1], module_args["Gather"][0])),
"MatMul":
lambda x: torch.matmul(x[0], x[1]),
"Mean":
lambda x: torch.mean(x,
dim=module_args["Mean"][0],
keepdim=(module_args["Mean"][1] == 1)),
"Reshape":
lambda x: x[0].reshape(x[1]),
"Shape":
lambda x: torch.tensor(x.size()).float(),
"Sub":
lambda x: x[0] - x[1],
"Sum":
lambda x: torch.sum(x,
dim=module_args["Sum"][0],
keepdim=(module_args["Sum"][1] == 1)),
"Squeeze":
lambda x: x.squeeze(module_args["Squeeze"][0]),
"Transpose":
lambda x: torch.from_numpy(x.numpy().transpose(module_args[
"Transpose"][0])),
"Unsqueeze":
lambda x: x.unsqueeze(module_args["Unsqueeze"][0]),
}
input_sizes = {
"Add": (10, 12),
"Concat": (2, 2),
"Constant": (1, ),
"Exp": (10, 10, 10),
"Gather": (4, 4, 4, 4),
"MatMul": (4, 4),
"Mean": (3, 3, 3),
"Reshape": (1, 4),
"Shape": (8, 3, 2),
"Sub": (10, 12),
"Sum": (3, 3, 3),
"Squeeze": (1, 12, 6),
"Transpose": (1, 2, 3, 4),
"Unsqueeze": (8, 3),
}
additional_inputs = {
"Gather": torch.tensor([[1, 2], [0, 3]]),
"Reshape": torch.Size([2, 2]),
}
module_attributes = {
# each attribute has two parameters: the name, and a bool indicating
# whether the value should be wrapped into a list when the module is created
"Add": [],
"Exp": [],
"Concat": [("axis", False)],
"Constant": [("value", False)],
"Gather": [("axis", False)],
"MatMul": [],
"Mean": [("axes", False), ("keepdims", False)],
"Reshape": [],
"Shape": [],
"Sub": [],
"Sum": [("axes", False), ("keepdims", False)],
"Squeeze": [("axes", True)],
"Transpose": [("perm", False)],
"Unsqueeze": [("axes", True)],
}
# loop over all modules:
for module_name in module_args.keys():
# create encrypted CrypTen module:
encr_module = getattr(crypten.nn,
module_name)(*module_args[module_name])
encr_module.encrypt()
self.assertTrue(encr_module.encrypted, "module not encrypted")
# generate inputs:
inputs, encr_inputs = None, None
ex_zero_values = module_name in ex_zero_modules
if module_name in binary_modules:
inputs = [
get_random_test_tensor(
size=input_sizes[module_name],
is_float=True,
ex_zero=ex_zero_values,
) for _ in range(2)
]
encr_inputs = [crypten.cryptensor(input) for input in inputs]
elif module_name not in no_input_modules:
inputs = get_random_test_tensor(size=input_sizes[module_name],
is_float=True,
ex_zero=ex_zero_values)
encr_inputs = crypten.cryptensor(inputs)
# some modules take additonal indices as input:
if module_name in additional_inputs:
if not isinstance(inputs, (list, tuple)):
inputs, encr_inputs = [inputs], [encr_inputs]
inputs.append(additional_inputs[module_name])
# encrypt only torch tensor inputs, not shapes
if torch.is_tensor(inputs[-1]):
encr_inputs.append(crypten.cryptensor(inputs[-1]))
else:
encr_inputs.append(inputs[-1])
# compare model outputs:
reference = module_lambdas[module_name](inputs)
encr_output = encr_module(encr_inputs)
self._check(encr_output, reference, "%s failed" % module_name)
# create attributes for static from_onnx function
local_attr = {}
for i, attr_tuple in enumerate(module_attributes[module_name]):
attr_name, wrap_attr_in_list = attr_tuple
if wrap_attr_in_list:
local_attr[attr_name] = [module_args[module_name][i]]
else:
local_attr[attr_name] = module_args[module_name][i]
# Update ReduceSum/ReduceMean module attributes, since the module and
# from_onnx path are different
if module_name == "ReduceSum":
local_attr["keepdims"] = 1 if module_args["ReduceSum"][
1] is True else 0
if module_name == "ReduceMean":
local_attr["keepdims"] = (
1 if module_args["ReduceMean"][1] is True else 0)
# compare model outputs using the from_onnx static function
module = getattr(crypten.nn,
module_name).from_onnx(attributes=local_attr)
encr_module_onnx = module.encrypt()
encr_output = encr_module_onnx(encr_inputs)
self._check(encr_output, reference, "%s failed" % module_name)
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 run_experiment(
model_name,
imagenet_folder=None,
tensorboard_folder="/tmp",
num_samples=None,
context_manager=None,
):
"""Runs inference using specified vision model on specified dataset."""
crypten.init()
# check inputs:
assert hasattr(models,
model_name), ("torchvision does not provide %s model" %
model_name)
if imagenet_folder is None:
imagenet_folder = tempfile.gettempdir()
download = True
else:
download = False
if context_manager is None:
context_manager = NoopContextManager()
# load dataset and model:
with context_manager:
model = getattr(models, model_name)(pretrained=True)
model.eval()
dataset = datasets.ImageNet(imagenet_folder,
split="val",
download=download)
# define appropriate transforms:
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406],
std=[0.229, 0.224, 0.225]),
])
to_tensor_transform = transforms.ToTensor()
# encrypt model:
dummy_input = to_tensor_transform(dataset[0][0])
dummy_input.unsqueeze_(0)
encrypted_model = crypten.nn.from_pytorch(model, dummy_input=dummy_input)
encrypted_model.encrypt()
# show encrypted model in tensorboard:
if SummaryWriter is not None:
writer = SummaryWriter(log_dir=tensorboard_folder)
writer.add_graph(encrypted_model)
writer.close()
# loop over dataset:
meter = AccuracyMeter()
for idx, sample in enumerate(dataset):
# preprocess sample:
image, target = sample
image = transform(image)
image.unsqueeze_(0)
target = torch.tensor([target], dtype=torch.long)
# perform inference using encrypted model on encrypted sample:
encrypted_image = crypten.cryptensor(image)
encrypted_output = encrypted_model(encrypted_image)
# measure accuracy of prediction
output = encrypted_output.get_plain_text()
meter.add(output, target)
# progress:
logging.info("[sample %d of %d] Accuracy: %f" %
(idx + 1, len(dataset), meter.value()[1]))
if num_samples is not None and idx == num_samples - 1:
break
# print final accuracy:
logging.info("Accuracy on all %d samples: %f" %
(len(dataset), meter.value()[1]))
def party(): # pragma: no cover
t = crypten.cryptensor(expected)
return t.get_plain_text()
def upload_en1(g_pe, list_up): # 第一轮以后，上传到字典里的是h的特征
for i in list_up:
a = g_pe.nodes[g_pe.map_to_subgraph_nid(i)].data['h'] # 对应到子图上的节点
a = a[0]
b = crypten.cryptensor(a)
common_nodes[i] = b # 为对应的点 添加值
def test_wrap_error_detection(self):
"""Force a wrap error and test whether it raises in debug mode."""
encrypted_tensor = crypten.cryptensor(0)
encrypted_tensor.share = tensor(2**63 - 1)
with self.assertRaises(ValueError):
encrypted_tensor.div(2)
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_model_mpc():
mem_before = get_process_memory()
runtime = 0
pid = comm.get().get_rank()
ws = comm.get().world_size
name = participants[pid]
if pid == 0:
print(f"Hello from the main process (rank#{pid} of {ws})!")
print(f"My name is {name}.")
print(f"My colleagues today are: ")
print(participants)
results = {
"total": 0,
"per_iter": [],
"inference": {
"total": 0,
"per_batch": [],
"per_image": [],
"average_per_image": 0
},
"mem_before": mem_before,
"mem_after": None
}
predictions = []
targets = []
class_correct = [0] * NUM_CLASSES
class_total = [0] * NUM_CLASSES
# Setup log files per process
postfix = f"{DATASET_NAME}_{ws}p_{pid}.log"
memory_log = memory_dir / postfix
runtimes_log = runtimes_dir / postfix
results_log = results_dir / postfix
#convert_legacy_config() # LEGACY
#model_mpc = crypten.nn.from_pytorch(model, dummy_image)
# Instantiate and load the model
model = Net()
# Load model
dummy_image = torch.empty([1, NUM_CHANNELS, IMG_WIDTH,
IMG_HEIGHT]) # is that the right way around? :D
#model = crypten.load(model_file_name, dummy_model=model)
model.load_state_dict(torch.load(model_file_name))
#model = crypten.load(model_file_name, dummy_model=model, src=0)
model_mpc = crypten.nn.from_pytorch(model, dummy_image)
model_mpc.encrypt(src=0)
if pid == 0:
print("Gonna evaluate now...")
test_loss = 0.0
model_mpc.eval() # prep model for evaluation
before_test.wait()
start = time()
iters = 0
for data, target in tqdm(test_loader, position=0): #, desc=f"{name}"):
start_iter = time()
data_enc = []
if ws > 2:
for idx, batch in enumerate(
split_data_even(data, ws - 1, data.shape[0])):
data_enc.append(crypten.cryptensor(batch, src=idx + 1))
#data_enc = crypten.cat(data_enc, dim=0)
else:
data_enc.append(crypten.cryptensor(data, src=1))
target_enc = crypten.cryptensor(target, src=0)
# forward pass: compute predicted outputs by passing inputs to the model
output = []
start_batch_inference = time()
# In each batch, each participant except the model holder has an equal share of the batch
# Iterate over each participants share
for dat in data_enc:
output.append(model_mpc(dat))
stop_batch_inference = time()
output = crypten.cat(output, dim=0)
# convert output probabilities to predicted class
pred = output.argmax(dim=1, one_hot=False)
# calculate the loss
if pid == 0:
if pred.shape != target_enc.shape:
print((pred.shape, target_enc.shape))
loss = criterion(pred, target_enc).get_plain_text()
# update test loss
test_loss += loss.item() * data.size(0)
### compare predictions to true label
# decrypt predictions
pred = pred.get_plain_text()
correct = np.squeeze(pred.eq(target.data.view_as(pred)))
# calculate test accuracy for each object class
predictions.append(pred)
targets.append(target)
for i in range(len(target)):
label = target.data[i]
class_correct[label] += correct[i].item()
class_total[label] += 1
results["per_iter"].append(time() - start_iter)
results["inference"]["per_batch"].append(stop_batch_inference - start_batch_inference)
iters += 1
iter_sync.wait()
log_memory(memory_log)
stop = time()
runtime = stop - start
results["total"] = runtime
results["average_per_iter"] = np.mean(results["per_iter"])
results["inference"]["total"] = np.sum(results["inference"]["per_batch"])
results["inference"]["per_image"] = [x/batch_size for x in results["inference"]["per_batch"]]
results["inference"]["average_per_image"] = np.mean(results["inference"]["per_image"])
# results = {
# "total": 0,
# "per_iter": [],
# "inference": {
# "total": 0,
# "per_batch": [],
# "per_image": [],
# "average_per_image": 0
# }
# }
if pid == 0:
print("Done evaluating...")
after_test.wait()
if pid == 0:
print("Ouputing information...")
# calculate and print avg test loss
test_loss = test_loss / len(test_loader.sampler)
# if pid == 0:
# print(f"Test runtime: {runtime:5.2f}s\n\n")
# print(f"Test Loss: {test_loss:.6}\n")
# # Print accuracy per class
# for i in range(NUM_CLASSES):
# if class_total[i] > 0:
# print(
# f"Test Accuracy of {i:5}: "
# f"{100 * class_correct[i] / class_total[i]:3.0f}% "
# f"({np.sum(class_correct[i]):4} / {np.sum(class_total[i]):4} )"
# )
# else:
# print(
# f"Test Accuracy of {classes[i]}: N/A (no training examples)"
# )
# # Print overall accuracy
# print(
# f"\nTest Accuracy (Overall): {100. * np.sum(class_correct) / np.sum(class_total):3.0f}% "
# f"( {np.sum(class_correct)} / {np.sum(class_total)} )")
# Gather log
LOG_STR = f"Rank: {pid}\nWorld_Size: {ws}\n\n"
LOG_STR += f"Test runtime: {runtime:5.2f}s\n"
LOG_STR += f"Test Loss: {test_loss:.6}\n"
LOG_STR += "\n"
for i in range(NUM_CLASSES):
if class_total[i] > 0:
LOG_STR += f"Test Accuracy of {i:5}: " \
f"{100 * class_correct[i] / class_total[i]:3.0f}% " \
f"({np.sum(class_correct[i]):4} / {np.sum(class_total[i]):4} )"
LOG_STR += "\n"
else:
LOG_STR += f"Test Accuracy of {classes[i]}: N/A (no training examples)"
LOG_STR += "\n"
LOG_STR += f"\nTest Accuracy (Overall): {100. * np.sum(class_correct) / np.sum(class_total):3.0f}% " + \
f"( {np.sum(class_correct)} / {np.sum(class_total)} )"
if pid == 0:
print(LOG_STR)
with open(log_dir / f"stdout_{pid}", "w") as f:
f.write(LOG_STR)
done.wait()
mem_after = get_process_memory()
results["mem_after"] = mem_after
with open(results_log, 'w') as f:
f.write(str(results))
if pid == 0:
with open(results_dir / f'latest_{pid}.txt', 'w') as f:
f.write(str(results))
return results
def _check_max_pool2d_forward_backward(self,
image,
kernel_size,
padding,
stride,
dilation,
ceil_mode,
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
ceil_mode (bool): determines whether output size is rounded down or up
"""
# 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,
dilation=dilation,
ceil_mode=ceil_mode,
)
out_enc = image_enc.max_pool2d(
kernel_size,
padding=padding,
stride=stride,
dilation=dilation,
ceil_mode=ceil_mode,
)
if out.isinf().any():
# PyTorch can produce improperly sized outputs with Inf values using ceil_mode in some cases
if ceil_mode:
return
self.assertTrue(out.size() == out_enc.size(),
"max_pool2d forward incorrect")
return # backward will break if output is -inf
else:
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)
def bernoulli(*size):
x = crypten.cryptensor(p * torch.ones(*size))
return x.bernoulli()
def f_return_cryptensor(*args, **kwargs):
return crypten.cryptensor(th.zeros([]))
print(epoch)
t0 = time.time()
net_A1.train()
net_B1.train()
net_A2.train()
net_B2.train()
net_A3.train()
net_B3.train()
g1.update_all(gcn_msg, gcn_reduce) # A聚合,针对 h
logits_A1 = net_A1(g1, g1.ndata['h']) # A的第一层,线性变换
for parameters in net_A1.parameters(): # 线性变换的参数
par1 = parameters
par1 = crypten.cryptensor(par1) # 参数加密上传
download_en(g1, dic1, par1) # 下载, 针对 h，是将h加上服务器上的加密值乘以参数再解密
logits_A1 = net_A2(g1, logits_A1) # Relu
g2.update_all(gcn_msg, gcn_reduce) # B聚合,针对 h
logits_B1 = net_B1(g2, g2.ndata['h']) # B的第一层
for parameters in net_B1.parameters(): # 线性变换的参数
par2 = parameters
par2 = crypten.cryptensor(par2) # 参数加密上传
download_en(g2, dic2, par2) # 下载, 针对 h，是将h加上服务器上的d
logits_B1 = net_B2(g2, logits_B1) # relu
###
upload_en1(g1, list_up1) # 第一个图同步上传，传到特征h
upload_en1(g2, list_up2) # 第二个图同步上传
# print('第一层后%s' % common_nodes[0].size())
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 _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_non_pytorch_modules(self):
"""
Tests all non-container Modules in crypten.nn that do not have
equivalent modules in PyTorch.
"""
# input arguments for modules and input sizes:
no_input_modules = ["Constant"]
binary_modules = ["Add", "Sub", "Concat"]
module_args = {
"Add": (),
"Concat": (0, ),
"Constant": (1.2, ),
"Gather": (0, ),
"Reshape": (),
"Shape": (),
"Sub": (),
"Squeeze": (0, ),
"Unsqueeze": (0, ),
}
module_lambdas = {
"Add":
lambda x: x[0] + x[1],
"Concat":
lambda x: torch.cat((x[0], x[1])),
"Constant":
lambda _: torch.tensor(module_args["Constant"][0]),
"Gather":
lambda x: torch.from_numpy(x[0].numpy().take(
x[1], module_args["Gather"][0])),
"Reshape":
lambda x: x[0].reshape(x[1].tolist()),
"Shape":
lambda x: torch.tensor(x.size()).float(),
"Sub":
lambda x: x[0] - x[1],
"Squeeze":
lambda x: x.squeeze(module_args["Squeeze"][0]),
"Unsqueeze":
lambda x: x.unsqueeze(module_args["Unsqueeze"][0]),
}
input_sizes = {
"Add": (10, 12),
"Concat": (2, 2),
"Constant": (1, ),
"Gather": (4, 4, 4, 4),
"Reshape": (1, 4),
"Shape": (8, 3, 2),
"Sub": (10, 12),
"Squeeze": (1, 12, 6),
"Unsqueeze": (8, 3),
}
additional_inputs = {
"Gather": torch.tensor([[1, 2], [0, 3]]),
"Reshape": torch.tensor([2, 2]),
}
module_attributes = {
"Add": [],
"Concat": [("axis", int)],
"Constant": [("value", int)],
"Gather": [("axis", int)],
"Reshape": [],
"Shape": [],
"Sub": [],
"Squeeze": [("axes", list)],
"Unsqueeze": [("axes", list)],
}
# loop over all modules:
for module_name in module_args.keys():
# create encrypted CrypTen module:
encr_module = getattr(crypten.nn,
module_name)(*module_args[module_name])
encr_module.encrypt()
self.assertTrue(encr_module.encrypted, "module not encrypted")
# generate inputs:
inputs, encr_inputs = None, None
if module_name in binary_modules:
inputs = [
get_random_test_tensor(size=input_sizes[module_name],
is_float=True) for _ in range(2)
]
encr_inputs = [crypten.cryptensor(input) for input in inputs]
elif module_name not in no_input_modules:
inputs = get_random_test_tensor(size=input_sizes[module_name],
is_float=True)
encr_inputs = crypten.cryptensor(inputs)
# some modules take additonal indices as input:
if module_name in additional_inputs:
if not isinstance(inputs, (list, tuple)):
inputs, encr_inputs = [inputs], [encr_inputs]
inputs.append(additional_inputs[module_name])
encr_inputs.append(crypten.cryptensor(inputs[-1]))
# compare model outputs:
reference = module_lambdas[module_name](inputs)
encr_output = encr_module(encr_inputs)
self._check(encr_output, reference, "%s failed" % module_name)
# create attributes for static from_onnx function
local_attr = {}
for i, attr_tuple in enumerate(module_attributes[module_name]):
attr_name, attr_type = attr_tuple
if attr_type == list:
local_attr[attr_name] = [module_args[module_name][i]]
else:
local_attr[attr_name] = module_args[module_name][i]
# compare model outputs using the from_onnx static function
module = getattr(crypten.nn,
module_name).from_onnx(attributes=local_attr)
encr_module_onnx = module.encrypt()
encr_output = encr_module_onnx(encr_inputs)
self._check(encr_output, reference, "%s failed" % module_name)
net_A1.train()
net_B1.train()
net_C1.train()
net_A2.train()
net_B2.train()
net_C2.train()
net_A3.train()
net_B3.train()
net_C3.train()
g1.update_all(gcn_msg, gcn_reduce) # A聚合,针对 h
logits_A1 = net_A1(g1, g1.ndata['h']) # A的第一层,线性变换
for parameters in net_A1.parameters(): # 线性变换的参数
par1 = parameters
par1 = crypten.cryptensor(par1) # 参数加密上传
download_en(g1, dic1, par1) # 下载, 针对 h，是将h加上服务器上的加密值乘以参数再解密
logits_A1 = net_A2(g1, logits_A1) # Relu
g2.update_all(gcn_msg, gcn_reduce) # B聚合,针对 h
logits_B1 = net_B1(g2, g2.ndata['h']) # B的第一层
for parameters in net_B1.parameters(): # 线性变换的参数
par2 = parameters
par2 = crypten.cryptensor(par2) # 参数加密上传
download_en(g2, dic2, par2) # 下载, 针对 h，是将h加上服务器上的d
logits_B1 = net_B2(g2, logits_B1) # relu
g3.update_all(gcn_msg, gcn_reduce) # C
logits_C1 = net_C1(g3, g3.ndata['h'])
for parameters in net_C1.parameters():
par3 = parameters
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 online_learner(
sampler,
backend="mpc",
nr_iters=7,
score_func=None,
monitor_func=None,
checkpoint_func=None,
checkpoint_every=0,
):
"""
Online learner that minimizes linear least squared loss.
Args:
sampler: An iterator that returns one sample at a time. Samples are
assumed to be `dict`s with a `'context'` and a `'rewards'` field.
backend: Which privacy protocol to use (default 'mpc').
score_func: A closure that can be used to plug in exploration mechanisms.
monitor_func: A closure that does logging.
checkpoint_func: A closure that does checkpointing.
nr_iters: Number of Newton-Rhapson iterations to use for private
reciprocal.
"""
# initialize some variables:
total_reward = 0.0
# initialize constructor for tensors:
crypten.set_default_backend(backend)
# loop over dataset:
idx = 0
for sample in sampler():
start_t = time.time()
# unpack sample:
assert "context" in sample and "rewards" in sample, (
"invalid sample: %s" % sample)
context = crypten.cryptensor(sample["context"])
num_features = context.nelement()
num_arms = sample["rewards"].nelement()
# initialization of model parameters:
if idx == 0:
# initialize accumulators for linear least squares:
A_inv = [
torch.eye(num_features).unsqueeze(0) for _ in range(num_arms)
]
A_inv = crypten.cat([crypten.cryptensor(A) for A in A_inv])
b = crypten.cryptensor(torch.zeros(num_arms, num_features))
# compute initial weights for all arms:
weights = b.unsqueeze(1).matmul(A_inv).squeeze(1)
# compute score of all arms:
scores = weights.matmul(context)
# plug in exploration mechanism:
if score_func is not None:
score_func(scores, A_inv, b, context)
onehot = scores.argmax()
# In practice only one party opens the onehot vector in order to
# take the action.
selected_arm = onehot.get_plain_text().argmax()
# Once the action is taken, the reward (a scalar) is observed by some
# party and secret shared. Here we simulate that by selecting the
# reward from the rewards vector and then sharing it.
reward = crypten.cryptensor((sample["rewards"][selected_arm] >
random.random()).view(1).float())
# update linear least squares accumulators (using Sherman–Morrison
# formula):
A_inv_context = A_inv.matmul(context)
numerator = A_inv_context.unsqueeze(1).mul(A_inv_context.unsqueeze(2))
denominator = A_inv_context.matmul(context).add(1.0).view(-1, 1, 1)
with crypten.mpc.ConfigManager("reciprocal_nr_iters", nr_iters):
update = numerator.mul_(denominator.reciprocal())
A_inv.sub_(update.mul_(onehot.view(-1, 1, 1)))
b.add_(context.mul(reward).unsqueeze(0).mul_(onehot.unsqueeze(0)))
# update model weights:
weights = b.unsqueeze(1).matmul(A_inv).squeeze(1)
# monitor learning progress: we use the plain reward only for
# monitoring
reward = reward.get_plain_text().item()
total_reward += reward
iter_time = time.time() - start_t
if monitor_func is not None:
monitor_func(idx, reward, total_reward, iter_time)
idx += 1
# checkpointing:
if checkpoint_func is not None and idx % checkpoint_every == 0:
checkpoint_func(
idx,
{
"A_inv": [AA.get_plain_text() for AA in A_inv],
"b": [bb.get_plain_text() for bb in b],
},
)
# signal monitoring closure that we are done:
if monitor_func is not None:
monitor_func(idx, None, None, None, finished=True)
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)
def forward(self, x):
size = torch.tensor(x.size())
if crypten.is_encrypted_tensor(x):
size = crypten.cryptensor(size.float())
return size
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(self):
resnet_model = resnet32(1)
transform = transforms.Compose([
lambda x: Image.open(x).convert('RGB'),
transforms.Resize((32, 32)),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))
])
if self.rank==0: # ALICE
model = crypten.load('./SCML_Project/training/models/covid_resnet32/checkpoint_cpu_cpu.pt', dummy_model=resnet_model, src=0)
dummy_input = torch.empty((1, 3, 32, 32))
private_model = crypten.nn.from_pytorch(model.double(), dummy_input.double())
private_model.encrypt(src=0)
data_enc = crypten.cryptensor(dummy_input, src=1)
private_model.eval()
start = time.time()
output_enc = private_model(data_enc)
comm.get().send_obj(output_enc.share, 1)
end = time.time()
elif self.rank==1: # BOB
model = crypten.load('./SCML_Project/training/models/covid_resnet32/resnet_random.pt', dummy_model=resnet_model, src=0)
dummy_input = torch.empty((1, 3, 32, 32))
private_model = crypten.nn.from_pytorch(model.double(), dummy_input.double())
private_model.encrypt(src=0)
data_enc = crypten.cryptensor(transform('./SCML_Project/training/dataset/COVID/COVID-1.png').unsqueeze(0), src=1)
private_model.eval()
start = time.time()
output_enc = private_model(data_enc)
done0 = comm.get().recv_obj(0)
done2 = comm.get().recv_obj(2)
done3 = comm.get().recv_obj(3)
print("Bob received: ", output_enc.share+done0+done2+done3)
end = time.time()
else:
model = crypten.load('./SCML_Project/training/models/covid_resnet32/resnet_random.pt', dummy_model=resnet_model, src=0)
dummy_input = torch.empty((1, 3, 32, 32))
private_model = crypten.nn.from_pytorch(model.double(), dummy_input.double())
private_model.encrypt(src=0)
data_enc = crypten.cryptensor(dummy_input, src=1)
private_model.eval()
start = time.time()
output_enc = private_model(data_enc)
comm.get().send_obj(output_enc.share, 1)
end = time.time()
print('Time: ', end-start)
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 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 compute_gradients 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)
encr_input.requires_grad = compute_gradients
# create PyTorch module:
module = getattr(torch.nn, module_name)(prob)
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")
# 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.sum().backward()
if compute_gradients:
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_tensorflow_model_conversion(self):
import tensorflow as tf
import tf2onnx
# create simple model
model_tf1 = tf.keras.Sequential([
tf.keras.layers.Dense(
10,
activation=tf.nn.relu,
kernel_initializer="ones",
bias_initializer="ones",
input_shape=(4, ),
),
tf.keras.layers.Dense(
10,
activation=tf.nn.relu,
kernel_initializer="ones",
bias_initializer="ones",
),
tf.keras.layers.Dense(3, kernel_initializer="ones"),
])
model_tf2 = tf.keras.Sequential([
tf.keras.layers.Conv2D(
32,
3,
activation="relu",
strides=1,
kernel_initializer="ones",
bias_initializer="ones",
input_shape=(32, 32, 3),
),
tf.keras.layers.MaxPooling2D(3),
tf.keras.layers.GlobalAveragePooling2D(),
tf.keras.layers.Dropout(0.5),
])
model_tf3 = tf.keras.Sequential([
tf.keras.layers.Conv1D(
32,
1,
activation="relu",
strides=1,
kernel_initializer="ones",
bias_initializer="ones",
input_shape=(6, 128),
),
tf.keras.layers.AvgPool1D(1),
])
feature_sizes = [(1, 4), (1, 32, 32, 3), (1, 6, 128)]
label_sizes = [(1, 3), (1, 32), (1, 6, 32)]
for i, curr_model_tf in enumerate([model_tf1, model_tf2, model_tf3]):
# create a random feature vector
features = get_random_test_tensor(size=feature_sizes[i],
is_float=True,
min_value=1,
max_value=3)
labels = get_random_test_tensor(size=label_sizes[i],
is_float=True,
min_value=1)
# convert to a TF tensor via numpy
features_tf = tf.convert_to_tensor(features.numpy())
labels_tf = tf.convert_to_tensor(labels.numpy())
# compute the tensorflow predictions
curr_model_tf.compile("sgd",
loss=tf.keras.losses.MeanSquaredError())
curr_model_tf.fit(features_tf, labels_tf)
result_tf = curr_model_tf(features_tf, training=False)
# convert TF model to CrypTen model
# write as a SavedModel, then load GraphDef from it
import tempfile
saved_model_dir = tempfile.NamedTemporaryFile(delete=True).name
os.makedirs(saved_model_dir, exist_ok=True)
curr_model_tf.save(saved_model_dir)
graph_def, inputs, outputs = tf2onnx.tf_loader.from_saved_model(
saved_model_dir, None, None)
model_enc = crypten.nn.from_tensorflow(graph_def,
list(inputs.keys()),
list(outputs.keys()))
# encrypt model and run it
model_enc.encrypt()
features_enc = crypten.cryptensor(features)
result_enc = model_enc(features_enc)
# compare the results
result = torch.tensor(result_tf.numpy())
self._check(result_enc, result, "nn.from_tensorflow failed")
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),
"Conv1d": (3, 6, 5),
"Conv2d": (3, 6, 5),
"Linear": (400, 120),
"MaxPool2d": (2, ),
"ReLU": (),
"Sigmoid": (),
"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),
"Conv1d": (1, 3, 32),
"Conv2d": (1, 3, 32, 32),
"Linear": (1, 400),
"MaxPool2d": (1, 2, 32, 32),
"ReLU": (1, 3, 32, 32),
"Sigmoid": (8, 3, 32, 32),
"Softmax": (5, 5, 5),
"LogSoftmax": (5, 5, 5),
}
# loop over all modules:
for module_name in module_args.keys():
for compute_gradients 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)
encr_input.requires_grad = compute_gradients
# 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.sum().backward()
encr_output.sum().backward()
if compute_gradients:
self._check(
encr_input.grad,
input.grad,
"%s backward on input failed" % module_name,
)
else:
self.assertIsNone(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_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,
)
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_autograd_accumulation(self):
"""Tests accumulation in autograd."""
# graphs that have nodes with multiple parents, dead leafs, etc.:
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
def test_case4(input, encr_input):
intermediate1 = input.mul(3.0).add(2.0).pow(2.0) # PyTorch
intermediate2 = intermediate1.add(1.0).add(intermediate1.mul(2.0))
output = intermediate2.pow(2.0).sum()
encr_intermediate1 = encr_input.mul(3.0).add(
2.0).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
def test_case5(input, encr_input):
intermediate1 = input.mul(3.0) # PyTorch
intermediate2 = input.add(2.0).pow(2.0)
intermediate3 = input.pow(2.0)
output = (torch.cat([intermediate1, intermediate2,
intermediate3]).mul(0.5).sum())
encr_intermediate1 = encr_input.mul(3.0) # CrypTen
encr_intermediate2 = encr_input.add(2.0).square()
encr_intermediate3 = encr_input.pow(2.0)
encr_output = (crypten.cat(
[encr_intermediate1, encr_intermediate2,
encr_intermediate3]).mul(0.5).sum())
return output, encr_output
def test_case6(input, encr_input):
idx1 = torch.tensor([[0, 2, 4, 3, 8]], dtype=torch.long)
idx2 = torch.tensor([[5, 1, 3, 5, 2]], dtype=torch.long)
idx3 = torch.tensor([[2, 3, 1]], dtype=torch.long)
intermediate1 = input.gather(0,
idx1).gather(1,
idx3).pow(2.0) # PyTorch
intermediate2 = input.gather(0, idx2).gather(1, idx3).add(-2.0)
output = torch.cat([intermediate1, intermediate2]).mul(0.5).sum()
encr_intermediate1 = (encr_input.gather(0, idx1).gather(
1, idx3).square()) # CrypTen
encr_intermediate2 = encr_input.gather(0,
idx2).gather(1,
idx3).add(-2.0)
encr_output = (crypten.cat(
[encr_intermediate1, encr_intermediate2],
dim=0).mul(0.5).sum())
return output, encr_output
def test_case7(input, encr_input):
intermediate1 = input.add(3.0) # PyTorch
intermediate2 = input.add(2.0).pow(2.0)
intermediate3 = intermediate1.add(intermediate2)
intermediate4 = intermediate1.add(intermediate2)
output = intermediate3.add(intermediate4).sum()
encr_intermediate1 = encr_input.add(3.0) # CrypTen
encr_intermediate2 = encr_input.add(2.0).pow(2.0)
encr_intermediate3 = encr_intermediate1.add(encr_intermediate2)
encr_intermediate4 = encr_intermediate1.add(encr_intermediate2)
encr_output = encr_intermediate3.add(encr_intermediate4).sum()
return output, encr_output
def test_case8(input, encr_input):
intermediate1 = input.add(3.0)
intermediate2 = torch.cat([input, intermediate1])
intermediate3 = intermediate2.pow(2.0)
output = torch.cat([input, intermediate2,
intermediate3]).add(-1).sum()
encr_intermediate1 = encr_input.add(3.0)
encr_intermediate2 = crypten.cat([encr_input, encr_intermediate1])
encr_intermediate3 = encr_intermediate2.pow(2.0)
encr_output = (crypten.cat(
[encr_input, encr_intermediate2,
encr_intermediate3]).add(-1).sum())
return output, encr_output
def test_case9(input, encr_input):
intermediate1 = torch.cat([input, input])
intermediate2 = intermediate1.mean(0, keepdim=True)
output = torch.cat([intermediate2, intermediate1], dim=0).sum()
encr_intermediate1 = crypten.cat([encr_input, encr_input])
encr_intermediate2 = encr_intermediate1.mean(0, keepdim=True)
encr_output = crypten.cat([encr_intermediate2,
encr_intermediate1]).sum()
return output, encr_output
# loop over test cases:
test_cases = [
value for key, value in locals().items()
if callable(value) and key.startswith("test_case")
]
for idx, test_case in enumerate(test_cases):
# get input tensors:
input = get_random_test_tensor(size=(12, 5), is_float=True)
input.requires_grad = True
encr_input = crypten.cryptensor(input, requires_grad=True)
# 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 = crypten.cryptensor(input, requires_grad=True)
# 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)