examples/notebooks/mnist_vae.ipynb

{
    "cells": [
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# what is the variational?"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Credits to [geohot](https://github.com/geohot/ai-notebooks/blob/master/mnist_gan.ipynb) for most of this code"
            ]
        },
        {
            "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\", \"tf2onnx\"])\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",
                "\n",
                "import os\n",
                "import time\n",
                "import random\n",
                "\n",
                "import tensorflow as tf\n",
                "import tensorflow.keras.backend as K\n",
                "from tensorflow.keras.optimizers import Adam\n",
                "from tensorflow.keras.layers import *\n",
                "from tensorflow.keras.models import Model\n",
                "from tensorflow.keras.losses import mse\n",
                "from tensorflow.keras.datasets import mnist\n",
                "(x_train, y_train), (x_test, y_test) = mnist.load_data()\n",
                "x_train, x_test = [x/255.0 for x in [x_train, x_test]]\n",
                "y_train, y_test = [tf.keras.utils.to_categorical(x) for x in [y_train, y_test]]"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "ZDIM = 4\n",
                "\n",
                "def get_encoder():\n",
                "  x = in1 = Input((28,28))\n",
                "  x = Reshape((28,28,1))(x)\n",
                "\n",
                "  x = Conv2D(64, (5,5), padding='same', strides=(2,2))(x)\n",
                "  x = BatchNormalization()(x)\n",
                "  x = ELU()(x)\n",
                "\n",
                "  x = Conv2D(128, (5,5), padding='same', strides=(2,2))(x)\n",
                "  x = BatchNormalization()(x)\n",
                "  x = ELU()(x)\n",
                "\n",
                "  x = Flatten()(x)\n",
                "  x = Dense(ZDIM)(x)\n",
                "  return Model(in1, x)\n",
                "\n",
                "def get_decoder():\n",
                "  x = in1 = Input((ZDIM,))\n",
                "\n",
                "  x = Dense(7*7*64)(x)\n",
                "  x = BatchNormalization()(x)\n",
                "  x = ELU()(x)\n",
                "  x = Reshape((7,7,64))(x)\n",
                "\n",
                "  x = Conv2DTranspose(128, (5,5), strides=(2,2), padding='same')(x)\n",
                "  x = BatchNormalization()(x)\n",
                "  x = ELU()(x)\n",
                "\n",
                "  x = Conv2DTranspose(1, (5,5), strides=(2,2), padding='same')(x)\n",
                "  x = Activation('sigmoid')(x)\n",
                "  x = Reshape((28,28))(x)\n",
                "  return Model(in1, x)"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "### Regular Autoencoder"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# normal autoencoder without the variational\n",
                "enc = get_encoder()\n",
                "dec = get_decoder()\n",
                "ae = Model(enc.input, dec(enc.output))\n",
                "ae.compile('adam', 'mse')\n",
                "ae.summary()\n",
                "# make the epochs larger for better results\n",
                "ae.fit(x_train, x_train, batch_size=128, epochs=1, shuffle=1, validation_data=(x_test, x_test))"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# while the autoencoder might work without the variational, the sampling doesn't\n",
                "import numpy as np\n",
                "from matplotlib.pyplot import figure, imshow\n",
                "imshow(np.concatenate(ae.predict(np.array([random.choice(x_test) for i in range(10)])), axis=1))\n",
                "figure(figsize=(16,16))\n",
                "imshow(np.concatenate(ae.layers[-1].predict(np.random.normal(size=(10, ZDIM))), axis=1))"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import os \n",
                "\n",
                "model_path = os.path.join('ae.onnx')\n",
                "compiled_model_path = os.path.join('ae.compiled')\n",
                "pk_path = os.path.join('ae.pk')\n",
                "vk_path = os.path.join('ae.vk')\n",
                "settings_path = os.path.join('ae_settings.json')\n",
                "srs_path = os.path.join('ae_kzg.srs')\n",
                "witness_path = os.path.join('ae_witness.json')\n",
                "data_path = os.path.join('ae_input.json')"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "Now we export the decoder (which presumably is what we want) -- to onnx"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "\n",
                "import numpy as np\n",
                "import tf2onnx\n",
                "import tensorflow as tf\n",
                "import json\n",
                "\n",
                "shape = [1, ZDIM]\n",
                "# After training, export to onnx (network.onnx) and create a data file (input.json)\n",
                "x = 0.1*np.random.rand(1,*shape)\n",
                "\n",
                "spec = tf.TensorSpec(shape, tf.float32, name='input_0')\n",
                "\n",
                "\n",
                "tf2onnx.convert.from_keras(dec, input_signature=[spec], inputs_as_nchw=['input_0'], opset=12, output_path=model_path)\n",
                "\n",
                "data_array = x.reshape([-1]).tolist()\n",
                "\n",
                "data = dict(input_data = [data_array])\n",
                "\n",
                "    # Serialize data into file:\n",
                "json.dump( data, open(data_path, 'w' ))\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import ezkl\n",
                "\n",
                "!RUST_LOG=trace\n",
                "res = ezkl.gen_settings(model_path, settings_path)\n",
                "assert res == True\n"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "cal_path = os.path.join(\"calibration.json\")\n",
                "\n",
                "data_array = (0.1 * np.random.rand(20, *shape)).reshape([-1]).tolist()\n",
                "\n",
                "data = dict(input_data = [data_array])\n",
                "\n",
                "# Serialize data into file:\n",
                "json.dump(data, open(cal_path, 'w'))\n",
                "\n",
                "\n",
                "await ezkl.calibrate_settings(cal_path, model_path, settings_path, \"resources\")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "res = ezkl.compile_circuit(model_path, compiled_model_path, settings_path)\n",
                "assert res == True"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# srs path\n",
                "res = await ezkl.get_srs( settings_path)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# now generate the witness file\n",
                "witness_path = \"ae_witness.json\"\n",
                "\n",
                "res = await ezkl.gen_witness(data_path, compiled_model_path, witness_path)\n",
                "assert os.path.isfile(witness_path)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "res = ezkl.mock(witness_path, compiled_model_path)\n",
                "assert res == True"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "\n",
                "# HERE WE SETUP THE CIRCUIT PARAMS\n",
                "# WE GOT KEYS\n",
                "# WE GOT CIRCUIT PARAMETERS\n",
                "# EVERYTHING ANYONE HAS EVER NEEDED FOR ZK\n",
                "\n",
                "res = ezkl.setup(\n",
                "        compiled_model_path,\n",
                "        vk_path,\n",
                "        pk_path,\n",
                "        \n",
                "    )\n",
                "\n",
                "assert res == True\n",
                "assert os.path.isfile(vk_path)\n",
                "assert os.path.isfile(pk_path)\n",
                "assert os.path.isfile(settings_path)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# GENERATE A PROOF\n",
                "\n",
                "\n",
                "proof_path = os.path.join('ae.pf')\n",
                "\n",
                "res = ezkl.prove(\n",
                "        witness_path,\n",
                "        compiled_model_path,\n",
                "        pk_path,\n",
                "        proof_path,\n",
                "        \n",
                "        \"single\",\n",
                "    )\n",
                "\n",
                "print(res)\n",
                "assert os.path.isfile(proof_path)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# VERIFY IT\n",
                "res = ezkl.verify(\n",
                "        proof_path,\n",
                "        settings_path,\n",
                "        vk_path,\n",
                "        \n",
                "    )\n",
                "\n",
                "assert res == True\n",
                "print(\"verified\")"
            ]
        },
        {
            "cell_type": "markdown",
            "metadata": {},
            "source": [
                "### Variational Autoencoder"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "in1 = Input((28,28))\n",
                "x = get_encoder()(in1)\n",
                "\n",
                "# add the variational\n",
                "z_mu = Dense(ZDIM)(x)\n",
                "z_log_var = Dense(ZDIM)(x)\n",
                "z = Lambda(lambda x: x[0] + K.exp(0.5 * x[1]) * K.random_normal(shape=K.shape(x[0])))([z_mu, z_log_var])\n",
                "dec = get_decoder()\n",
                "dec.output_names=['output']\n",
                "\n",
                "out = dec(z)\n",
                "\n",
                "mse_loss = mse(Reshape((28*28,))(in1), Reshape((28*28,))(out)) * 28 * 28\n",
                "kl_loss = 1 + z_log_var - K.square(z_mu) - K.exp(z_log_var)\n",
                "kl_loss = -0.5 * K.mean(kl_loss, axis=-1)\n",
                "\n",
                "vae = Model(in1, out)\n",
                "vae.add_loss(K.mean(mse_loss + kl_loss))\n",
                "vae.compile('adam')"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# z is sampled from z_mu and z_log_var with gaussian noise\n",
                "test = Model(in1, [z, z_mu, z_log_var])\n",
                "test.predict(x_train[0:1])"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "vae.fit(x_train, batch_size=128, epochs=1, shuffle=1, validation_data=(x_test, None))"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "imshow(np.concatenate(vae.predict(np.array([random.choice(x_test) for i in range(10)])), axis=1))\n",
                "figure(figsize=(16,16))\n",
                "imshow(np.concatenate(vae.layers[5].predict(np.random.normal(size=(10, ZDIM))), axis=1))"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import os \n",
                "\n",
                "model_path = os.path.join('vae.onnx')\n",
                "compiled_model_path = os.path.join('vae.compiled')\n",
                "pk_path = os.path.join('vae.pk')\n",
                "vk_path = os.path.join('vae.vk')\n",
                "settings_path = os.path.join('vae_settings.json')\n",
                "srs_path = os.path.join('vae_kzg.srs')\n",
                "witness_path = os.path.join('vae_witness.json')\n",
                "data_path = os.path.join('vae_input.json')"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "\n",
                "import numpy as np\n",
                "import tf2onnx\n",
                "import tensorflow as tf\n",
                "import json\n",
                "\n",
                "# After training, export to onnx (network.onnx) and create a data file (input.json)\n",
                "x = 0.1*np.random.rand(1,*[1, ZDIM])\n",
                "\n",
                "spec = tf.TensorSpec([1, ZDIM], tf.float32, name='input_0')\n",
                "\n",
                "\n",
                "tf2onnx.convert.from_keras(dec, input_signature=[spec], inputs_as_nchw=['input_0'], opset=12, output_path=model_path)\n",
                "\n",
                "data_array = x.reshape([-1]).tolist()\n",
                "\n",
                "data = dict(input_data = [data_array])\n",
                "\n",
                "    # Serialize data into file:\n",
                "json.dump( data, open(data_path, 'w' ))"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "import ezkl\n",
                "\n",
                "!RUST_LOG=trace\n",
                "res = ezkl.gen_settings(model_path, settings_path)\n",
                "assert res == True\n",
                "\n",
                "res = await ezkl.calibrate_settings(data_path, model_path, settings_path, \"resources\")\n",
                "assert res == True\n",
                "print(\"verified\")"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "res = ezkl.compile_circuit(model_path, compiled_model_path, settings_path)\n",
                "assert res == True"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# srs path\n",
                "res = await ezkl.get_srs( settings_path)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# now generate the witness file \n",
                "witness_path = \"vae_witness.json\"\n",
                "\n",
                "res = await ezkl.gen_witness(data_path, compiled_model_path, witness_path)\n",
                "assert os.path.isfile(witness_path)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# uncomment to mock prove\n",
                "# res = ezkl.mock(witness_path, compiled_model_path)\n",
                "# assert res == True"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "\n",
                "# HERE WE SETUP THE CIRCUIT PARAMS\n",
                "# WE GOT KEYS\n",
                "# WE GOT CIRCUIT PARAMETERS\n",
                "# EVERYTHING ANYONE HAS EVER NEEDED FOR ZK\n",
                "\n",
                "res = ezkl.setup(\n",
                "        compiled_model_path,\n",
                "        vk_path,\n",
                "        pk_path,\n",
                "        \n",
                "    )\n",
                "\n",
                "\n",
                "assert res == True\n",
                "assert os.path.isfile(vk_path)\n",
                "assert os.path.isfile(pk_path)\n",
                "assert os.path.isfile(settings_path)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# GENERATE A PROOF\n",
                "\n",
                "\n",
                "proof_path = os.path.join('test.pf')\n",
                "\n",
                "res = ezkl.prove(\n",
                "        witness_path,\n",
                "        compiled_model_path,\n",
                "        pk_path,\n",
                "        proof_path,\n",
                "        \n",
                "        \"single\",\n",
                "    )\n",
                "\n",
                "print(res)\n",
                "assert os.path.isfile(proof_path)"
            ]
        },
        {
            "cell_type": "code",
            "execution_count": null,
            "metadata": {},
            "outputs": [],
            "source": [
                "# VERIFY IT\n",
                "res = ezkl.verify(\n",
                "        proof_path,\n",
                "        settings_path,\n",
                "        vk_path,\n",
                "        \n",
                "    )\n",
                "\n",
                "assert res == True\n",
                "print(\"verified\")"
            ]
        }
    ],
    "metadata": {
        "kernelspec": {
            "display_name": "Python 3",
            "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.9.15"
        }
    },
    "nbformat": 4,
    "nbformat_minor": 2
}