examples/notebooks/proof_splitting.ipynb

{
    "cells": [
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "# Proof splitting\n",
                "\n",
                "Here we showcase how to split a larger circuit into multiple smaller proofs. This is useful if you want to prove over multiple machines, or if you want to split a proof into multiple parts to reduce the memory requirements.\n",
                "\n",
                "We showcase how to do this in the case where:\n",
                "- intermediate calculations can be public (i.e. they do not need to be kept secret) and we can stitch the circuits together using instances\n",
                "- intermediate calculations need to be kept secret (but not blinded !)  and we need to use the low overhead kzg commitment scheme detailed [here](https://blog.ezkl.xyz/post/commits/) to stitch the circuits together. \n"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "First we import the necessary dependencies and set up logging to be as informative as possible. "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# check if notebook is in colab\n",
                "try:\n",
                "    # install ezkl\n",
                "    import google.colab\n",
                "    import subprocess\n",
                "    import sys\n",
                "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"ezkl\"])\n",
                "    subprocess.check_call([sys.executable, \"-m\", \"pip\", \"install\", \"onnx\"])\n",
                "\n",
                "# rely on local installation of ezkl if the notebook is not in colab\n",
                "except:\n",
                "    pass\n",
                "\n",
                "from torch import nn\n",
                "import ezkl\n",
                "import os\n",
                "import json\n",
                "import logging\n",
                "\n",
                "# uncomment for more descriptive logging \n",
                "# FORMAT = '%(levelname)s %(name)s %(asctime)-15s %(filename)s:%(lineno)d %(message)s'\n",
                "# logging.basicConfig(format=FORMAT)\n",
                "# logging.getLogger().setLevel(logging.INFO)\n"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Now we define our model. It is a humble model with but a conv layer and a $ReLU$ non-linearity, but it is a model nonetheless"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import torch\n",
                "# Defines the model\n",
                "# we got convs, we got relu, \n",
                "# What else could one want ????\n",
                "\n",
                "class MyModel(nn.Module):\n",
                "    def __init__(self):\n",
                "        super(MyModel, self).__init__()\n",
                "\n",
                "        self.conv1 = nn.Conv2d(in_channels=3, out_channels=1, kernel_size=2, stride=4)\n",
                "        self.conv2 = nn.Conv2d(in_channels=1, out_channels=1, kernel_size=2, stride=4)\n",
                "        self.relu = nn.ReLU()\n",
                "\n",
                "    def forward(self, x):\n",
                "        x = self.conv1(x)\n",
                "        x = self.relu(x)\n",
                "        x = self.conv2(x)\n",
                "        x = self.relu(x)\n",
                "\n",
                "        return x\n",
                "    \n",
                "    def split_1(self, x):\n",
                "        x = self.conv1(x)\n",
                "        x = self.relu(x)\n",
                "        return x\n",
                "\n",
                "\n",
                "circuit = MyModel()\n",
                "\n",
                "# this is where you'd train your model\n",
                "\n",
                "\n"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "We omit training for purposes of this demonstration. We've marked where training would happen in the cell above. \n",
                "Now we export the model to onnx and create a corresponding (randomly generated) input file.\n",
                "\n",
                "You can replace the random `x` with real data if you so wish. "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "x = torch.rand(1,*[3, 8, 8], requires_grad=True)\n",
                "\n",
                "# Flips the neural net into inference mode\n",
                "circuit.eval()\n",
                "\n",
                "    # Export the model\n",
                "torch.onnx.export(circuit,               # model being run\n",
                "                      x,                   # model input (or a tuple for multiple inputs)\n",
                "                      \"network.onnx\",            # where to save the model (can be a file or file-like object)\n",
                "                      export_params=True,        # store the trained parameter weights inside the model file\n",
                "                      opset_version=10,          # the ONNX version to export the model to\n",
                "                      do_constant_folding=True,  # whether to execute constant folding for optimization\n",
                "                      input_names = ['input'],   # the model's input names\n",
                "                      output_names = ['output'], # the model's output names\n",
                "                      dynamic_axes={'input' : {0 : 'batch_size'},    # variable length axes\n",
                "                                    'output' : {0 : 'batch_size'}})\n",
                "\n",
                "\n",
                "data_path = os.path.join(os.getcwd(), \"input_0.json\")\n",
                "data = dict(input_data = [((x).detach().numpy()).reshape([-1]).tolist()])\n",
                "json.dump( data, open(data_path, 'w' ))\n",
                "\n",
                "inter_1 = circuit.split_1(x)\n",
                "data_path = os.path.join(os.getcwd(), \"input_1.json\")\n",
                "data = dict(input_data = [((inter_1).detach().numpy()).reshape([-1]).tolist()])\n",
                "json.dump( data, open(data_path, 'w' ))\n",
                "\n"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Now we split the model into two parts. The first part is the first conv layer and the second part is the rest of the model."
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import onnx\n",
                "\n",
                "input_path = \"network.onnx\"\n",
                "output_path = \"network_split_0.onnx\"\n",
                "input_names = [\"input\"]\n",
                "output_names = [\"/relu/Relu_output_0\"]\n",
                "# first model\n",
                "onnx.utils.extract_model(input_path, output_path, input_names, output_names)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import onnx\n",
                "\n",
                "input_path = \"network.onnx\"\n",
                "output_path = \"network_split_1.onnx\"\n",
                "input_names = [\"/relu/Relu_output_0\"]\n",
                "output_names = [\"output\"]\n",
                "# second model\n",
                "onnx.utils.extract_model(input_path, output_path, input_names, output_names)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "### Public intermediate calculations"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "This is where the magic happens. We define our `PyRunArgs` objects which contains the visibility parameters for out model. \n",
                "- `input_visibility` defines the visibility of the model inputs\n",
                "- `param_visibility` defines the visibility of the model weights and constants and parameters \n",
                "- `output_visibility` defines the visibility of the model outputs\n",
                "\n",
                "There are currently 5 visibility settings:\n",
                "- `public`: known to both the verifier and prover (a subtle nuance is that this may not be the case for model parameters but until we have more rigorous theoretical results we don't want to make strong claims as to this). \n",
                "- `private`: known only to the prover\n",
                "- `hashed`: the hash pre-image is known to the prover, the prover and verifier know the hash. The prover proves that the they know the pre-image to the hash. \n",
                "- `encrypted`: the non-encrypted element and the secret key used for decryption are known to the prover. The prover and the verifier know the encrypted element, the public key used to encrypt, and the hash of the decryption hey. The prover proves that they know the pre-image of the hashed decryption key and that this key can in fact decrypt the encrypted message.\n",
                "- `polycommit`: unblinded advice column which generates a kzg commitment. This doesn't appear in the instances of the circuit and must instead be modified directly within the proof bytes.  \n",
                "\n",
                "Here we create the following setup:\n",
                "- `input_visibility`: \"public\"\n",
                "- `param_visibility`: \"public\"\n",
                "- `output_visibility`: public\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import ezkl\n",
                "\n",
                "\n",
                "data_path = os.path.join('input.json')\n",
                "\n",
                "run_args = ezkl.PyRunArgs()\n",
                "run_args.input_visibility = \"public\"\n",
                "run_args.param_visibility = \"fixed\"\n",
                "run_args.output_visibility = \"public\"\n",
                "run_args.input_scale = 2\n",
                "run_args.logrows = 8\n",
                "\n",
                "ezkl.get_srs(logrows=run_args.logrows, commitment=ezkl.PyCommitments.KZG)"
            ]
        },
        {
            "attachments": {},
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Now we generate a settings file. This file basically instantiates a bunch of parameters that determine their circuit shape, size etc... Because of the way we represent nonlinearities in the circuit (using Halo2's [lookup tables](https://zcash.github.io/halo2/design/proving-system/lookup.html)), it is often best to _calibrate_ this settings file as some data can fall out of range of these lookups.\n",
                "\n",
                "You can pass a dataset for calibration that will be representative of real inputs you might find if and when you deploy the prover. Here we create a dummy calibration dataset for demonstration purposes. \n",
                "\n",
                "As we use Halo2 with KZG-commitments we need an SRS string from (preferably) a multi-party trusted setup ceremony. For an overview of the procedures for such a ceremony check out [this page](https://blog.ethereum.org/2023/01/16/announcing-kzg-ceremony). The `get_srs` command retrieves a correctly sized SRS given the calibrated settings file from [here](https://github.com/han0110/halo2-kzg-srs). \n",
                "\n",
                "These SRS were generated with [this](https://github.com/privacy-scaling-explorations/perpetualpowersoftau) ceremony. \n",
                "\n",
                "We also need to generate the (partial) circuit witness. These are the model outputs (and any hashes) that are generated when feeding the previously generated `input.json` through the circuit / model. "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# iterate over each submodel gen-settings, compile circuit and setup zkSNARK\n",
                "\n",
                "async def setup(i):\n",
                "    # file names\n",
                "    model_path = os.path.join('network_split_'+str(i)+'.onnx')\n",
                "    settings_path = os.path.join('settings_split_'+str(i)+'.json')\n",
                "    data_path =  os.path.join('input_'+str(i)+'.json')\n",
                "    compiled_model_path = os.path.join('network_split_'+str(i)+'.compiled')\n",
                "    pk_path = os.path.join('test_split_'+str(i)+'.pk')\n",
                "    vk_path = os.path.join('test_split_'+str(i)+'.vk')\n",
                "    witness_path = os.path.join('witness_split_'+str(i)+'.json')\n",
                "\n",
                "    if i > 0:\n",
                "         prev_witness_path = os.path.join('witness_split_'+str(i-1)+'.json')\n",
                "         witness = json.load(open(prev_witness_path, 'r'))\n",
                "         data = dict(input_data = witness['outputs'])\n",
                "         # Serialize data into file:\n",
                "         json.dump(data, open(data_path, 'w' ))\n",
                "    else:\n",
                "         data_path = os.path.join('input_0.json')\n",
                "\n",
                "    # generate settings for the current model\n",
                "    res = ezkl.gen_settings(model_path, settings_path, py_run_args=run_args)\n",
                "    res = await ezkl.calibrate_settings(data_path, model_path, settings_path, \"resources\", scales=[run_args.input_scale], max_logrows=run_args.logrows)\n",
                "    assert res == True\n",
                "\n",
                "    # load settings and print them to the console\n",
                "    settings = json.load(open(settings_path, 'r'))\n",
                "    settings['run_args']['logrows'] = run_args.logrows\n",
                "    json.dump(settings, open(settings_path, 'w' ))\n",
                "\n",
                "    res = ezkl.compile_circuit(model_path, compiled_model_path, settings_path)\n",
                "\n",
                "\n",
                "    res = ezkl.setup(\n",
                "         compiled_model_path,\n",
                "         vk_path,\n",
                "         pk_path,\n",
                "      )\n",
                "\n",
                "    assert res == True\n",
                "    assert os.path.isfile(vk_path)\n",
                "    assert os.path.isfile(pk_path)\n",
                "\n",
                "    res = await ezkl.gen_witness(data_path, compiled_model_path, witness_path, vk_path)\n",
                "    run_args.input_scale = settings[\"model_output_scales\"][0]\n",
                "\n",
                "for i in range(2):\n",
                "    await setup(i)\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# GENERATE A PROOF\n",
                "def prove_model(i):\n",
                "    proof_path = os.path.join('proof_split_'+str(i)+'.json')\n",
                "    witness_path = os.path.join('witness_split_'+str(i)+'.json')\n",
                "    compiled_model_path = os.path.join('network_split_'+str(i)+'.compiled')\n",
                "    pk_path = os.path.join('test_split_'+str(i)+'.pk')\n",
                "    vk_path = os.path.join('test_split_'+str(i)+'.vk')\n",
                "    settings_path = os.path.join('settings_split_'+str(i)+'.json')\n",
                "\n",
                "    res = ezkl.prove(\n",
                "            witness_path,\n",
                "            compiled_model_path,\n",
                "            pk_path,\n",
                "            proof_path,\n",
                "            \"for-aggr\",\n",
                "        )\n",
                "\n",
                "    print(res)\n",
                "    res_1_proof = res[\"proof\"]\n",
                "    assert os.path.isfile(proof_path)\n",
                "\n",
                "    # # Verify the proof\n",
                "    if i > 0:\n",
                "        print(\"swapping commitments\")\n",
                "        # swap the proof commitments if we are not the first model\n",
                "        prev_witness_path = os.path.join('witness_split_'+str(i-1)+'.json')\n",
                "        prev_witness = json.load(open(prev_witness_path, 'r'))\n",
                "\n",
                "        witness = json.load(open(witness_path, 'r'))\n",
                "\n",
                "        print(prev_witness[\"processed_outputs\"])\n",
                "        print(witness[\"processed_inputs\"])\n",
                "        witness[\"processed_inputs\"] = prev_witness[\"processed_outputs\"]\n",
                "\n",
                "        # now save the witness\n",
                "        with open(witness_path, \"w\") as f:\n",
                "            json.dump(witness, f)\n",
                "\n",
                "        res = ezkl.swap_proof_commitments(proof_path, witness_path)\n",
                "        print(res)\n",
                "        \n",
                "        # load proof and then print \n",
                "        proof = json.load(open(proof_path, 'r'))\n",
                "        res_2_proof = proof[\"hex_proof\"]\n",
                "        # show diff in hex strings\n",
                "        print(res_1_proof)\n",
                "        print(res_2_proof)\n",
                "        assert res_1_proof == res_2_proof\n",
                "\n",
                "    res = ezkl.verify(\n",
                "            proof_path,\n",
                "            settings_path,\n",
                "            vk_path,\n",
                "        )\n",
                "\n",
                "    assert res == True\n",
                "    print(\"verified\")\n",
                "\n",
                "for i in range(2):\n",
                "    prove_model(i)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "###  KZG commitment intermediate calculations\n",
                "\n",
                "This time the visibility parameters are:\n",
                "- `input_visibility`: \"polycommit\"\n",
                "- `param_visibility`: \"public\"\n",
                "- `output_visibility`: polycommit"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import ezkl\n",
                "\n",
                "run_args = ezkl.PyRunArgs()\n",
                "run_args.input_visibility = \"polycommit\"\n",
                "run_args.param_visibility = \"fixed\"\n",
                "run_args.output_visibility = \"polycommit\"\n",
                "run_args.variables = [(\"batch_size\", 1)]\n",
                "run_args.input_scale = 2\n",
                "run_args.logrows = 8\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "for i in range(2):\n",
                "    await setup(i)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "for i in range(2):\n",
                "    prove_model(i)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "You can also mock aggregate the split proofs into a single proof. This is useful if you want to verify the proof on chain at a lower cost. Here we mock aggregate the proofs to save time. You can use other notebooks to see how to aggregate in full ! "
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# now mock aggregate the proofs\n",
                "proofs = []\n",
                "for i in range(2):\n",
                "    proof_path = os.path.join('proof_split_'+str(i)+'.json')\n",
                "    proofs.append(proof_path)\n",
                "\n",
                "ezkl.mock_aggregate(proofs, logrows=22, split_proofs = True)"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "ezkl",
            "language": "python",
            "name": "python3"
        },
        "language_info": {
            "codemirror_mode": {
                "name": "ipython",
                "version": 3
            },
            "file_extension": ".py",
            "mimetype": "text/x-python",
            "name": "python",
            "nbconvert_exporter": "python",
            "pygments_lexer": "ipython3",
            "version": "3.12.2"
        },
        "orig_nbformat": 4
    },
    "nbformat": 4,
    "nbformat_minor": 2
}