compile.py

"A script to manually compile all filters."

import json
import shutil
import onnx

from common import AVAILABLE_FILTERS, FILTERS_PATH, KEYS_PATH
from client_server_interface import FHEClient, FHEDev

print("Starting compiling the filters.")

for filter_name in AVAILABLE_FILTERS:
    print("\nCompiling filter:", filter_name)

    # Retrieve the deployment files associated to the current filter
    deployment_path = FILTERS_PATH / f"{filter_name}/deployment"

    # Retrieve the client associated to the current filter
    filter = FHEClient(deployment_path, KEYS_PATH).filter

    # Load the onnx graph
    onnx_graph = onnx.load(FILTERS_PATH / f"{filter_name}/server.onnx")

    # Compile the filter on a representative inputset, using the loaded onnx graph
    filter.compile(onnx_graph=onnx_graph)

    processing_json_path = deployment_path / "serialized_processing.json"

    # Load the serialized_processing.json file
    with open(processing_json_path, "r") as f:
        serialized_processing = json.load(f)

    # Delete the deployment folder and its content if it exist
    if deployment_path.is_dir():
        shutil.rmtree(deployment_path)

    # Save the development files needed for deployment
    fhe_dev = FHEDev(filter, deployment_path)
    fhe_dev.save()

    # Write the serialized_processing.json file in the deployment directory
    with open(processing_json_path, "w") as f:
        json.dump(serialized_processing, f)

print("Done!")