tfe-lowlevel.py

import tensorflow as tf
import tf_encrypted as tfe

def provide_weights(): """Load model weight from disk using TensorFlow."""
def provide_input(): """Load and preprocess input data locally on the client."""
def receive_output(logits): return tf.print(tf.argmax(logits))

w0, b0, w1, b1, w2, b2 = provide_weights()

# run provide_input locally on the client and encrypt
x = tfe.define_private_input("prediction-client", provide_input)

# compute prediction on the encrypted input
layer0 = tfe.relu((tfe.matmul(x, w0) + b0))
layer1 = tfe.relu((tfe.matmul(layer0, w1) + b1))
logits = tfe.matmul(layer1, w2) + b2

# send result back to client, decrypt, and run receive_output locally
prediction_op = tfe.define_output("prediction-client", receive_output, logits)

with tfe.Session() as sess:
    sess.run(prediction_op)