tfe-keras.py

import tensorflow as tf
import tf_encrypted as tfe

class PredictionClient:

    @tfe.private_input
    def provide_input(self):
        """Load and preprocess input data."""

    @tfe.private_output
    def receive_output(self, logits):
        return tf.print(tf.argmax(logits))

model = tfe.models.Sequential([
    tfe.keras.layers.Dense(activation='relu'),
    tfe.keras.layers.Dense(activation='relu'),
    tfe.keras.layers.Dense(activation=None)
])

prediction_client = PredictionClient()
x = prediction_client.provide_input()

y = model.predict(x)

prediction_client.receive_output(y)