dataset.txt

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import numpy.random as rd\n",
    "import sys\n",
    "\n",
    "\n",
    "class CueAccumulationDataset(torch.utils.data.Dataset):\n",
    "    \"\"\"Adapted from the original TensorFlow e-prop implemation from TU Graz, available at https://github.com/IGITUGraz/eligibility_propagation\"\"\"\n",
    "\n",
    "    def __init__(self, args, type):\n",
    "      \n",
    "        n_cues     = 7\n",
    "        f0         = 40\n",
    "        t_cue      = 100\n",
    "        t_wait     = 1200\n",
    "        n_symbols  = 4\n",
    "        p_group    = 0.3\n",
    "        \n",
    "        self.dt         = 1e-3\n",
    "        self.t_interval = 150\n",
    "        self.seq_len    = n_cues*self.t_interval + t_wait\n",
    "        self.n_in       = 40\n",
    "        self.n_out      = 2    # This is a binary classification task, so using two output units with a softmax activation redundant\n",
    "        n_channel       = self.n_in // n_symbols\n",
    "        prob0           = f0 * self.dt\n",
    "        t_silent        = self.t_interval - t_cue\n",
    "        \n",
    "        if (type == 'train'):\n",
    "            length = args.train_len\n",
    "        else:\n",
    "            length = args.test_len\n",
    "            \n",
    "    \n",
    "        # Randomly assign group A and B\n",
    "        prob_choices = np.array([p_group, 1 - p_group], dtype=np.float32)\n",
    "        idx = rd.choice([0, 1], length)\n",
    "        probs = np.zeros((length, 2), dtype=np.float32)\n",
    "        # Assign input spike probabilities\n",
    "        probs[:, 0] = prob_choices[idx]\n",
    "        probs[:, 1] = prob_choices[1 - idx]\n",
    "    \n",
    "        cue_assignments = np.zeros((length, n_cues), dtype=np.int)\n",
    "        # For each example in batch, draw which cues are going to be active (left or right)\n",
    "        for b in range(length):\n",
    "            cue_assignments[b, :] = rd.choice([0, 1], n_cues, p=probs[b])\n",
    "    \n",
    "        # Generate input spikes\n",
    "        input_spike_prob = np.zeros((length, self.seq_len, self.n_in))\n",
    "        t_silent = self.t_interval - t_cue\n",
    "        for b in range(length):\n",
    "            for k in range(n_cues):\n",
    "                # Input channels only fire when they are selected (left or right)\n",
    "                c = cue_assignments[b, k]\n",
    "                input_spike_prob[b, t_silent+k*self.t_interval:t_silent+k*self.t_interval+t_cue, c*n_channel:(c+1)*n_channel] = prob0\n",
    "    \n",
    "        # Recall cue and background noise\n",
    "        input_spike_prob[:, -self.t_interval:, 2*n_channel:3*n_channel] = prob0\n",
    "        input_spike_prob[:, :, 3*n_channel:] = prob0/4.\n",
    "        input_spikes = generate_poisson_noise_np(input_spike_prob)\n",
    "        self.x = torch.tensor(input_spikes).float()\n",
    "    \n",
    "        # Generate targets\n",
    "        target_nums = np.zeros((length, self.seq_len), dtype=np.int)\n",
    "        target_nums[:, :] = np.transpose(np.tile(np.sum(cue_assignments, axis=1) > int(n_cues/2), (self.seq_len, 1)))\n",
    "        self.y = torch.tensor(target_nums).long()\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.y)\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        return self.x[index], self.y[index]\n",
    "\n",
    "\n",
    "\n",
    "def setup(args):\n",
    "    args.cuda = not args.cpu and torch.cuda.is_available()\n",
    "    if args.cuda:\n",
    "        print(\"=== The available CUDA GPU will be used for computations.\")\n",
    "        device = torch.cuda.current_device()\n",
    "    else:\n",
    "        device = torch.device('cpu')\n",
    "    \n",
    "    kwargs = {'num_workers': 0, 'pin_memory': True} if args.cuda else {}\n",
    "\n",
    "    if args.dataset == \"cue_accumulation\":\n",
    "        print(\"=== Loading cue evidence accumulation dataset...\")\n",
    "        (train_loader, traintest_loader, test_loader) = load_dataset_cue_accumulation(args, kwargs)\n",
    "    else:\n",
    "        print(\"=== ERROR - Unsupported dataset ===\")\n",
    "        sys.exit(1)\n",
    "        \n",
    "    print(\"Training set length: \"+str(args.full_train_len))\n",
    "    print(\"Test set length: \"+str(args.full_test_len))\n",
    "    \n",
    "    return (device, train_loader, traintest_loader, test_loader)\n",
    "\n",
    "\n",
    "def load_dataset_cue_accumulation(args, kwargs):\n",
    "\n",
    "    trainset = CueAccumulationDataset(args,\"train\")\n",
    "    testset  = CueAccumulationDataset(args,\"test\")\n",
    "\n",
    "    train_loader     = torch.utils.data.DataLoader(trainset, batch_size=args.batch_size,      shuffle=args.shuffle, **kwargs)\n",
    "    traintest_loader = torch.utils.data.DataLoader(trainset, batch_size=args.test_batch_size, shuffle=False       , **kwargs)\n",
    "    test_loader      = torch.utils.data.DataLoader(testset , batch_size=args.test_batch_size, shuffle=False       , **kwargs)\n",
    "    \n",
    "    args.n_classes      = trainset.n_out\n",
    "    args.n_steps        = trainset.seq_len\n",
    "    args.n_inputs       = trainset.n_in\n",
    "    args.dt             = trainset.dt\n",
    "    args.classif        = True\n",
    "    args.full_train_len = len(trainset)\n",
    "    args.full_test_len  = len(testset)\n",
    "    args.delay_targets  = trainset.t_interval\n",
    "    args.skip_test      = False\n",
    "    \n",
    "    return (train_loader, traintest_loader, test_loader)\n",
    "\n",
    "\n",
    "def generate_poisson_noise_np(prob_pattern, freezing_seed=None):\n",
    "    if isinstance(prob_pattern, list):\n",
    "        return [generate_poisson_noise_np(pb, freezing_seed=freezing_seed) for pb in prob_pattern]\n",
    "\n",
    "    shp = prob_pattern.shape\n",
    "\n",
    "    if not(freezing_seed is None): rng = rd.RandomState(freezing_seed)\n",
    "    else: rng = rd.RandomState()\n",
    "\n",
    "    spikes = prob_pattern > rng.rand(prob_pattern.size).reshape(shp)\n",
    "    return spikes\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "\n",
    "\n",
    "def args():\n",
    "    parser = argparse.ArgumentParser(\n",
    "        description=\"Train a reservoir based SNN on biosignals\"\n",
    "    )\n",
    "\n",
    "    parser.add_argument('--cpu', action='store_true', default=True, help='Disable CUDA training and run training on CPU')\n",
    "    parser.add_argument('--dataset', type=str, choices = ['cue_accumulation'], default='cue_accumulation', help='Choice of the dataset')\n",
    "    parser.add_argument('--shuffle', type=bool, default=True, help='Enables shuffling sample order in datasets after each epoch')\n",
    "    parser.add_argument('--trials', type=int, default=1, help='Nomber of trial experiments to do (i.e. repetitions with different initializations)')\n",
    "    parser.add_argument('--epochs', type=int, default=10, help='Number of epochs to train')\n",
    "    parser.add_argument('--optimizer', type=str, choices = ['SGD', 'NAG', 'Adam', 'RMSProp'], default='Adam', help='Choice of the optimizer')\n",
    "    parser.add_argument('--loss', type=str, choices = ['MSE', 'BCE', 'CE'], default='BCE', help='Choice of the loss function (only for performance monitoring purposes, does not influence learning)')\n",
    "    parser.add_argument('--lr', type=float, default=1e-4, help='Initial learning rate')\n",
    "    parser.add_argument('--lr-layer-norm', type=float, nargs='+', default=(0.05,0.05,1.0), help='Per-layer modulation factor of the learning rate')\n",
    "    parser.add_argument('--batch-size', type=int, default=5, help='Batch size for training (limited by the available GPU memory)')\n",
    "    parser.add_argument('--test-batch-size', type=int, default=5, help='Batch size for testing (limited by the available GPU memory)')\n",
    "    parser.add_argument('--train-len', type=int, default=80, help='Number of training set samples')\n",
    "    parser.add_argument('--test-len', type=int, default=20, help='Number of test set samples')\n",
    "    parser.add_argument('--visualize', type=bool, default=True, help='Enable network visualization')\n",
    "    parser.add_argument('--visualize-light', type=bool, default=True, help='Enable light mode in network visualization, plots traces only for a single neuron')\n",
    "    # Network model parameters\n",
    "    parser.add_argument('--n_rec', type=int, default=100, help='Number of recurrent units')\n",
    "    parser.add_argument('--model', type=str, choices = ['LIF'], default='LIF', help='Neuron model in the recurrent layer. Support for the ALIF neuron model has been removed.')\n",
    "    parser.add_argument('--threshold', type=float, default=0.6, help='Firing threshold in the recurrent layer')\n",
    "    parser.add_argument('--tau-mem', type=float, default=2000e-3, help='Membrane potential leakage time constant in the recurrent layer (in seconds)')\n",
    "    parser.add_argument('--tau-out', type=float, default=20e-3, help='Membrane potential leakage time constant in the output layer (in seconds)')\n",
    "    parser.add_argument('--bias-out', type=float, default=0.0, help='Bias of the output layer')\n",
    "    parser.add_argument('--gamma', type=float, default=0.3, help='Surrogate derivative magnitude parameter')\n",
    "    parser.add_argument('--w-init-gain', type=float, nargs='+', default=(0.5,0.1,0.5), help='Gain parameter for the He Normal initialization of the input, recurrent and output layer weights')\n",
    "    \n",
    "\n",
    "    my_args = parser.parse_args()\n",
    "\n",
    "    return my_args"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'function' object has no attribute 'cpu'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "\u001b[1;32mC:\\Users\\NIKHIL~1.GAR\\AppData\\Local\\Temp/ipykernel_10048/3045256771.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0mtrain\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m      2\u001b[0m \u001b[1;32mimport\u001b[0m \u001b[0msetup\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m----> 3\u001b[1;33m \u001b[1;33m(\u001b[0m\u001b[0mdevice\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtrain_loader\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtraintest_loader\u001b[0m\u001b[1;33m,\u001b[0m \u001b[0mtest_loader\u001b[0m\u001b[1;33m)\u001b[0m \u001b[1;33m=\u001b[0m \u001b[0msetup\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0msetup\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;32mc:\\Users\\nikhil.garg\\SNN-Toolbox\\setup.py\u001b[0m in \u001b[0;36msetup\u001b[1;34m(args)\u001b[0m\n\u001b[0;32m    108\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    109\u001b[0m \u001b[1;32mdef\u001b[0m \u001b[0msetup\u001b[0m\u001b[1;33m(\u001b[0m\u001b[0margs\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m--> 110\u001b[1;33m     \u001b[0margs\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcuda\u001b[0m \u001b[1;33m=\u001b[0m \u001b[1;32mnot\u001b[0m \u001b[0margs\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcpu\u001b[0m \u001b[1;32mand\u001b[0m \u001b[0mtorch\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcuda\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mis_available\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m    111\u001b[0m     \u001b[1;32mif\u001b[0m \u001b[0margs\u001b[0m\u001b[1;33m.\u001b[0m\u001b[0mcuda\u001b[0m\u001b[1;33m:\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m    112\u001b[0m         \u001b[0mprint\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;34m\"=== The available CUDA GPU will be used for computations.\"\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mAttributeError\u001b[0m: 'function' object has no attribute 'cpu'"
     ]
    }
   ],
   "source": [
    "import train\n",
    "import setup\n",
    "(device, train_loader, traintest_loader, test_loader) = setup.setup(args) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "for batch_idx, (data, label) in enumerate(loader):"
   ]
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "66583042dd22f0377bf472c38da1b9014a0e7c4bef1042f28d5f3ecd26668a1e"
  },
  "kernelspec": {
   "display_name": "Python 3.9.7 ('base')",
   "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.7"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}