mirror of
https://github.com/luxfhe/docs.git
synced 2026-10-30 18:13:39 +00:00
11 KiB
11 KiB
In [1]:
import platform
import time
from shutil import copyfile
from tempfile import TemporaryDirectory
import numpy
from sklearn.datasets import load_breast_cancer
from concrete.ml.deployment import FHEModelClient, FHEModelDev, FHEModelServer
from concrete.ml.sklearn import XGBClassifierIn [2]:
class OnDiskNetwork:
"""Simulate a network on disk."""
def __init__(self):
# Create 3 temporary folder for server, client and dev with tempfile
self.server_dir = TemporaryDirectory() # pylint: disable=consider-using-with
self.client_dir = TemporaryDirectory() # pylint: disable=consider-using-with
self.dev_dir = TemporaryDirectory() # pylint: disable=consider-using-with
def client_send_evaluation_key_to_server(self, serialized_evaluation_keys):
"""Send the public key to the server."""
with open(self.server_dir.name + "/serialized_evaluation_keys.ekl", "wb") as f:
f.write(serialized_evaluation_keys)
def client_send_input_to_server_for_prediction(self, encrypted_input):
"""Send the input to the server and execute on the server in FHE."""
with open(self.server_dir.name + "/serialized_evaluation_keys.ekl", "rb") as f:
serialized_evaluation_keys = f.read()
time_begin = time.time()
encrypted_prediction = FHEModelServer(self.server_dir.name).run(
encrypted_input, serialized_evaluation_keys
)
time_end = time.time()
with open(self.server_dir.name + "/encrypted_prediction.enc", "wb") as f:
f.write(encrypted_prediction)
return time_end - time_begin
def dev_send_model_to_server(self):
"""Send the model to the server."""
copyfile(self.dev_dir.name + "/server.zip", self.server_dir.name + "/server.zip")
def server_send_encrypted_prediction_to_client(self):
"""Send the encrypted prediction to the client."""
with open(self.server_dir.name + "/encrypted_prediction.enc", "rb") as f:
encrypted_prediction = f.read()
return encrypted_prediction
def dev_send_clientspecs_and_modelspecs_to_client(self):
"""Send the clientspecs and evaluation key to the client."""
copyfile(self.dev_dir.name + "/client.zip", self.client_dir.name + "/client.zip")
def cleanup(self):
"""Clean up the temporary folders."""
self.server_dir.cleanup()
self.client_dir.cleanup()
self.dev_dir.cleanup()In [3]:
from concrete.compiler import check_gpu_available
# Let's first get some data and train a model.
X, y = load_breast_cancer(return_X_y=True)
# Split X into X_model_owner and X_client
X_model_owner, X_client = X[:-10], X[-10:]
y_model_owner, y_client = y[:-10], y[-10:]
# Some issues on macOS, if too many estimators
n_estimators = 10
if platform.system() == "Darwin":
n_estimators = 9
use_gpu_if_available = False
device = "cuda" if use_gpu_if_available and check_gpu_available() else "cpu"
# Train the model and compile it
model_dev = XGBClassifier(n_bits=2, n_estimators=n_estimators, max_depth=3)
model_dev.fit(X_model_owner, y_model_owner)
model_dev.compile(X_model_owner, device=device)
print("Model trained and compiled.")Model trained and compiled.
In [4]:
# Let's instantiate the network
network = OnDiskNetwork()
# Now that the model has been trained we want to save it to send it to a server
fhemodel_dev = FHEModelDev(network.dev_dir.name, model_dev)
fhemodel_dev.save()In [5]:
# Print all files in the temporary directory along with their sizes in KB
!ls -lh $network.dev_dir.nametotal 64K -rw-r--r-- 1 root root 59K Sep 5 14:28 client.zip -rw-r--r-- 1 root root 2.6K Sep 5 14:28 server.zip
In [6]:
# Let's send the model to the server
network.dev_send_model_to_server()
!ls -lh $network.server_dir.nametotal 4.0K -rw-r--r-- 1 root root 2.6K Sep 5 14:28 server.zip
In [7]:
# Let's send the clientspecs and evaluation key to the client
network.dev_send_clientspecs_and_modelspecs_to_client()
!ls -lh $network.client_dir.nametotal 60K -rw-r--r-- 1 root root 59K Sep 5 14:28 client.zip
In [8]:
# Let's create the client and load the model
fhemodel_client = FHEModelClient(network.client_dir.name, key_dir=network.client_dir.name)
# The client first need to create the private and evaluation keys.
serialized_evaluation_keys = fhemodel_client.get_serialized_evaluation_keys()KeySetCache: miss, regenerating /tmp/tmpsf57p3ed/3497183917896914639
In [9]:
# Evaluation keys can be quite large files but only have to be shared once with the server.
# Check the size of the evaluation keys (in MB)
print(f"Evaluation keys size: {len(serialized_evaluation_keys) / (10**6):.2f} MB")
# Let's send this evaluation key to the server (this has to be done only once)
network.client_send_evaluation_key_to_server(serialized_evaluation_keys)Evaluation keys size: 29.14 MB
In [10]:
# Now we have everything for the client to interact with the server
# We create a loop to send the input to the server and receive the encrypted prediction
decrypted_predictions = []
execution_time = []
for i in range(X_client.shape[0]):
clear_input = X_client[[i], :]
encrypted_input = fhemodel_client.quantize_encrypt_serialize(clear_input)
execution_time += [network.client_send_input_to_server_for_prediction(encrypted_input)]
encrypted_prediction = network.server_send_encrypted_prediction_to_client()
decrypted_prediction = fhemodel_client.deserialize_decrypt_dequantize(encrypted_prediction)[0]
decrypted_predictions.append(decrypted_prediction)
# Check MB size with sys of the encrypted data vs clear data
print(
f"Encrypted data is "
f"{len(encrypted_input)/clear_input.nbytes:.2f}"
" times larger than the clear data"
)
# Show execution time
print(f"The average execution time is {numpy.mean(execution_time):.2f} seconds per sample.")Encrypted data is 4.10 times larger than the clear data The average execution time is 0.62 seconds per sample.
In [11]:
# Let's check the results and compare them against the clear model
clear_prediction_classes = model_dev.predict_proba(X_client).argmax(axis=1)
decrypted_predictions_classes = numpy.array(decrypted_predictions).argmax(axis=1)
accuracy = (clear_prediction_classes == decrypted_predictions_classes).mean()
print(f"Accuracy between FHE prediction and clear model is: {accuracy*100:.0f}%")Accuracy between FHE prediction and clear model is: 100%