def test_rshift_private():
tf.reset_default_graph()
prot = ABY3()
tfe.set_protocol(prot)
x = tfe.define_private_variable(tf.constant([[1, 2, 3], [4, 5, 6]]),
share_type=BOOLEAN)
y = tfe.define_private_variable(tf.constant([[-1, -2, -3], [-4, 5, 6]]),
share_type=BOOLEAN,
apply_scaling=False)
z = x >> 1
w = y >> 1
s = y.logical_rshift(1)
with tfe.Session() as sess:
# initialize variables
sess.run(tfe.global_variables_initializer())
# reveal result
result = sess.run(z.reveal())
close(result, np.array(
[[0.5, 1, 1.5],
[2, 2.5,
3]])) # NOTE: x is scaled and treated as fixed-point number
result = sess.run(w.reveal())
close(result, np.array([[-1, -1, -2], [-2, 2, 3]]))
result = sess.run(s.reveal())
close(
result,
np.array([[(-1 & ((1 << prot.nbits) - 1)) >> 1,
(-2 & ((1 << prot.nbits) - 1)) >> 1,
(-3 & ((1 << prot.nbits) - 1)) >> 1],
[(-4 & ((1 << prot.nbits) - 1)) >> 1, 2, 3]]))
print("test_rshift_private succeeds")
def test_ot():
tf.reset_default_graph()
prot = ABY3()
tfe.set_protocol(prot)
m0 = prot.define_constant(np.array([[1, 2, 3], [4, 5, 6]]),
apply_scaling=False).unwrapped[0]
m1 = prot.define_constant(np.array([[2, 3, 4], [5, 6, 7]]),
apply_scaling=False).unwrapped[0]
c_on_receiver = prot.define_constant(
np.array([[1, 0, 1], [0, 1, 0]]),
apply_scaling=False,
factory=prot.bool_factory).unwrapped[0]
c_on_helper = prot.define_constant(np.array([[1, 0, 1], [0, 1, 0]]),
apply_scaling=False,
factory=prot.bool_factory).unwrapped[0]
m_c = prot._ot(prot.servers[1], prot.servers[2], prot.servers[0], m0, m1,
c_on_receiver, c_on_helper, prot.pairwise_keys[1][0],
prot.pairwise_keys[0][1], prot.pairwise_nonces[0])
with tfe.Session() as sess:
# initialize variables
sess.run(tfe.global_variables_initializer())
# reveal result
result = sess.run(prot._decode(m_c, False))
close(result, np.array([[2, 2, 4], [4, 6, 6]]))
print("test_ot succeeds")
def test_ppa_private_private():
tf.reset_default_graph()
prot = ABY3()
tfe.set_protocol(prot)
x = tfe.define_private_variable(tf.constant([[1, 2, 3], [4, 5, 6]]),
share_type=BOOLEAN)
y = tfe.define_private_variable(tf.constant([[7, 8, 9], [10, 11, 12]]),
share_type=BOOLEAN)
# Parallel prefix adder. It is simply an adder for boolean sharing.
z1 = tfe.B_ppa(x, y, topology="sklansky")
z2 = tfe.B_ppa(x, y, topology="kogge_stone")
with tfe.Session() as sess:
# initialize variables
sess.run(tfe.global_variables_initializer())
# reveal result
result = sess.run(z1.reveal())
close(result, np.array([[8, 10, 12], [14, 16, 18]]))
result = sess.run(z2.reveal())
close(result, np.array([[8, 10, 12], [14, 16, 18]]))
print("test_ppa_private_private succeeds")
def test_not_private():
tf.reset_default_graph()
prot = ABY3()
tfe.set_protocol(prot)
x = tfe.define_private_variable(tf.constant([[1, 2, 3], [4, 5, 6]]),
share_type=BOOLEAN,
apply_scaling=False)
y = tfe.define_private_variable(tf.constant([[1, 0, 0], [0, 1, 0]]),
apply_scaling=False,
share_type=BOOLEAN,
factory=prot.bool_factory)
z1 = ~x
z2 = ~y
with tfe.Session() as sess:
# initialize variables
sess.run(tfe.global_variables_initializer())
# reveal result
result = sess.run(z1.reveal())
close(result, np.array([[-2, -3, -4], [-5, -6, -7]]))
result = sess.run(z2.reveal())
close(result, np.array([[0, 1, 1], [1, 0, 1]]))
print("test_not_private succeeds")
def test_boolean_sharing():
tf.reset_default_graph()
prot = ABY3()
tfe.set_protocol(prot)
x = tfe.define_private_variable(tf.constant([[1, 2, 3], [4, 5, 6]]),
share_type=BOOLEAN)
y = tfe.define_private_variable(tf.constant([[7, 8, 9], [10, 11, 12]]),
share_type=BOOLEAN)
z1 = tfe.B_xor(x, y)
z2 = tfe.B_and(x, y)
with tfe.Session() as sess:
# initialize variables
sess.run(tfe.global_variables_initializer())
# reveal result
result = sess.run(z1.reveal())
close(result, np.array([[6, 10, 10], [14, 14, 10]]))
result = sess.run(z2.reveal())
close(result, np.array([[1, 0, 1], [0, 1, 4]]))
print("test_boolean_sharing succeeds")
def test_matmul_public_private():
tf.reset_default_graph()
prot = ABY3()
tfe.set_protocol(prot)
def provide_input():
# normal TensorFlow operations can be run locally
# as part of defining a private input, in this
# case on the machine of the input provider
return tf.constant(np.array([[1.1, 1.2], [1.3, 1.4], [1.5, 1.6]]))
# define inputs
x = tfe.define_private_variable(tf.ones(shape=(2, 2)))
y = tfe.define_public_input('input-provider', provide_input)
v = tfe.define_constant(np.ones((2, 2)))
# define computation
w = y.matmul(x) # matmul_public_private
z = w.matmul(v) # matmul_private_public
with tfe.Session() as sess:
# initialize variables
sess.run(tfe.global_variables_initializer())
# reveal result
result = sess.run(w.reveal())
close(result, np.array([[2.3, 2.3], [2.7, 2.7], [3.1, 3.1]]))
result = sess.run(z.reveal())
close(result, np.array([[4.6, 4.6], [5.4, 5.4], [6.2, 6.2]]))
print("test_matmul_public_private succeeds")
def test_mul_private_private():
tf.reset_default_graph()
prot = ABY3()
tfe.set_protocol(prot)
def provide_input():
# normal TensorFlow operations can be run locally
# as part of defining a private input, in this
# case on the machine of the input provider
return tf.ones(shape=(2, 2)) * 1.3
# define inputs
x = tfe.define_private_variable(tf.ones(shape=(2, 2)) * 2)
y = tfe.define_private_input("input-provider", provide_input)
# define computation
z = y * x
with tfe.Session() as sess:
# initialize variables
sess.run(tfe.global_variables_initializer())
# reveal result
result = sess.run(z.reveal())
close(result, np.array([[2.6, 2.6], [2.6, 2.6]]))
print("test_mul_private_private succeeds")
def test_iterate_private():
tf.reset_default_graph()
prot = ABY3()
tfe.set_protocol(prot)
def provide_input():
return tf.reshape(tf.range(0, 8), [4, 2])
# define inputs
x = tfe.define_private_input('input-provider', provide_input)
write_op = x.write("x.tfrecord")
with tfe.Session() as sess:
# initialize variables
sess.run(tfe.global_variables_initializer())
# reveal result
sess.run(write_op)
x = tfe.read("x.tfrecord", batch_size=5, n_columns=2)
y = tfe.iterate(x, batch_size=3, repeat=True, shuffle=False)
z = tfe.iterate(x, batch_size=3, repeat=True, shuffle=True)
with tfe.Session() as sess:
sess.run(tfe.global_variables_initializer())
print(sess.run(x.reveal()))
print(sess.run(y.reveal()))
print(sess.run(y.reveal()))
print(sess.run(x.reveal()))
print(sess.run(z.reveal()))
def test_read_private():
tf.reset_default_graph()
prot = ABY3()
tfe.set_protocol(prot)
def provide_input():
return tf.reshape(tf.range(0, 8), [4, 2])
# define inputs
x = tfe.define_private_input('input-provider', provide_input)
write_op = x.write("x.tfrecord")
with tfe.Session() as sess:
# initialize variables
sess.run(tfe.global_variables_initializer())
# reveal result
sess.run(write_op)
x = tfe.read("x.tfrecord", batch_size=5, n_columns=2)
with tfe.Session() as sess:
result = sess.run(x.reveal())
close(result, np.array(list(range(0, 8)) + [0, 1]).reshape([5, 2]))
print("test_read_private succeeds")
def __exit__(
self,
# type is `Optional[Type[BaseException]]`, but declaring `Type` breaks readthedocs.
exception_type,
exception_value: Optional[Exception],
traceback: Optional[TracebackType],
) -> Optional[bool]:
tfe.set_protocol(self.last_protocol)
def start_master(cluster_config_file=None):
print("Starting alice...")
remote_config = tfe.RemoteConfig.load(cluster_config_file)
tfe.set_config(remote_config)
tfe.set_protocol(tfe.protocol.Pond())
players = remote_config.players
server0 = remote_config.server(players[0].name)
main(server0)
def start_slave(cluster_config_file):
print("Starting crypto producer...")
remote_config = tfe.RemoteConfig.load(cluster_config_file)
tfe.set_config(remote_config)
tfe.set_protocol(tfe.protocol.Pond())
players = remote_config.players
server = remote_config.server(players[2].name)
print("server_name = ", players[2].name)
server.join()
def start_slave(cluster_config_file):
print("Starting bob...")
remote_config = tfe.RemoteConfig.load(cluster_config_file)
tfe.set_config(remote_config)
tfe.set_protocol(tfe.protocol.Pond())
players = remote_config.players
bob = remote_config.server(players[1].name)
print("server_name = ", players[1].name)
bob.join()
def main(server):
num_rows = 7000
num_features = 32
num_epoch = 5
batch_size = 200
num_batches = (num_rows // batch_size) * num_epoch
#who shall receive the output
model_owner = ModelOwner('alice')
data_schema0 = DataSchema([tf.float64] * 16, [0.0] * 16)
data_schema1 = DataSchema([tf.int64] + [tf.float64] * 16, [0] + [0.0] * 16)
data_owner_0 = DataOwner('alice',
'aliceTrainFile.csv',
data_schema0,
batch_size=batch_size)
data_owner_1 = DataOwner('bob',
'bobTrainFileWithLabel.csv',
data_schema1,
batch_size=batch_size)
tfe.set_protocol(
tfe.protocol.Pond(
tfe.get_config().get_player(data_owner_0.player_name),
tfe.get_config().get_player(data_owner_1.player_name)))
x_train_0 = tfe.define_private_input(data_owner_0.player_name,
data_owner_0.provide_data)
x_train_1 = tfe.define_private_input(data_owner_1.player_name,
data_owner_1.provide_data)
y_train = tfe.gather(x_train_1, 0, axis=1)
y_train = tfe.reshape(y_train, [batch_size, 1])
#Remove bob's first column (which is label)
x_train_1 = tfe.strided_slice(x_train_1, [0, 1],
[x_train_1.shape[0], x_train_1.shape[1]],
[1, 1])
x_train = tfe.concat([x_train_0, x_train_1], axis=1)
model = LogisticRegression(num_features)
reveal_weights_op = model_owner.receive_weights(model.weights)
with tfe.Session() as sess:
sess.run(tfe.global_variables_initializer(), tag='init')
start_time = time.time()
model.fit(sess, x_train, y_train, num_batches)
end_time = time.time()
# TODO(Morten)
# each evaluation results in nodes for a forward pass being added to the graph;
# maybe there's some way to avoid this, even if it means only if the shapes match
model.evaluate(sess, x_train, y_train, data_owner_0)
model.evaluate(sess, x_train, y_train, data_owner_1)
print(sess.run(reveal_weights_op, tag='reveal'),
((end_time - start_time) * 1000))
def start_master(cluster_config_file=None):
print("Starting alice...")
remote_config = tfe.RemoteConfig.load(cluster_config_file)
tfe.set_config(remote_config)
tfe.set_protocol(tfe.protocol.Pond())
players = remote_config.players
server0 = remote_config.server(players[0].name)
st = time.perf_counter()
main(server0)
ed = time.perf_counter()
print(f'Elapsed time: {ed - st}s')
def _configure_tfe(cluster):
if not cluster or len(cluster.workers) != 3:
raise RuntimeError(
"TF Encrypted expects three parties for its sharing protocols.")
config = cluster.tfe_config
tfe.set_config(config)
prot = tfe.protocol.SecureNN(config.get_player("server0"),
config.get_player("server1"),
config.get_player("server2"))
tfe.set_protocol(prot)
tfe.clear_initializers()
def connect_to_model(self, input_shape, output_shape, *workers):
config, _ = self.config_from_workers(workers)
tfe.set_config(config)
prot = tfe.protocol.SecureNN(config.get_player("server0"),
config.get_player("server1"),
config.get_player("server2"))
tfe.set_protocol(prot)
self._tf_client = tfe.serving.QueueClient(input_shape=input_shape,
output_shape=output_shape)
sess = tfe.Session(config=config)
self._tf_session = sess
def test_simple_lr_model():
tf.reset_default_graph()
import time
start = time.time()
prot = ABY3()
tfe.set_protocol(prot)
# define inputs
x_raw = tf.random.uniform(minval=-0.5,
maxval=0.5,
shape=[99, 10],
seed=1000)
x = tfe.define_private_variable(x_raw, name="x")
y_raw = tf.cast(tf.reduce_mean(x_raw, axis=1, keepdims=True) > 0,
dtype=tf.float32)
y = tfe.define_private_variable(y_raw, name="y")
w = tfe.define_private_variable(tf.random_uniform([10, 1],
-0.01,
0.01,
seed=100),
name="w")
b = tfe.define_private_variable(tf.zeros([1]), name="b")
learning_rate = 0.01
with tf.name_scope("forward"):
out = tfe.matmul(x, w) + b
y_hat = tfe.sigmoid(out)
with tf.name_scope("loss-grad"):
dy = y_hat - y
batch_size = x.shape.as_list()[0]
with tf.name_scope("backward"):
dw = tfe.matmul(tfe.transpose(x), dy) / batch_size
db = tfe.reduce_sum(dy, axis=0) / batch_size
upd1 = dw * learning_rate
upd2 = db * learning_rate
assign_ops = [tfe.assign(w, w - upd1), tfe.assign(b, b - upd2)]
with tfe.Session() as sess:
# initialize variables
sess.run(tfe.global_variables_initializer())
for i in range(1):
sess.run(assign_ops)
print(sess.run(w.reveal()))
end = time.time()
print("Elapsed time: {} seconds".format(end - start))
def test_bit_extract(self):
tf.reset_default_graph()
prot = ABY3()
tfe.set_protocol(prot)
x = tfe.define_private_variable(
np.array([[1, -2, 3], [-4, -5, 6]]),
share_type=ARITHMETIC,
)
y = tfe.define_private_variable(
np.array([[1, -2, 3], [-4, -5, 6]]),
share_type=ARITHMETIC,
apply_scaling=False,
)
z = tfe.bit_extract(
x, 63
) # The sign bit. Since x is scaled, you should be more careful about extracting other bits.
w = tfe.bit_extract(y, 1) # y is not scaled
s = tfe.msb(x) # Sign bit
with tfe.Session() as sess:
# initialize variables
sess.run(tfe.global_variables_initializer())
# reveal result
result = sess.run(z.reveal())
np.testing.assert_allclose(
result.astype(int),
np.array([[0, 1, 0], [1, 1, 0]]),
rtol=0.0,
atol=0.01,
)
result = sess.run(w.reveal())
np.testing.assert_allclose(
result.astype(int),
np.array([[0, 1, 1], [0, 1, 1]]),
rtol=0.0,
atol=0.01,
)
result = sess.run(s.reveal())
np.testing.assert_allclose(
result.astype(int),
np.array([[0, 1, 0], [1, 1, 0]]),
rtol=0.0,
atol=0.01,
)
def _configure_tfe(workers):
if not workers or len(workers) != 3:
raise RuntimeError(
"TF Encrypted expects three parties for its sharing protocols.")
tfe_worker_cls = workers[0].__class__
config, player_to_worker_mapping = tfe_worker_cls.config_from_workers(
workers)
tfe.set_config(config)
prot = tfe.protocol.SecureNN(config.get_player("server0"),
config.get_player("server1"),
config.get_player("server2"))
tfe.set_protocol(prot)
return player_to_worker_mapping
def test_lshift_private():
tf.reset_default_graph()
prot = ABY3()
tfe.set_protocol(prot)
x = tfe.define_private_variable(tf.constant([[1, 2, 3], [4, 5, 6]]),
share_type=BOOLEAN)
z = x << 1
with tfe.Session() as sess:
# initialize variables
sess.run(tfe.global_variables_initializer())
# reveal result
result = sess.run(z.reveal())
close(result, np.array([[2, 4, 6], [8, 10, 12]]))
print("test_lshift_private succeeds")
def __init__(
self,
registry,
config: Optional[Config] = None,
protocol: Optional[Protocol] = None,
model_provider: Optional[Union[str, Player]] = None,
) -> None:
self.config = config if config is not None else get_config()
if protocol is not None:
tfe.set_protocol(protocol)
if model_provider is None:
self.model_provider = self.config.get_player("model-provider")
elif isinstance(model_provider, str):
self.model_provider = self.config.get_player(model_provider)
else:
self.model_provider = model_provider
self.registry = registry
self.outputs = {}
def test_a2b_private():
tf.reset_default_graph()
prot = ABY3()
tfe.set_protocol(prot)
x = tfe.define_private_variable(tf.constant([[1, 2, 3], [4, 5, 6]]),
share_type=ARITHMETIC)
z = tfe.A2B(x)
assert z.share_type == BOOLEAN
with tfe.Session() as sess:
# initialize variables
sess.run(tfe.global_variables_initializer())
# reveal result
result = sess.run(z.reveal())
close(result, np.array([[1, 2, 3], [4, 5, 6]]))
print("test_a2b_private succeeds")
def test_matmul_private_private(self):
tf.reset_default_graph()
prot = ABY3()
tfe.set_protocol(prot)
# 2-D matrix mult
x = tfe.define_private_variable(tf.constant([[1, 2, 3], [4, 5, 6]]))
y = tfe.define_private_variable(tf.constant([[7, 8], [9, 10], [11, 12]]))
z = tfe.matmul(x, y)
with tfe.Session() as sess:
# initialize variables
sess.run(tfe.global_variables_initializer())
# reveal result
result = sess.run(z.reveal())
np.testing.assert_allclose(
result, np.array([[58, 64], [139, 154]]), rtol=0.0, atol=0.01
)
def test_error():
tf.reset_default_graph()
prot = ABY3()
tfe.set_protocol(prot)
a = tf.random_uniform([6, 10000], -3, 3)
b = tf.random_uniform([10000, 1], -0.1, 0.1)
x = tfe.define_private_input("input-provider", lambda: a)
y = tfe.define_private_input("input-provider", lambda: b)
z = tfe.matmul(x, y)
w = tf.matmul(a, b)
with tfe.Session() as sess:
sess.run(tfe.global_variables_initializer())
for i in range(1000):
result1, result2 = sess.run([z.reveal(), w])
close(result1, result2, 0.1)
print("test_error succeeds")
def test_add_private_public(self):
tf.reset_default_graph()
prot = ABY3()
tfe.set_protocol(prot)
# define inputs
x = tfe.define_private_variable(tf.ones(shape=(2, 2)))
y = tfe.define_constant(np.array([[0.6, 0.7], [0.8, 0.9]]))
# define computation
z = x + y
z = y + z
with tfe.Session() as sess:
# initialize variables
sess.run(tfe.global_variables_initializer())
# reveal result
result = sess.run(z.reveal())
expected = np.array([[2.2, 2.4], [2.6, 2.8]])
np.testing.assert_allclose(result, expected, rtol=0.0, atol=0.01)
def test_sub_private_private():
tf.reset_default_graph()
prot = ABY3()
tfe.set_protocol(prot)
def provide_input():
return tf.ones(shape=(2, 2)) * 1.3
x = tfe.define_private_variable(tf.ones(shape=(2, 2)))
y = tfe.define_private_input('input-provider', provide_input)
z = x - y
with tfe.Session() as sess:
# initialize variables
sess.run(tfe.global_variables_initializer())
# reveal result
result = sess.run(z.reveal())
close(result, np.array([[-0.3, -0.3], [-0.3, -0.3]]))
print("test_sub_private_private succeeds")
def test_b2a_private(self):
tf.reset_default_graph()
prot = ABY3()
tfe.set_protocol(prot)
x = tfe.define_private_variable(
tf.constant([[1, 2, 3], [4, 5, 6]]), share_type=BOOLEAN
)
z = tfe.B2A(x)
assert z.share_type == ARITHMETIC
with tfe.Session() as sess:
# initialize variables
sess.run(tfe.global_variables_initializer())
# reveal result
result = sess.run(z.reveal())
np.testing.assert_allclose(
result, np.array([[1, 2, 3], [4, 5, 6]]), rtol=0.0, atol=0.01
)
def test_sub_private_private(self):
tf.reset_default_graph()
prot = ABY3()
tfe.set_protocol(prot)
def provide_input():
return tf.ones(shape=(2, 2)) * 1.3
x = tfe.define_private_variable(tf.ones(shape=(2, 2)))
y = tfe.define_private_input("input-provider", provide_input)
z = x - y
with tfe.Session() as sess:
# initialize variables
sess.run(tfe.global_variables_initializer())
# reveal result
result = sess.run(z.reveal())
expected = np.array([[-0.3, -0.3], [-0.3, -0.3]])
np.testing.assert_allclose(result, expected, rtol=0.0, atol=0.01)
def connect_to_model(self, input_shape, output_shape, cluster, sess=None):
"""
Connect to a TF Encrypted model being served by the given cluster.
This must be done before querying the model.
"""
config = cluster.tfe_config
tfe.set_config(config)
prot = tfe.protocol.SecureNN(config.get_player("server0"),
config.get_player("server1"),
config.get_player("server2"))
tfe.set_protocol(prot)
self._tf_client = tfe.serving.QueueClient(input_shape=input_shape,
output_shape=output_shape)
if sess is None:
sess = tfe.Session(config=config)
self._tf_session = sess
