Layer-wise SSL - MLP - fixed non-linear projectors - simultaneous perturbations stochastic approximation - stochastic sequential objective.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Layer-wise Self-supervised training of a 4-layers MLP on MNIST "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import matplotlib.pyplot as plt\n",
    "from tqdm import tqdm\n",
    "import torchvision.transforms as transforms\n",
    "from torch.utils.data import Dataset\n",
    "import torchvision\n",
    "import optuna\n",
    "import os\n",
    "import argparse\n",
    "from tqdm import tqdm\n",
    "\n",
    "from utils_mlp_spsa_stochastic_objective import*"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Hyperparameters of the simulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "parser = argparse.ArgumentParser(description='Train a MLP with layer-wise SSL on MNIST - no BP but SPSA')\n",
    "# general params\n",
    "parser.add_argument(\n",
    "    '--device',\n",
    "    type=int,\n",
    "    default=2,\n",
    "    help='GPU name to use cuda')\n",
    "parser.add_argument(\n",
    "    '--dataset',\n",
    "    type=str,\n",
    "    default='mnist',\n",
    "    help='Dataset we use for training (default=mnist, others: None for this specific notebook)')\n",
    "# Architecture params\n",
    "parser.add_argument(\n",
    "    '--nlayers',\n",
    "    type=int,\n",
    "    default=2,\n",
    "help='Number of layers for the main MLP (default: 4)')\n",
    "parser.add_argument(\n",
    "    '--nneurons',\n",
    "    type=int,\n",
    "    default=1000,\n",
    "help='Number of neuron per layer of the main MLP (default: 1000)')\n",
    "parser.add_argument(\n",
    "    '--nlayers_proj',\n",
    "    type=int,\n",
    "    default=3,\n",
    "help='Number of layers for each non-linear projector (default: 3)')\n",
    "parser.add_argument(\n",
    "    '--nneurons_proj',\n",
    "    type=int,\n",
    "    default=256,\n",
    "help='Number of neuron per layer of each non-linear projector (default: 256)')\n",
    "# Optimization params\n",
    "parser.add_argument(\n",
    "    '--epochs',\n",
    "    type=int,\n",
    "    default=200,\n",
    "help='Number of epochs to train each layer (default: 500)')\n",
    "parser.add_argument(\n",
    "    '--epochs_classifier',\n",
    "    type=int,\n",
    "    default=20,\n",
    "help='Number of epochs to train a linear classifier on top of each layer (default: 20)')\n",
    "parser.add_argument(\n",
    "    '--batchSize_pretrain',\n",
    "    type=int,\n",
    "    default=128,\n",
    "    help='Batch size for pre-training (default=256)')\n",
    "parser.add_argument(\n",
    "    '--n_average',\n",
    "    type=int,\n",
    "    default=10,\n",
    "    help='Number of sub-mini-batches in a mini-batches (default=10)')\n",
    "parser.add_argument(\n",
    "    '--batchSize_classifier',\n",
    "    type=int,\n",
    "    default=64,\n",
    "    help='Batch size for training the linear classifier (default=64)')\n",
    "parser.add_argument(\n",
    "    '--test_batchSize',\n",
    "    type=int,\n",
    "    default=512,\n",
    "    help='Testing Batch size (default=512)')\n",
    "parser.add_argument(\n",
    "    '--lr',\n",
    "    type=float,\n",
    "    default=1e-4,\n",
    "    help='Learning rate for Adam optimizer for pre-training (default=1e-4)')\n",
    "parser.add_argument(\n",
    "    '--lr_classifier',\n",
    "    type=float,\n",
    "    default=1e-3,\n",
    "    help='Learning rate for Adam optimizer for training the linear classifier (default=1e-4)')\n",
    "parser.add_argument(\n",
    "    '--initial_step',\n",
    "    type=float,\n",
    "    default=1e-2,\n",
    "    help='Initial step of the random perturbation (default=1e-2)')\n",
    "parser.add_argument(\n",
    "    '--gamma_perturbation',\n",
    "    type=float,\n",
    "    default=0.95,\n",
    "    help='Decay rate for the size of the random perturbation applied at each inference (default=0.8)')\n",
    "parser.add_argument(\n",
    "    '--gamma_optimizer',\n",
    "    type=float,\n",
    "    default=0.9,\n",
    "    help='Decay rate for the learning rate for pre-training (default=0.9)')\n",
    "# SSL settings\n",
    "parser.add_argument(\n",
    "    '--nviews',\n",
    "    type=int,\n",
    "    default=2,\n",
    "    help='Number of views for computing the SSL objective (default=2)')\n",
    "parser.add_argument(\n",
    "    '--deg',\n",
    "    type=int,\n",
    "    default=15,\n",
    "    help='Maximum angle for the RandomRotation transform applied to the input image (default=15)')\n",
    "parser.add_argument(\n",
    "    '--pad',\n",
    "    type=int,\n",
    "    default=2,\n",
    "    help='Padding applied before RandomCrop (default=2)')\n",
    "parser.add_argument(\n",
    "    '--contrast',\n",
    "    type=float,\n",
    "    default=0.5,\n",
    "    help='Contrast value for ColorJittering (default=0.5)')\n",
    "parser.add_argument(\n",
    "    '--hue',\n",
    "    type=float,\n",
    "    default=0.5,\n",
    "    help='Hue value for ColorJittering (default=0.5)')\n",
    "parser.add_argument(\n",
    "    '--scaleaffine',\n",
    "    nargs='+',\n",
    "    type=float,\n",
    "    default=[0.7, 1.3],\n",
    "    help='Scale parameters for the RandomAffine transform applied (default=(0.7, 1.3))')\n",
    "# SSL objective params (for the sigmoidal parametrization)\n",
    "parser.add_argument(\n",
    "    '--scale',\n",
    "    nargs='+',\n",
    "    type=float,\n",
    "    default=[20, 20, 1, 0],\n",
    "    help='Scale value for the sigmoidal parametrization (default=(10, 20, 1, 0))')\n",
    "parser.add_argument(\n",
    "    '--slope',\n",
    "    nargs='+',\n",
    "    type=float,\n",
    "    default=[0,0,0,0],\n",
    "    help='Slope value for the sigmoidal parametrization (default=(10, 20, 1, 0))')\n",
    "parser.add_argument(\n",
    "    '--threshold',\n",
    "    nargs='+',\n",
    "    type=float,\n",
    "    default=[3,3,2,0],\n",
    "    help='Threshold value for the sigmoidal parametrization (default=(10, 20, 1, 0))')\n",
    "parser.add_argument(\n",
    "    '--bias',\n",
    "    nargs='+',\n",
    "    type=float,\n",
    "    default=[5,15,1,0],\n",
    "    help='Bias value for the sigmoidal parametrization (default=(10, 20, 1, 0))')\n",
    "\n",
    "\n",
    "args, _ = parser.parse_known_args()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dataset creation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# Data and target transfomations\n",
    "class ReshapeTransform:\n",
    "    def __init__(self, new_size):\n",
    "        self.new_size = new_size\n",
    "\n",
    "    def __call__(self, img):\n",
    "        return torch.reshape(img, self.new_size)\n",
    "        \n",
    "        \n",
    "class ReshapeTransformTarget:\n",
    "    def __init__(self, number_classes):\n",
    "        self.number_classes = number_classes\n",
    "    \n",
    "    def __call__(self, target):\n",
    "        target=torch.tensor(target).unsqueeze(0).unsqueeze(1)\n",
    "        target_onehot = torch.zeros((1,self.number_classes))      \n",
    "        return target_onehot.scatter_(1, target, 1).squeeze(0)\n",
    "\n",
    "    \n",
    "class ContrastiveTransformations(object):\n",
    "    def __init__(self, base_transforms, n_views=2, n_average = 1):\n",
    "        self.base_transforms = base_transforms #random transformations\n",
    "        self.n_views = n_views # number of differents copies with different \n",
    "        self.n_average = n_average\n",
    "        \n",
    "    def __call__(self, x):\n",
    "        return [[self.base_transforms(x) for i in range(self.n_views)] for n in range(self.n_average)]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Dataset creation\n",
    "def get_dataloader(args):\n",
    "    '''\n",
    "    Function that returns the dataloaders given the current hyperparameters \n",
    "    '''\n",
    "    # random data augmentations for the pre-training stage\n",
    "    contrast_transforms =  transforms.Compose([torchvision.transforms.RandomRotation(degrees = args.deg, fill=0), #random rotation\n",
    "                                               torchvision.transforms.RandomCrop((28,28), padding = args.pad), #random crop\n",
    "                                               torchvision.transforms.RandomAffine(degrees=(0, 0), translate=(0.0, 0.0), scale=(args.scaleaffine[0], args.scaleaffine[1])),\n",
    "                                               torchvision.transforms.ColorJitter(brightness=0, contrast = args.contrast, saturation=0, hue = args.hue),\n",
    "                                               torchvision.transforms.ToTensor(),\n",
    "                                               ReshapeTransform((-1,))])\n",
    "    \n",
    "    # fixed transformation for training testing the linear classifier\n",
    "    transforms_test = transforms.Compose([torchvision.transforms.ToTensor(),\n",
    "                                           ReshapeTransform((-1,))])\n",
    "\n",
    "    # Train loader for pre-training\n",
    "    train_loader = torch.utils.data.DataLoader(\n",
    "    torchvision.datasets.MNIST(root='./data', train=True, download=True,\n",
    "                                transform = ContrastiveTransformations(contrast_transforms, n_views=args.nviews, n_average=args.n_average),\n",
    "                                target_transform=ReshapeTransformTarget(10)), batch_size = args.batchSize_pretrain, shuffle=True)\n",
    "\n",
    "    # Small train loader for pre-training\n",
    "    train_loader_small = torch.utils.data.DataLoader(\n",
    "    torchvision.datasets.MNIST(root='./data', train=False, download=True,\n",
    "                                transform = ContrastiveTransformations(contrast_transforms, n_views=args.nviews, n_average=args.n_average),\n",
    "                                target_transform=ReshapeTransformTarget(10)), batch_size = args.batchSize_pretrain, shuffle=True)\n",
    "\n",
    "    \n",
    "    # Train loader for the linear classifier\n",
    "    train_loader_classifier = torch.utils.data.DataLoader(torchvision.datasets.MNIST(root='./data', train=True, download=True,\n",
    "                                transform = transforms_test,\n",
    "                                target_transform=ReshapeTransformTarget(10)), batch_size = args.batchSize_classifier, shuffle=True)\n",
    "\n",
    "    # Test loader for the linear classifier\n",
    "    test_loader = torch.utils.data.DataLoader(torchvision.datasets.MNIST(root='./data', train=False, download=True,\n",
    "                                transform = transforms_test,\n",
    "                                target_transform=ReshapeTransformTarget(10)), batch_size = args.test_batchSize, shuffle=False)\n",
    "    \n",
    "    \n",
    "    return train_loader, train_loader_small, train_loader_classifier, test_loader\n",
    " "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "train_loader_pretrain, train_loader_pretrain_small, train_loader_classifier, test_loader = get_dataloader(args)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Visualize the random transformations applied on mnist data\n",
    "# Data from the pretraining dataloader is a list with 2 elements! we can stack them to have a global mini-batch in the training loop ?\n",
    "\n",
    "data, target = next(iter(train_loader_pretrain))\n",
    "print(len(data))\n",
    "plt.figure()\n",
    "plt.imshow(data[0][0][0].view(28,28), cmap = \"gray\")\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(data[0][1][0].view(28,28), cmap = \"gray\")\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(data[0][0][0].view(28,28) - data[0][1][0].view(28,28), cmap = \"gray\")\n",
    "\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SSL (VicREG) Objective"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def off_diagonal(x):\n",
    "    n, m = x.shape\n",
    "    assert n == m\n",
    "    return x.flatten()[:-1].view(n - 1, n + 1)[:, 1:].flatten()\n",
    "\n",
    "# new version of the VicRegLoss: add an intra-sample variance objective! we want the different views of the\n",
    "# same inputs to be initially quite different then they should be tend to be more similar as we go deeper\n",
    "# initial version is from the Meta github repo\n",
    "\n",
    "class VicRegLoss(torch.nn.Module):\n",
    "    '''\n",
    "    Class that implement the VicReg loss\n",
    "    '''\n",
    "    def __init__(self, device, sim_coeff = 25, std_coeff = 25, cov_coeff = 1, intra_std_coeff = 1, intra_std_target = 1):\n",
    "        super(VicRegLoss, self).__init__()\n",
    "        self.device = device\n",
    "        self.sim_coeff = sim_coeff\n",
    "        self.std_coeff = std_coeff\n",
    "        self.cov_coeff = cov_coeff\n",
    "        self.intra_std_coeff = intra_std_coeff\n",
    "        self.intra_std_target = intra_std_target\n",
    "        self.flatten = nn.Flatten()\n",
    "        \n",
    "    def forward(self, x, y):\n",
    "        '''\n",
    "        Args:\n",
    "            x,y: embeddings vectors - are \"flattened\" so they have dimension (batch_size, num_features)\n",
    "        '''\n",
    "        batch_size = x.size(0)\n",
    "        num_features = x.size(1)\n",
    "        repr_loss = F.mse_loss(x, y) \n",
    "\n",
    "        x = x - x.mean(dim=0) \n",
    "        y = y - y.mean(dim=0) \n",
    "                \n",
    "        std_x = torch.sqrt(x.var(dim=0) + 0.0001) \n",
    "        std_y = torch.sqrt(y.var(dim=0) + 0.0001) \n",
    "        std_loss = torch.mean(F.relu(1 - std_x)) / 2 + torch.mean(F.relu(1 - std_y)) / 2\n",
    "        \n",
    "        cov_x = (x.T @ x) / (batch_size - 1)\n",
    "        cov_y = (y.T @ y) / (batch_size - 1)\n",
    "        \n",
    "        diag = torch.eye(num_features, device=self.device)\n",
    "\n",
    "        cov_loss = off_diagonal(cov_x).pow_(2).sum().div(\n",
    "            num_features # num elements here?\n",
    "        ) + off_diagonal(cov_y).pow_(2).sum().div(num_features)\n",
    "\n",
    "        return repr_loss, std_loss, cov_loss"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Model definition\n",
    "- MLP with Relu\n",
    "- each layer is trained with a non-linear projector (MLP wiht Relu)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Network class\n",
    "class Network(nn.Module):\n",
    "    ''' \n",
    "    Define the network used\n",
    "    '''\n",
    "    #def __init__(self, thetas_coefs, thetas_exponents, run_gpu):\n",
    "    def __init__(self, args, ysim, yvar, ycovar, yintravar):\n",
    "\n",
    "        super(Network, self).__init__()\n",
    "        \n",
    "        self.n_views = args.nviews\n",
    "        \n",
    "        self.n_layers = args.nlayers\n",
    "        self.n_neurons = args.nneurons\n",
    "\n",
    "        self.n_layers_proj = args.nlayers_proj\n",
    "        self.n_neurons_proj = args.nneurons_proj\n",
    "        \n",
    "        self.initial_step = args.initial_step\n",
    "        self.step = args.initial_step #this will be updated after each epoch: step(epoch + 1) = step(epoch)*gamma_perturbation\n",
    "        self.gamma_perturbation = args.gamma_perturbation\n",
    "        \n",
    "        self.layers = [nn.Linear(784, self.n_neurons, bias = False)]\n",
    "        self.layers += [nn.Linear(self.n_neurons, self.n_neurons, bias = False) for k in range(self.n_layers-1)]\n",
    "        self.layers = nn.Sequential(*self.layers)\n",
    "        \n",
    "        # need to optimize how we create the non-linear projectors to parametrize the numbers we create\n",
    "        self.projs1 = [nn.Linear(self.n_neurons, self.n_neurons_proj, bias = False)]\n",
    "        self.projs1 += [nn.Linear(self.n_neurons_proj, self.n_neurons_proj, bias = False) for k in range(self.n_layers_proj - 1)]\n",
    "        self.projs1 = nn.Sequential(*self.projs1)\n",
    "        \n",
    "        self.projs2 = [nn.Linear(self.n_neurons, self.n_neurons_proj, bias = False)]\n",
    "        self.projs2 += [nn.Linear(self.n_neurons_proj, self.n_neurons_proj, bias = False) for k in range(self.n_layers_proj - 1)]\n",
    "        self.projs2 = nn.Sequential(*self.projs2)\n",
    "        \n",
    "        self.projs3 = [nn.Linear(self.n_neurons, self.n_neurons_proj, bias = False)]\n",
    "        self.projs3 += [nn.Linear(self.n_neurons_proj, self.n_neurons_proj, bias = False) for k in range(self.n_layers_proj - 1)]\n",
    "        self.projs3 = nn.Sequential(*self.projs3)\n",
    "        \n",
    "        self.projs4 = [nn.Linear(self.n_neurons, self.n_neurons_proj, bias = False)]\n",
    "        self.projs4 += [nn.Linear(self.n_neurons_proj, self.n_neurons_proj, bias = False) for k in range(self.n_layers_proj - 1)]\n",
    "        self.projs4 = nn.Sequential(*self.projs4)\n",
    "        \n",
    "        self.projs = [self.projs1, self.projs2, self.projs3, self.projs4]\n",
    "        \n",
    "        self.f = nn.ReLU()\n",
    "        self.dropout = torch.nn.Dropout(p=0.5)\n",
    "        \n",
    "        #put model on GPU is available and asked\n",
    "        if args.device >= 0 and torch.cuda.is_available():\n",
    "           # device = torch.device(\"cuda:\"+str(args.device)+\")\")\n",
    "            device = torch.device(args.device)\n",
    "            self.cuda = True\n",
    "        else:\n",
    "            device = torch.device(\"cpu\")\n",
    "            self.cuda = False\n",
    "\n",
    "        self.device = device\n",
    "        self = self.to(device)\n",
    "\n",
    "        self.losses = [VicRegLoss(self.device, sim_coeff = ysim[idx], \n",
    "                                               std_coeff = yvar[idx], \n",
    "                                               cov_coeff = ycovar[idx],\n",
    "                                               intra_std_coeff = yintravar[idx]) \n",
    "                                               for idx in range(self.n_layers)]\n",
    "        \n",
    "        self.optimizer = torch.optim.Adam(self.parameters(), lr=args.lr, betas=(0.9, 0.999))\n",
    "        self.scheduler = torch.optim.lr_scheduler.ExponentialLR(self.optimizer, gamma=args.gamma_optimizer, last_epoch=-1, verbose=False)\n",
    "    \n",
    "    def forward(self, x, train_layer, perturbation):\n",
    "        '''\n",
    "         for the forward pass, we store the gradient of each layer, but we use detach() between each layer so the computational graph \n",
    "         is only at the layer-level\n",
    "         and we store the activations before the normalization\n",
    "         we send the layer-norm activation vector to the next layer (we normalize along the 2 axis)\n",
    "         we compute the layer-wise loss at each layer during the forward pass that allows us to use the .detach() features\n",
    "         the two views of the input are treated in the same mini-batch and we split the mini-batch for computing the SSL loss\n",
    "        '''\n",
    "        \n",
    "        loss = 0\n",
    "        for idx, fc in enumerate(self.layers):\n",
    "            if idx != train_layer: #if we don't train that layer with SPSA, we use the standard MVM of pytorch\n",
    "                # 1. compute forward pass for every layer\n",
    "                x = self.f(fc(x)) \n",
    "            \n",
    "            # 2. compute the SSL loss of that layer\n",
    "            elif idx == train_layer:\n",
    "                #1. first do the MVM + non-linearity\n",
    "                x = self.f(torch.matmul(x, (self.layers[idx].weight + perturbation).t()))\n",
    "                \n",
    "                # 3. feed to a 3 layers MLP (non-linear)\n",
    "                y = x\n",
    "                \n",
    "                for _, proj in enumerate(self.projs[idx][:-1]): # 4. iterate on the layers of the projector except the last one (that has no Relu)\n",
    "                    y = self.f(proj(y))\n",
    "                y = self.projs[idx][-1](y)\n",
    "                \n",
    "                y = self.multi_views_proj(y) # No Relu on the last layer\n",
    "                y0, y1 = y[0], y[1]\n",
    "                \n",
    "                loss_sim, loss_std, loss_cov = self.losses[idx](y0, y1) # 5. Compute the SSL loss\n",
    "                break #break if we have computed the loss: avoid to compute the next layers if we don't train them\n",
    "\n",
    "            x = x.detach() # 6. detach the activations vector to have a layer-wise computational graph here\n",
    "\n",
    "        return loss_sim.item(), loss_std.item(), loss_cov.item()\n",
    "    \n",
    "    \n",
    "    def forward_simple(self, x):\n",
    "        '''\n",
    "        Forward pass without computing the loss - \n",
    "        '''\n",
    "        loss = 0\n",
    "        states = []\n",
    "        for idx, fc in enumerate(self.layers):\n",
    "            x = self.f(fc(x)) \n",
    "            states.append(x)\n",
    "            x = x.detach()\n",
    "\n",
    "        return states\n",
    "    \n",
    "        \n",
    "    def single_batch(self, x):\n",
    "        '''\n",
    "        return a single big batch given the two views of the input data\n",
    "        '''\n",
    "        x = torch.stack(x) #here data as a new first dimension which is the number of views for the same input image\n",
    "        x = x.view(x.size()[0]*x.size()[1], -1) #we change the first dim to be n_views*batch_size\n",
    "        \n",
    "        return x\n",
    "    \n",
    "    \n",
    "    def multi_views(self, x):\n",
    "        '''\n",
    "        return a view of the tensor that has first dim n_views, second dim batch_size and third dim the layers dimension\n",
    "        '''\n",
    "        \n",
    "        return x.view(self.n_views, -1, self.n_neurons)\n",
    "\n",
    "\n",
    "    def multi_views_proj(self, x):\n",
    "        '''\n",
    "        return a view of the tensor that has first dim n_views, second dim batch_size and third dim the layers dimension\n",
    "        '''\n",
    "        \n",
    "        return x.view(self.n_views, -1, self.n_neurons_proj)\n",
    "    \n",
    "    \n",
    "    def reset_perturbation_step(self):\n",
    "        '''\n",
    "        Reset to initial_step the perturbation step\n",
    "        '''\n",
    "        \n",
    "        self.step = self.initial_step\n",
    "        \n",
    "        return 0\n",
    "    \n",
    "    \n",
    "    def update_perturbation_step(self):\n",
    "        '''\n",
    "        Do one step of exponential decay of the perturbation step\n",
    "        '''\n",
    "        \n",
    "        self.step *= self.gamma_perturbation\n",
    "        \n",
    "        return 0\n",
    "    \n",
    "    \n",
    "    def generate_perturbation(self, layer = 0):\n",
    "        '''\n",
    "        Generate a random perturbation \"vector\" - in fact a matrix - for the weights indicated by the layer index\n",
    "        1. Generate a random vector of -1/1\n",
    "        2. Scale down the vector by self.step\n",
    "        '''\n",
    "        \n",
    "        perturb = torch.randint(low=0, high = 2, size = self.layers[layer].weight.size()) #randint between 0 and 1 - 2 is exclusive\n",
    "        perturb = self.step * (2*(perturb-0.5)) #scale the perturbation vector to -1/+1 and scale down by the current step value\n",
    "        \n",
    "    \n",
    "        return perturb.to(self.device)\n",
    "    \n",
    "    \n",
    "    def compute_spsa(self, perturbation, pos_obj, neg_obj, layer = 0):\n",
    "        '''\n",
    "        Compute the gradient for all the parameters with the SPSA rule\n",
    "        g(obj|theta)=(obj(theta+ pertubation)-obj(theta-perturbation))/perturbation (SPSA)\n",
    "        theta = theta - eta*g(obj|theta) (SGD)\n",
    "        '''\n",
    "        \n",
    "        grad = (pos_obj-neg_obj)*(torch.ones(perturbation.size()).to(self.device)) #numerator - scaled to the size of the weight matrix\n",
    "        grad /= perturbation\n",
    "        \n",
    "        assert grad.size() == self.layers[layer].weight.size(), \"the size of the gradient computed does not match the size of the weight matrix\"\n",
    "            \n",
    "        return grad\n",
    "    \n",
    "    \n",
    "    def apply_spsa(self, grads, layer = 0):\n",
    "        '''\n",
    "        Apply the averaged gradients prescribed by SPSA to the respective weights\n",
    "        '''\n",
    "        #add -grad because we compute the gradient of the objective but we want -grad? maybe not - pytorch return the grad, the optimier then do -grad?\n",
    "        self.layers[layer].weight.grad = grads #fill in the grad attribute of the specific weight to use the optimizer with the integrated scheduler for the learning rate\n",
    "    \n",
    "        return 0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate simulations environment + data folders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BASE_PATH = createPath(archi = \"MLP-SPSA-Stochastic-Objective\", dataset = \"MNIST\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "saveHyperparameters(args, BASE_PATH)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dataframe = initDataframe_pretraining(BASE_PATH, dataframe_to_init = 'pre_training.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "ysim, yvar, ycovar, yintravar = param_sigmoid(args, plot = True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "net = Network(args, ysim, yvar, ycovar, yintravar)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "net, pretraining_loss = pretraining_loop(BASE_PATH, args, net, train_loader_pretrain, train_loader_classifier, test_loader, epochs = args.epochs)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python [conda env:NN_v2]",
   "language": "python",
   "name": "conda-env-NN_v2-py"
  },
  "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.7.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}