associative_memory.ipynb

{
 "metadata": {
  "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
  },
  "orig_nbformat": 2,
  "kernelspec": {
   "name": "python_defaultSpec_1596038169325",
   "display_name": "Python 3.6.8 64-bit ('pytorch_env': virtualenv)"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2,
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import random\n",
    "import matplotlib.pyplot as plt\n",
    "import importlib\n",
    "import matplotlib\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "from torchvision import datasets\n",
    "import sys\n",
    "import os\n",
    "import snn_lib.utilities as utilities\n",
    "from snn_lib.snn_layers import *"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Configurations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset_path = './dataset/associative_target.npy'\n",
    "\n",
    "checkpoint_base_name = \"associative_memory_checkpoint_\"\n",
    "checkpoint_base_path = './associative_memory_checkpoint/'\n",
    "\n",
    "torch.manual_seed(2)\n",
    "np.random.seed(0)\n",
    "\n",
    "length = 300\n",
    "input_size = 300\n",
    "batch_size = 64\n",
    "synapse_type = 'dual_exp'\n",
    "epoch = 50\n",
    "tau_m = 8\n",
    "tau_s = 2\n",
    "filter_tau_m = tau_m\n",
    "filter_tau_s = tau_s\n",
    "dropout_rate = 0.3\n",
    "\n",
    "modify_input_prob = 0.6\n",
    "remove_prob = 0.5\n",
    "remove_row_prob = 0.2\n",
    "remove_col_prob = 0.2\n",
    "remove_block_prob = 0.2\n",
    "mutate_prob = 0.7\n",
    "noise_prob = 0.005\n",
    "\n",
    "remove_block_h = 30\n",
    "remove_block_w = 40\n",
    "remove_block_h = 40\n",
    "remove_col_w = 30\n",
    "remove_row_h = 30\n",
    "\n",
    "optimizer_choice = 1\n",
    "scheduler_choice = 1\n",
    "\n",
    "optimizer_config = {0: (torch.optim.Adam,       0.001),\n",
    "                    1: (torch.optim.AdamW,      0.001),\n",
    "                    2: (torch.optim.SGD,        0.0001)}\n",
    "\n",
    "scheduler_config = {0: None,\n",
    "                    1: (torch.optim.lr_scheduler.MultiStepLR, [50,100, 150], 0.1),\n",
    "                    # order: milestones, gamma=0.1\n",
    "                    2: (torch.optim.lr_scheduler.CosineAnnealingWarmRestarts, 1000),\n",
    "                    # T_0\n",
    "                    3: (torch.optim.lr_scheduler.CyclicLR, 0.001, 0.01, 200)\n",
    "                    }"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Prepare dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# %% utility functions\n",
    "def add_noise_spike(spike_array,probability = 0.003):\n",
    "    '''\n",
    "    :param spike_array: 2d array [spike train num, length]\n",
    "    :param probability:\n",
    "    :return:\n",
    "    '''\n",
    "    noise_mat = np.random.rand(*spike_array.shape).astype(np.float32)\n",
    "\n",
    "    noise_mat[np.where(noise_mat > 1-probability)] = 1\n",
    "    noise_mat[np.where(noise_mat <= 1 - probability)] = 0\n",
    "\n",
    "    new_arr = spike_array + noise_mat\n",
    "\n",
    "    new_arr[np.where(new_arr > 1)] = 1\n",
    "\n",
    "    return  new_arr\n",
    "\n",
    "def remove_row(spike_array,remove_width, position = None):\n",
    "    '''\n",
    "    remove a few rows in the spike mat (set rows to 0)\n",
    "    :param spike_array: 2d array [spike train num, length]\n",
    "    :param remove_width: How many rows to remove\n",
    "    :param position: spike will be removed from row position to position+remove_width\n",
    "    :return:\n",
    "    '''\n",
    "\n",
    "    h,w = spike_array.shape\n",
    "    upper_row = np.random.randint(0,h-remove_width)\n",
    "\n",
    "    if position != None:\n",
    "        upper_row = position\n",
    "\n",
    "    new_arr = spike_array\n",
    "    new_arr[upper_row:upper_row+remove_width,:] = 0\n",
    "\n",
    "    return new_arr\n",
    "\n",
    "def remove_col(spike_array,remove_width):\n",
    "    '''\n",
    "    remove a few columns in spike mat (set columns to 0)\n",
    "    :param spike_array: 2d array [spike train num, length]\n",
    "    :param width:\n",
    "    :return:\n",
    "    '''\n",
    "\n",
    "    h,w = spike_array.shape\n",
    "    left_col = np.random.randint(0,w-remove_width)\n",
    "\n",
    "    new_arr = spike_array\n",
    "    new_arr[:,left_col:left_col+remove_width] = 0\n",
    "\n",
    "    return new_arr\n",
    "\n",
    "def remove_block(spike_array,remove_hight, remove_width):\n",
    "    '''\n",
    "    set a block region in spike mat to 0\n",
    "    :param spike_array: 2d array [spike train num, length]\n",
    "    :param width:\n",
    "    :return:\n",
    "    '''\n",
    "\n",
    "    h,w = spike_array.shape\n",
    "    top_left_row = np.random.randint(0,w-remove_hight)\n",
    "    top_left_col = np.random.randint(0,w-remove_width)\n",
    "\n",
    "    new_arr = spike_array\n",
    "    new_arr[top_left_row:top_left_row+remove_hight,top_left_col:top_left_col+remove_width] = 0\n",
    "\n",
    "    return new_arr\n",
    "\n",
    "class PatternDataset(Dataset):\n",
    "    \"\"\"random pattern dataset\"\"\"\n",
    "\n",
    "    def __init__(self, input_pattern, target, filtered_target,length):\n",
    "        '''\n",
    "\n",
    "        :param input_pattern: [number, hight, width/time]\n",
    "        :param label_cat:\n",
    "        :param target:\n",
    "        :param length:\n",
    "        '''\n",
    "        self.input_pattern = input_pattern\n",
    "        self.target = target\n",
    "        self.filtered_target = filtered_target\n",
    "        self.length = length\n",
    "\n",
    "\n",
    "        self.modify_input_prob = modify_input_prob\n",
    "        self.remove_prob = remove_prob\n",
    "        self.remove_row_prob = remove_row_prob\n",
    "        self.remove_col_prob = remove_col_prob\n",
    "        self.noise_prob = noise_prob\n",
    "        self.mutate_prob = mutate_prob\n",
    "\n",
    "        self.remove_block_h = remove_block_h\n",
    "        self.remove_block_w = remove_block_w\n",
    "        self.remove_block_h =remove_block_h\n",
    "        self.remove_col_w = remove_col_w\n",
    "        self.remove_row_h = remove_row_h\n",
    "\n",
    "    def __len__(self):\n",
    "        #input only has 10 samples, so increase length by 100\n",
    "        #each sample is a class\n",
    "        return len(self.input_pattern)*100\n",
    "\n",
    "    def __getitem__(self, idx):\n",
    "        if torch.is_tensor(idx):\n",
    "            idx = idx.tolist()\n",
    "\n",
    "        #input data only has 10 sample\n",
    "        #by mod get the actual idx, which is also the class\n",
    "        real_idx = idx % 10\n",
    "\n",
    "        label = real_idx\n",
    "        input_pattern = self.input_pattern[real_idx].copy()\n",
    "        output_target = self.target[real_idx]\n",
    "        filtered_target = self.filtered_target[real_idx]\n",
    "\n",
    "        if np.random.rand() > self.modify_input_prob:\n",
    "            #modify input data, remove, delete, add noise, mutate\n",
    "            if np.random.rand() < self.mutate_prob:\n",
    "                input_pattern = utilities.mutate_spike_pattern(input_pattern, 0, 1)\n",
    "\n",
    "            #random select remove row, col, or block\n",
    "            choice = np.random.randint(2)\n",
    "\n",
    "            if np.random.rand() < self.remove_prob:\n",
    "                if choice == 0:\n",
    "                    input_pattern = remove_block(input_pattern,remove_block_w, remove_block_h)\n",
    "                elif choice == 1:\n",
    "                    input_pattern = remove_row(input_pattern, remove_block_h)\n",
    "\n",
    "        input_pattern = add_noise_spike(input_pattern, self.noise_prob)\n",
    "\n",
    "        return input_pattern, label, output_target, filtered_target"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Build the network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class mysnn(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        self.length = length\n",
    "        self.batch_size = batch_size\n",
    "\n",
    "        self.train_coefficients = True\n",
    "        self.train_decay_v = False\n",
    "        self.train_v0 = False\n",
    "        self.train_reset_v = False\n",
    "        self.train_reset_decay = False\n",
    "        self.train_threshold = False\n",
    "        self.train_bias = True\n",
    "        self.membrane_filter = False\n",
    "\n",
    "        self.axon1 = dual_exp_iir_layer((input_size,), self.length, \n",
    "        self.batch_size, tau_m, tau_s, self.train_coefficients)\n",
    "\n",
    "        self.dense1 = neuron_layer(input_size, 500, self.length, self.batch_size, \n",
    "        tau_m, self.train_bias, self.membrane_filter)\n",
    "\n",
    "        self.axon2 = dual_exp_iir_layer((500,), self.length, self.batch_size, tau_m, tau_s, \n",
    "        self.train_coefficients)\n",
    "        self.dense2 = neuron_layer(500, 200, self.length, self.batch_size, tau_m,  \n",
    "        self.train_bias, self.membrane_filter)\n",
    "                    \n",
    "        self.axon3 = dual_exp_iir_layer((200,), self.length, self.batch_size, tau_m, tau_s, \n",
    "        self.train_coefficients)\n",
    "\n",
    "        self.dense3 = neuron_layer(200, 500, self.length, self.batch_size, tau_m, \n",
    "        self.train_bias, self.membrane_filter)\n",
    "        \n",
    "        self.axon4 = dual_exp_iir_layer((500,), self.length, self.batch_size, tau_m, tau_s, \n",
    "        self.train_coefficients)\n",
    "\n",
    "        self.dense4 = neuron_layer(500, 300, self.length, self.batch_size, tau_m, \n",
    "        self.train_bias, self.membrane_filter)\n",
    "\n",
    "        self.output_filter = filter_layer(300, self.length, self.batch_size, filter_tau_m, filter_tau_s)                                \n",
    "    def forward(self, inputs):\n",
    "        \"\"\"\n",
    "        :param inputs: [batch, input_size, t]\n",
    "        :return:\n",
    "        \"\"\"\n",
    "        axon1_states = self.axon1.create_init_states()\n",
    "        dense1_states = self.dense1.create_init_states()\n",
    "\n",
    "        axon2_states = self.axon2.create_init_states()\n",
    "        dense2_states = self.dense2.create_init_states()\n",
    "\n",
    "        axon3_states = self.axon3.create_init_states()\n",
    "        dense3_states = self.dense3.create_init_states()\n",
    "\n",
    "        axon4_states = self.axon4.create_init_states()\n",
    "        dense4_states = self.dense4.create_init_states()\n",
    "\n",
    "        axon1_out, axon1_states = self.axon1(inputs, axon1_states)\n",
    "        spike_dense1, dense1_states = self.dense1(axon1_out, dense1_states)\n",
    "\n",
    "        axon2_out, axon2_states = self.axon2(spike_dense1, axon2_states)\n",
    "        spike_dense2, dense2_states = self.dense2(axon2_out, dense2_states)\n",
    "\n",
    "        axon3_out, axon3_states = self.axon3(spike_dense2, axon3_states)\n",
    "        spike_dense3, dense3_states = self.dense3(axon3_out, dense3_states)\n",
    "\n",
    "        axon4_out, axon4_states = self.axon4(spike_dense3, axon4_states)\n",
    "        spike_dense4, dense4_states = self.dense4(axon4_out, dense4_states)\n",
    "\n",
    "        filtered_output = self.output_filter(spike_dense4)\n",
    "\n",
    "        return spike_dense4, filtered_output"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train and test function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, optimizer, scheduler, data_loader):\n",
    "    '''\n",
    "\n",
    "    '''\n",
    "\n",
    "    model.train()\n",
    "\n",
    "    eval_image_number = 0\n",
    "    correct_total = 0\n",
    "    wrong_total = 0\n",
    "\n",
    "    criterion = torch.nn.MSELoss()\n",
    "\n",
    "    for i_batch, sample_batched in enumerate(data_loader):\n",
    "\n",
    "        x_train = sample_batched[0].to(device)\n",
    "        label = sample_batched[1].to(device)\n",
    "        target_pattern = sample_batched[2].to(device)\n",
    "        filtered_target_pattern = sample_batched[3].to(device)\n",
    "\n",
    "        out_spike, filtered_out_spike = model(x_train)\n",
    "\n",
    "        model.zero_grad()\n",
    "\n",
    "        loss = criterion(filtered_out_spike, filtered_target_pattern)\n",
    "\n",
    "        print('train loss: {}'.format(loss))\n",
    "\n",
    "        loss.backward()\n",
    "\n",
    "        optimizer.step()\n",
    "\n",
    "        if isinstance(scheduler, torch.optim.lr_scheduler.CosineAnnealingWarmRestarts) or \\\n",
    "                isinstance(scheduler, torch.optim.lr_scheduler.CyclicLR):\n",
    "            # isinstance(scheduler, torch.optim.lr_scheduler.OneCycleLR):\n",
    "            scheduler.step()\n",
    "\n",
    "    if isinstance(scheduler, torch.optim.lr_scheduler.MultiStepLR):\n",
    "        scheduler.step()\n",
    "\n",
    "    return loss\n",
    "\n",
    "def test(model, data_loader):\n",
    "\n",
    "    model.eval()\n",
    "\n",
    "    loss_list = []\n",
    "\n",
    "    criterion = torch.nn.MSELoss()\n",
    "\n",
    "    for i_batch, sample_batched in enumerate(data_loader):\n",
    "\n",
    "        criterion = torch.nn.MSELoss()\n",
    "        x_test = sample_batched[0].to(device)\n",
    "        label = sample_batched[1].to(device)\n",
    "        target_pattern = sample_batched[2].to(device)\n",
    "        filtered_target_pattern = sample_batched[3].to(device)\n",
    "\n",
    "        out_spike, filtered_out_spike = model(x_test)\n",
    "\n",
    "        loss = criterion(filtered_out_spike, filtered_target_pattern)\n",
    "\n",
    "        print('Test loss: {}'.format(loss))\n",
    "\n",
    "        loss_list.append(loss.cpu().detach().numpy())\n",
    "\n",
    "    #calculate loss of this epoch\n",
    "    loss_array = np.stack(loss_list)\n",
    "    average_loss = loss_array.mean()\n",
    "\n",
    "    return average_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Train the network"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if torch.cuda.is_available():\n",
    "    device = torch.device('cuda:0')\n",
    "else:\n",
    "    device = torch.device('cpu')\n",
    "\n",
    "original_pattern = np.load(dataset_path).astype(np.float32)\n",
    "\n",
    "delayed_target = np.zeros(original_pattern.shape, dtype=np.float32)\n",
    "\n",
    "delayed_target[:, :, 4:] = original_pattern[:, :, :-4]\n",
    "\n",
    "filtered_target = []\n",
    "for target_idx, delayed_target_pattern in enumerate(delayed_target):\n",
    "    filtered_target.append(utilities.filter_spike_multiple(delayed_target_pattern, filter_type='dual_exp',\n",
    "                                                            tau_m=filter_tau_m, tau_s=filter_tau_s))\n",
    "# shape [pattern idx, spike train idx, time]\n",
    "filtered_target = np.array(filtered_target)\n",
    "\n",
    "snn = mysnn().to(device)\n",
    "\n",
    "params = list(snn.parameters())\n",
    "\n",
    "optimizer_class = optimizer_config[optimizer_choice][0]\n",
    "learning_rate = optimizer_config[optimizer_choice][1]\n",
    "optimizer = optimizer_class(params, learning_rate)\n",
    "\n",
    "scheduler = None\n",
    "if scheduler_choice == 0:\n",
    "    scheduler = None\n",
    "elif scheduler_choice == 1:\n",
    "    scheduler = torch.optim.lr_scheduler.MultiStepLR(optimizer, \n",
    "    scheduler_config[scheduler_choice][1], \n",
    "    scheduler_config[scheduler_choice][2])\n",
    "elif scheduler_choice == 2:\n",
    "    scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, scheduler_config[scheduler_choice][1])\n",
    "elif scheduler_choice == 3:\n",
    "    scheduler = torch.optim.lr_scheduler.CyclicLR(optimizer, \n",
    "    scheduler_config[scheduler_choice][1], \n",
    "    scheduler_config[scheduler_choice][2], \n",
    "    scheduler_config[scheduler_choice][3])\n",
    "elif scheduler_choice == 4:\n",
    "    scheduler = torch.optim.lr_scheduler.OneCycleLR(optimizer, \n",
    "    scheduler_config[scheduler_choice][1], \n",
    "    scheduler_config[scheduler_choice][2], \n",
    "    scheduler_config[scheduler_choice][3])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train and test dataloader"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_data = PatternDataset(original_pattern, delayed_target, filtered_target, length)\n",
    "train_dataloader = DataLoader(train_data, batch_size=batch_size, shuffle=True, drop_last=True)\n",
    "\n",
    "test_data = PatternDataset(original_pattern, delayed_target, filtered_target, length)\n",
    "test_dataloader = DataLoader(test_data, batch_size=batch_size, shuffle=True, drop_last=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "outputPrepend"
    ]
   },
   "outputs": [],
   "source": [
    "best_epoch = 0\n",
    "lowest_loss = 100\n",
    "\n",
    "for j in range(epoch):\n",
    "    print('Train epoch: {}'.format(j))\n",
    "\n",
    "    snn.train()\n",
    "\n",
    "    train_loss = train(snn, optimizer, scheduler, train_dataloader)\n",
    "\n",
    "    save_checkpoint = True\n",
    "    if save_checkpoint:\n",
    "        checkpoint_path = os.path.join(checkpoint_base_path, checkpoint_base_name + str(j))\n",
    "        torch.save({\n",
    "            'epoch': j,\n",
    "            'snn_state_dict': snn.state_dict(),\n",
    "            'optimizer_state_dict': optimizer.state_dict(),\n",
    "            'loss': train_loss,\n",
    "        }, checkpoint_path)\n",
    "    \n",
    "    snn.eval()\n",
    "    average_loss = test(snn, test_dataloader)\n",
    "    if lowest_loss > average_loss:\n",
    "        lowest_loss = average_loss\n",
    "        best_epoch = j\n",
    "    \n",
    "    print('Test: {}, loss: {}'.format(j, average_loss))\n",
    "\n",
    "print('Best epoch: {}'.format(best_epoch))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run the best epoch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "input_list = []\n",
    "output_list = []\n",
    "\n",
    "# checkpoint_path = os.path.join(checkpoint_base_path, checkpoint_base_name + str(best_epoch))\n",
    "# checkpoint = torch.load(checkpoint_path)\n",
    "\n",
    "# You can replpace checkpoint_path by a pre-trained checkpoint\n",
    "checkpoint_path = './associative_memory_checkpoint/pretrained_associative_memory'\n",
    "checkpoint = torch.load(checkpoint_path)\n",
    "\n",
    "snn.load_state_dict(checkpoint[\"snn_state_dict\"])\n",
    "\n",
    "print('Checkpoint: {} loaded'.format(checkpoint_path))\n",
    "\n",
    "snn.eval()\n",
    "\n",
    "for i_batch, sample_batched in enumerate(test_dataloader):\n",
    "\n",
    "    criterion = torch.nn.MSELoss()\n",
    "    x_test = sample_batched[0].to(device)\n",
    "    label = sample_batched[1].to(device)\n",
    "    target_pattern = sample_batched[2].to(device)\n",
    "    filtered_target_pattern = sample_batched[3].to(device)\n",
    "\n",
    "    out_spike, filtered_out_spike = snn(x_test)\n",
    "\n",
    "    loss = criterion(filtered_out_spike, filtered_target_pattern)\n",
    "\n",
    "    print(i_batch, loss)\n",
    "\n",
    "    input_list.append(x_test.cpu().detach().numpy())\n",
    "    output_list.append(out_spike.cpu().detach().numpy())"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot some samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, axs = plt.subplots(2, 4)\n",
    "axs[0, 0].imshow(input_list[0][0])\n",
    "axs[0, 1].imshow(output_list[0][0])\n",
    "\n",
    "axs[1, 0].imshow(input_list[0][1])\n",
    "axs[1, 1].imshow(output_list[0][1])\n",
    "\n",
    "axs[0, 2].imshow(input_list[0][2])\n",
    "axs[0, 3].imshow(output_list[0][2])\n",
    "\n",
    "axs[1, 2].imshow(input_list[0][3])\n",
    "axs[1, 3].imshow(output_list[0][3])\n",
    "\n",
    "fig, axs = plt.subplots(2, 4)\n",
    "axs[0, 0].imshow(input_list[0][25])\n",
    "axs[0, 1].imshow(output_list[0][25])\n",
    "\n",
    "axs[1, 0].imshow(input_list[0][26])\n",
    "axs[1, 1].imshow(output_list[0][26])\n",
    "\n",
    "axs[0, 2].imshow(input_list[0][23])\n",
    "axs[0, 3].imshow(output_list[0][23])\n",
    "\n",
    "axs[1, 2].imshow(input_list[0][27])\n",
    "axs[1, 3].imshow(output_list[0][27])\n",
    "\n",
    "plt.figure()\n",
    "utilities.plot_raster_dot(output_list[0][0])\n",
    "plt.figure()\n",
    "utilities.plot_raster_dot(input_list[0][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ]
}