model_trainer = ModelTrainer()
prediction_client = PredictionClient()
# get model parameters as private tensors from model owner
params = tfe.define_private_input('model-trainer', model_trainer.provide_input, masked=True) # pylint: disable=E0632
# we'll use the same parameters for each prediction so we cache them to avoid re-training each time
params = tfe.cache(params)
# get prediction input from client
x, y = tfe.define_private_input('prediction-client', prediction_client.provide_input, masked=True) # pylint: disable=E0632
# helpers
conv = lambda x, w, s: tfe.conv2d(x, w, s, 'VALID')
pool = lambda x: tfe.avgpool2d(x, (2, 2), (2, 2), 'VALID')
# compute prediction
Wconv1, bconv1, Wfc1, bfc1, Wfc2, bfc2 = params
bconv1 = tfe.reshape(bconv1, [-1, 1, 1])
layer1 = pool(tfe.relu(conv(x, Wconv1, ModelTrainer.STRIDE) + bconv1))
layer1 = tfe.reshape(layer1, [-1, ModelTrainer.HIDDEN_FC1])
layer2 = tfe.matmul(layer1, Wfc1) + bfc1
logits = tfe.matmul(layer2, Wfc2) + bfc2
# send prediction output back to client
prediction_op = tfe.define_output('prediction-client', [logits, y], prediction_client.receive_output)
with tfe.Session() as sess:
print("Init")
def pool(self, x, pool_size, strides, padding):
return tfe.avgpool2d(x, pool_size, strides, padding)
def pool(x):
return tfe.avgpool2d(x, (2, 2), (2, 2), 'VALID')