[5_ary]_pretrain_5x2_squeeze_MNIST.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "oaYVuKnhWLB9"
   },
   "outputs": [],
   "source": [
    "device = \"cpu\"\n",
    "import time,os,copy,torch,pytorch_spiking,torchvision\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\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",
    "\n",
    "torch.manual_seed(42)\n",
    "np.random.seed(42)\n",
    "\n",
    "# Plotting\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# OpenMP: number of parallel threads.\n",
    "os.environ[\"OMP_NUM_THREADS\"] = \"16\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "VNDrzJlUTbv1"
   },
   "outputs": [],
   "source": [
    "spikeaware_model = torch.load(\"./MNIST/spikeaware_mnist.h5\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 164
    },
    "id": "fok2C-ecl4pV",
    "outputId": "cd22c63d-7c73-4e83-e1cd-a85250c4cab5"
   },
   "outputs": [],
   "source": [
    "display(spikeaware_model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "THJBENVKRYSJ"
   },
   "outputs": [],
   "source": [
    "n_qubits = 5           # Number of qubits\n",
    "nqubits=n_qubits\n",
    "batch_size = 4              # Number of samples for each training step\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",
    "start_time = time.time()    # Start of the computation timer\n",
    "tensor_length = n_qubits*(n_qubits-1)*q_depth+n_qubits"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "XYFDgRZ3YNcv"
   },
   "outputs": [],
   "source": [
    "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])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "AoUoGHoFV8rI"
   },
   "outputs": [],
   "source": [
    "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",
    "    \n",
    "    H_layer(n_qubits)\n",
    "\n",
    "    # Embed features in the quantum node\n",
    "    RZ_layer(q_input_features)\n",
    "\n",
    "    entangling_layer(nqubits,q_weights)\n",
    "    # Expectation values in the Z basis\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",
    "\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",
    "\n",
    "        return self.post_net(q_out)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Z9oe41C-7NiB"
   },
   "outputs": [],
   "source": [
    "i = 1\n",
    "if i==1:\n",
    "  draw = qml.draw(quantum_net)\n",
    "  draw([0,0,0,0,0],np.random.rand(tensor_length))\n",
    "  i+=1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "qmdwghpo7p_A",
    "outputId": "704cb71a-d7bc-4a33-c8dd-15494055fe23"
   },
   "outputs": [],
   "source": [
    "print(qml.draw(quantum_net)([0,0,0,0,0],np.random.rand(tensor_length)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "fqahBRmc6BzQ",
    "outputId": "9a65d81a-818b-43f9-f6a3-aaf56680712d"
   },
   "outputs": [],
   "source": [
    "quantum_net([0,0,0,0,0],np.random.rand(tensor_length))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Oe15V30yVwln"
   },
   "outputs": [],
   "source": [
    "model_hybrid = copy.deepcopy(torch.nn.Sequential(*(list(spikeaware_model.children())[:-1])))\n",
    "for param in model_hybrid.parameters():\n",
    "    param.requires_grad = False\n",
    "\n",
    "model_hybrid.fc = DressedQuantumNet()\n",
    "\n",
    "model_hybrid = model_hybrid.to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 219
    },
    "id": "4GrawGKtVFBV",
    "outputId": "5c8e00e5-4af6-494d-919b-bc65153291b9"
   },
   "outputs": [],
   "source": [
    "display(model_hybrid)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "BBqhpzaiuuWC"
   },
   "outputs": [],
   "source": [
    "\n",
    "class_names = [\n",
    "    \"5\",\n",
    "    \"6\",\n",
    "    \"7\",\n",
    "    \"8\",\n",
    "    \"9\",\n",
    "]\n",
    "num_classes = len(class_names)\n",
    "\n",
    "import idx2numpy \n",
    "import numpy as np \n",
    "arr3 = []\n",
    "arr4 = []\n",
    "file1 = './MNIST/TorchvisionDatasetWrapper/raw/train-images-idx3-ubyte' \n",
    "arr = idx2numpy.convert_from_file(file1)\n",
    "file2 = './MNIST/TorchvisionDatasetWrapper/raw/train-labels-idx1-ubyte' \n",
    "arr2 = idx2numpy.convert_from_file(file2)\n",
    "for i in range(len(arr)):\n",
    "    if arr2[i] in [5,6,7,8,9]:\n",
    "        arr3.append(arr[i])\n",
    "        arr4.append(arr2[i])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "k8WsgkKno38A",
    "outputId": "0d26035c-535f-4407-d22e-30401ce1fc36"
   },
   "outputs": [],
   "source": [
    "len(arr4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "9XrbmsLB_sbt"
   },
   "outputs": [],
   "source": [
    "# repeat the images for n_steps\n",
    "n_steps = 10\n",
    "\n",
    "train_images = np.asarray(arr3,dtype=np.float32) / 255.0\n",
    "train_labels = np.asarray(arr4,dtype=np.int64)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 35
    },
    "id": "NTmQ0jC-10dq",
    "outputId": "31e079bc-d486-4907-956d-3dbf74301501"
   },
   "outputs": [],
   "source": [
    "train_labels -= 5\n",
    "display(train_labels)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 785
    },
    "id": "5iOn-WphsNw8",
    "outputId": "ed68c026-5d0d-4d9a-d27b-778d664ec173"
   },
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10, 10))\n",
    "for i in range(25):\n",
    "    plt.subplot(5, 5, i + 1)\n",
    "    plt.imshow(train_images[i], cmap=plt.cm.binary)\n",
    "    plt.axis(\"off\")\n",
    "    plt.title(class_names[train_labels[i]])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "f-WhMWC7pHOf",
    "outputId": "9f46dcec-4c0c-4520-d650-dc0780ce575e"
   },
   "outputs": [],
   "source": [
    "\n",
    "split = np.array_split(train_images,[int(len(train_images)*0.7)])\n",
    "train_sequences = split[0]\n",
    "test_sequences = split[1]\n",
    "print(len(split[0]),len(split[1]))\n",
    "split = np.array_split(train_labels,[int(len(train_labels)*0.7)])\n",
    "train_labels = split[0]\n",
    "test_labels = split[1]\n",
    "print(len(split[0]),len(split[1]))\n",
    "datalength = 5000\n",
    "\n",
    "train_sequences = np.tile(train_sequences[:, None], (1, n_steps, 1, 1))[:datalength]\n",
    "test_sequences = np.tile(test_sequences[:, None], (1, n_steps, 1, 1))\n",
    "train_labels = train_labels[:datalength]\n",
    "test_labels = test_labels\n",
    "\n",
    "print(len(test_labels),len(test_sequences),len(train_labels),len(train_sequences))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "metric = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "QfTPVzsst41s"
   },
   "outputs": [],
   "source": [
    "def train(input_model, train_x, test_x):\n",
    "    minibatch_size = 32\n",
    "    optimizer = torch.optim.Adam(input_model.parameters())\n",
    " \n",
    "    input_model.train()\n",
    "    for j in range(100):\n",
    "        train_acc = 0\n",
    "        loss_acc = 0\n",
    "        for i in tqdm(range(train_x.shape[0] // minibatch_size)):\n",
    "            input_model.zero_grad()\n",
    " \n",
    "            batch_in = train_x[i * minibatch_size : (i + 1) * minibatch_size]\n",
    "            # flatten images\n",
    "            batch_in = batch_in.reshape((-1,) + train_x.shape[1:-2] + (784,))\n",
    "            batch_label = train_labels[i * minibatch_size : (i + 1) * minibatch_size]\n",
    "            output = input_model(torch.tensor(batch_in))\n",
    " \n",
    "            # compute sparse categorical cross entropy loss\n",
    "            logp = torch.nn.functional.log_softmax(output, dim=-1)\n",
    "            logpy = torch.gather(logp, 1, torch.tensor(batch_label).view(-1, 1))\n",
    "            loss = -logpy.mean()\n",
    "            loss_acc += loss \n",
    " \n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    " \n",
    "            train_acc += torch.mean(\n",
    "                torch.eq(torch.argmax(output, dim=1), torch.tensor(batch_label)).float()\n",
    "            )\n",
    "        train_acc /= i + 1  \n",
    "        print(\"Train accuracy (%d): \" % j, train_acc.numpy())\n",
    "        #metric_train=train_acc.numpy()\n",
    "\n",
    "        loss_acc /=i + 1\n",
    "        print(\"Loss (%d): \" % j, loss_acc.detach().numpy())\n",
    "        torch.save(model_hybrid,f=\"./model_saves/4_class_MNIST/6x2_zz_{0}_{1}_{2}.h5\".format(j,train_acc.numpy(),loss_acc.detach().numpy()))\n",
    "        #metric_loss= loss_acc.detach().numpy()\n",
    "        #metric.append([metric_train,metric_loss])\n",
    "        #np.save(\"metrics.npy\",np.array(metric))\n",
    " \n",
    "    # compute test accuracy\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",
    " \n",
    "    print(\"Test accuracy:\", test_acc.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "huwxaiS52-NG"
   },
   "outputs": [],
   "source": [
    "def test(input_model, test_x):\n",
    "    minibatch_size = 32\n",
    "    optimizer = torch.optim.Adam(input_model.parameters())\n",
    "\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",
    "\n",
    "    print(\"Test accuracy:\", test_acc.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "TwU0pUFe0JAg"
   },
   "outputs": [],
   "source": [
    "model_hybrid=torch.load(\"./model_saves/5_class_MNIST/5x2_zz_0_0.9458734035491943_1.3720676898956299.h5\")"
   ]
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "[5-ary] pretrain 4x1 squeeze.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3.9.12 64-bit",
   "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.12"
  },
  "vscode": {
   "interpreter": {
    "hash": "beaffcd69350167707082d796d83130d58bab3408b4f65f2e212ecd34479db88"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}