Testbench_KMNIST_4_class.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oaYVuKnhWLB9"
   },
   "outputs": [],
   "source": [
    "device = \"cpu\"\n",
    "import time,os,copy,torch,torchvision\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "\n",
    "from torchvision import datasets, transforms\n",
    "from tqdm import tqdm\n",
    "import pickle\n",
    "# Pennylane\n",
    "import pennylane as qml\n",
    "from pennylane import numpy as np\n",
    "seed = 42\n",
    "torch.manual_seed(seed)\n",
    "np.random.seed(seed)\n",
    "\n",
    "# Plotting\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# OpenMP: number of parallel threads.\n",
    "os.environ[\"OMP_NUM_THREADS\"] = \"16\"\n",
    "dataset = \"KMNIST_4\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "eRH52GXpbdPd",
    "outputId": "2694d772-d65c-488b-e4d0-07b28eeee835"
   },
   "outputs": [],
   "source": [
    "n_qubits = 6             # Number of qubits\n",
    "nqubits=n_qubits\n",
    "q_depth = 2            # Depth of the quantum circuit (number of variational layers)\n",
    "q_delta = 0.01              # Initial spread of random quantum weights\n",
    "tensor_length = n_qubits*(n_qubits-1)*q_depth+n_qubits*(q_depth-1)\n",
    "print(tensor_length)\n",
    "\n",
    "def H_layer(nqubits):\n",
    "    for idx in range(nqubits):\n",
    "        qml.Hadamard(wires=idx)\n",
    "\n",
    "def RZ_layer(w):\n",
    "    for idx, element in enumerate(w):\n",
    "        qml.RZ(element, wires=idx)\n",
    "\n",
    "def entangling_layer(nqubits,weights):\n",
    "    p = nqubits\n",
    "    weights_ = (weight for weight in weights)\n",
    "    for i in range(1,nqubits):\n",
    "        for j in range(i):\n",
    "            qml.CNOT(wires=[j,i])\n",
    "            param = next(weights_)\n",
    "            # vqc_params.append(param)\n",
    "            qml.RZ(param, wires=i)\n",
    "            p+=1\n",
    "            param = next(weights_)\n",
    "            # vqc_params.append(param)\n",
    "            qml.RX(param, wires=i)\n",
    "            p+=1\n",
    "            qml.CNOT(wires=[j,i])\n",
    "\n",
    "dev = qml.device(\"default.qubit\", wires=n_qubits)\n",
    "@qml.qnode(dev, interface=\"torch\")\n",
    "def quantum_net(q_input_features, q_weights_flat):\n",
    "\n",
    "    q_weights = q_weights_flat\n",
    "    H_layer(n_qubits)\n",
    "    RZ_layer(q_input_features)\n",
    "    entangling_layer(nqubits,q_weights)\n",
    "    exp_vals = [qml.expval(qml.PauliZ(position)) for position in range(n_qubits)]\n",
    "\n",
    "    return tuple(exp_vals)\n",
    "    \n",
    "class DressedQuantumNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.pre_net = nn.Linear(128, n_qubits)\n",
    "        self.q_params = nn.Parameter(q_delta * torch.randn(tensor_length))\n",
    "        self.post_net = nn.Linear(n_qubits, 5)\n",
    "\n",
    "    def forward(self, input_features):\n",
    "        pre_out = self.pre_net(input_features)\n",
    "        q_in = torch.tanh(pre_out) * np.pi / 2.0\n",
    "        q_out = torch.Tensor(0, n_qubits)\n",
    "        q_out = q_out.to(device)\n",
    "        for elem in q_in:\n",
    "            q_out_elem = quantum_net(elem, self.q_params).float().unsqueeze(0)\n",
    "            q_out = torch.cat((q_out, q_out_elem))\n",
    "        return self.post_net(q_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class DressedClassicalNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "        self.pre_net = nn.Linear(128, 6)\n",
    "        self.mid_net = nn.Linear(6,6)\n",
    "        self.post_net = nn.Linear(6, 5)\n",
    "\n",
    "    def forward(self, input_features):\n",
    "        pre_out = self.pre_net(input_features)\n",
    "        q_in = torch.tanh(pre_out) * np.pi / 2.0\n",
    "        mid_out = self.mid_net(q_in)\n",
    "        return self.post_net(mid_out)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 785
    },
    "id": "5iOn-WphsNw8",
    "outputId": "284cd43a-ad56-4e60-ba28-1965db2e32db"
   },
   "outputs": [],
   "source": [
    "class_names = [\n",
    "    \"Shirt\",\n",
    "    \"Sneaker\",\n",
    "    \"Bag\",\n",
    "    \"Ankle boot\",\n",
    "]\n",
    "num_classes = len(class_names)\n",
    "train_images = np.load(f\"./saved_data_{dataset}/train_seqs_{seed}.npz\")['arr_0']\n",
    "train_labels = np.load(f\"./saved_data_{dataset}/train_labels.npz\")['arr_0']\n",
    "plt.figure(figsize=(10, 10))\n",
    "for i in range(25):\n",
    "    plt.subplot(5, 5, i + 1)\n",
    "    plt.imshow(train_images[i][0], cmap=plt.cm.binary)\n",
    "    plt.axis(\"off\")\n",
    "    plt.title(class_names[train_labels[i]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "DUbi9Zk6ZEN8"
   },
   "source": [
    "We make an array to store all the models loaded from h5 files and other arrays named after the type of noise used which are populated by the performances of the models."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Model:\n",
    "    def __init__(self,path,desc):\n",
    "        self.path = path\n",
    "        self.desc = desc\n",
    "        self.gauss = {'0.01':0.,'0.02':0.,'0.03':0.,'0.04':0.,'0.05':0.,'0.06':0.,'0.07':0.,'0.08':0.,'0.09':0.,'0.1':0.,'0.2':0.,'0.3':0.,'0.4':0.,'0.5':0.,'0.6':0.,'0.7':0.,'0.8':0.,'0.9':0.}\n",
    "        self.snp = {'0.01':0.,'0.02':0.,'0.03':0.,'0.04':0.,'0.05':0.,'0.06':0.,'0.07':0.,'0.08':0.,'0.09':0.,'0.1':0.,'0.2':0.,'0.3':0.,'0.4':0.,'0.5':0.}\n",
    "        self.uniform = {'0.01':0.,'0.02':0.,'0.03':0.,'0.04':0.,'0.05':0.,'0.06':0.,'0.07':0.,'0.08':0.,'0.09':0.,'0.1':0.,'0.2':0.,'0.3':0.,'0.4':0.,'0.5':0.,'0.6':0.,'0.7':0.,'0.8':0.,'0.9':0.}\n",
    "        self.rayleigh ={'0.01':0.,'0.02':0.,'0.03':0.,'0.04':0.,'0.05':0.,'0.06':0.,'0.07':0.,'0.08':0.,'0.09':0.,'0.1':0.,'0.2':0.,'0.3':0.,'0.4':0.,'0.5':0.,'0.6':0.,'0.7':0.,'0.8':0.,'0.9':0.}\n",
    "        self.perlin = {'1x1':0.,'7x7':0.,'14x14':0.}\n",
    "\n",
    "    def dump(self):\n",
    "        np.savez_compressed(file=self.desc,gauss=self.gauss,snp=self.snp,uniform=self.uniform,rayleigh=self.rayleigh,perlin=self.perlin)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "qGFIXLfzCukJ",
    "outputId": "854abcee-06c6-48a7-c38b-84e070609afb"
   },
   "outputs": [],
   "source": [
    "modelnames = {\n",
    "    \"./KMNIST/random_kmnist_final.h5\":\"classical for the above\"\n",
    "}\n",
    "# \"path_to_file.h5\" : \"short description of this model\"\n",
    "models = [Model(k,v) for k,v in modelnames.items()]\n",
    "for model in models:\n",
    "    print(model.path,model.desc)\n",
    "    print('-------------------------- '*2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_labels = np.load(f\"./saved_data_{dataset}/test_labels.npz\")['arr_0']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def test(input_model,test_x,key):\n",
    "    minibatch_size = 32\n",
    "    optimizer = torch.optim.Adam(input_model.parameters())\n",
    "    input_model.eval()\n",
    "    test_acc = 0\n",
    "    for i in range(test_x.shape[0] // minibatch_size):\n",
    "        batch_in = test_x[i * minibatch_size : (i + 1) * minibatch_size]\n",
    "        batch_in = batch_in.reshape((-1,) + test_x.shape[1:-2] + (784,))\n",
    "        batch_label = test_labels[i * minibatch_size : (i + 1) * minibatch_size]\n",
    "        output = input_model(torch.tensor(batch_in))\n",
    "\n",
    "        test_acc += torch.mean(\n",
    "            torch.eq(torch.argmax(output, dim=1), torch.tensor(batch_label)).float()\n",
    "        )\n",
    "\n",
    "    test_acc /= i + 1\n",
    "    print(f\"Test accuracy for {key}:\", test_acc.numpy())\n",
    "    return test_acc.numpy()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ZhhDI8qczWjQ"
   },
   "source": [
    "Here we try to add salt&pepper noise to the test_sequences and then we test our model on them."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tRYGLE_fEouB"
   },
   "outputs": [],
   "source": [
    "for model in models:\n",
    "    print(f\"Model = {model.desc} at {model.path}\")\n",
    "    obj = torch.load(model.path)\n",
    "    print(\"training above model for salt and pepper noise\")\n",
    "    for i in tqdm([i/100 for i in range(1,10)]+[i/10 for i in range(1,6)]):\n",
    "         model.snp[str(i)] = test(obj,np.load(f\"./saved_data_{dataset}/snp_{i}_{seed}.npz\")['arr_0'],i)\n",
    "    print(\"training above model for gaussian noise\")\n",
    "    for i in tqdm([i/100 for i in range(1,10)]+[i/10 for i in range(1,10)]):\n",
    "        model.gauss[str(i)] = test(obj,np.load(f\"./saved_data_{dataset}/gauss_{i}_{seed}.npz\")['arr_0'],i)\n",
    "    print(\"training above model for rayleigh noise\")\n",
    "    for i in tqdm([i/100 for i in range(1,10)]+[i/10 for i in range(1,10)]):\n",
    "        model.rayleigh[str(i)] = test(obj,np.load(f\"./saved_data_{dataset}/rayl_{i}_{seed}.npz\")['arr_0'],i)\n",
    "    print(\"training above model for uniform noise\")\n",
    "    for i in tqdm([i/100 for i in range(1,10)]+[i/10 for i in range(1,10)]):\n",
    "        model.uniform[str(i)] = test(obj,np.load(f\"./saved_data_{dataset}/uniform_0_{i}_{seed}.npz\")['arr_0'],i)\n",
    "    print(\"training above model for perlin noise\")\n",
    "    for i in tqdm(['1','7','14']):\n",
    "        model.perlin[str(i)+'x'+str(i)] = test(obj,np.load(f\"./saved_data_{dataset}/perlin_{i}_x_{i}_{seed}.npz\")['arr_0'],i)\n",
    "    model.dump()"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Testbench",
   "provenance": []
  },
  "interpreter": {
   "hash": "63fd5069d213b44bf678585dea6b12cceca9941eaf7f819626cde1f2670de90d"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.8.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}