def _required_space_to_batch_paddings(converter, node, inputs: List[str]):
inputs_node = [converter.outputs[inputs[i]] for i in range(len(inputs))]
inputs_int32 = []
for x_in in inputs_node:
pvt_check = isinstance(x_in, PondPrivateTensor)
msk_check = isinstance(x_in, PondMaskedTensor)
if pvt_check or msk_check:
logging.warning(("Revealing private input: "
"required_space_to_batch_paddings assumes public "
"input."))
inputs_int32.append(tf.cast(x_in.reveal().decode(), tf.int32))
elif isinstance(x_in, tf.NodeDef):
inputs_int32.append(_nodef_to_numpy_array(x_in))
else:
raise TypeError("Unexpected input of type {}.".format(type(x_in)))
if len(inputs_int32) == 2:
input_shape, block_shape = inputs_int32
def inputter_pad():
pads, _ = tf.required_space_to_batch_paddings(
input_shape, block_shape)
return tf.cast(pads, tf.float64)
def inputter_crop():
_, crops = tf.required_space_to_batch_paddings(
input_shape, block_shape)
return tf.cast(crops, tf.float64)
else:
base_paddings, input_shape, block_shape = inputs_int32
def inputter_pad():
pads, _ = tf.required_space_to_batch_paddings(
input_shape,
block_shape,
base_paddings=base_paddings,
)
return tf.cast(pads, tf.float64)
def inputter_crop():
_, crops = tf.required_space_to_batch_paddings(
input_shape,
block_shape,
base_paddings=base_paddings,
)
return tf.cast(crops, tf.float64)
pad_private = tfe.define_public_input(
converter.model_provider,
inputter_pad,
)
crop_private = tfe.define_public_input(
converter.model_provider,
inputter_crop,
)
return (pad_private, crop_private)
def _nodef_to_public_pond(converter, x):
"""Map a NodeDef x to a PublicPondTensor."""
dtype = x.attr["dtype"].type
x_shape = [i.size for i in x.attr["value"].tensor.tensor_shape.dim]
if not x_shape:
if dtype == tf.float32:
nums = x.attr["value"].tensor.float_val
elif dtype == tf.float64:
nums = x.attr["value"].tensor.float_val
elif dtype == tf.int32:
nums = x.attr["value"].tensor.int_val
else:
raise TypeError("Unsupported dtype")
def inputter_fn():
return tf.constant(np.array(nums).reshape(1, 1))
else:
if dtype == tf.float32:
nums = array.array('f', x.attr["value"].tensor.tensor_content)
elif dtype == tf.float64:
nums = array.array('d', x.attr["value"].tensor.tensor_content)
elif dtype == tf.int32:
nums = array.array('i', x.attr["value"].tensor.tensor_content)
else:
raise TypeError("Unsupported dtype")
def inputter_fn():
return tf.constant(np.array(nums).reshape(x_shape))
x_public = tfe.define_public_input(
converter.model_provider, inputter_fn)
return x_public
def _rsqrt(converter, node: Any, inputs: List[str]) -> Any:
x_in = converter.outputs[inputs[0]]
if isinstance(x_in, tf.NodeDef):
tensor = x_in.attr["value"].tensor
shape = [i.size for i in tensor.tensor_shape.dim]
dtype = x_in.attr["dtype"].type
if dtype == tf.float32:
nums = array.array('f', tensor.tensor_content)
elif dtype == tf.float64:
nums = array.array('d', tensor.tensor_content)
else:
raise TypeError("Unsupported dtype for rsqrt")
def inputter_fn():
return tf.constant(1 / np.sqrt(np.array(nums).reshape(shape)))
else:
# XXX this is a little weird but the input into rsqrt is public and
# being used only for batchnorm at the moment
decoded = tfe._decode(x_in.value_on_0, True) # pylint: disable=protected-access
def inputter_fn():
return tf.rsqrt(decoded)
x = tfe.define_public_input(
converter.model_provider, inputter_fn)
return x
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")