Layer-wise SSL - MLP - non-linear projectors.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 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')\n",
    "# general params\n",
    "parser.add_argument(\n",
    "    '--device',\n",
    "    type=int,\n",
    "    default=0,\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=4,\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=256,\n",
    "    help='Batch size for pre-training (default=256)')\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",
    "# 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=[40, 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=[1.5,0,2,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=[4,3,2,2],\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=[10,15,0,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",
    "        self.base_transforms = base_transforms #random transformations\n",
    "        self.n_views = n_views # number of differents copies with different \n",
    "\n",
    "    def __call__(self, x):\n",
    "        return [self.base_transforms(x) for i in range(self.n_views)]"
   ]
  },
  {
   "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 loadr 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",
    "                                target_transform=ReshapeTransformTarget(10)), batch_size = args.batchSize_pretrain, shuffle=True)\n",
    "\n",
    "    # Train loadr 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",
    "                                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",
    "\n",
    "data, target = next(iter(train_loader_pretrain))\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(data[0][0].view(28,28), cmap = \"gray\")\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(data[1][0].view(28,28), cmap = \"gray\")\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(data[0][0].view(28,28) - data[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",
    "        stacked_views = torch.stack((x,y)) #we add the view dimension as the first dimension\n",
    "        intra_std = torch.sqrt(stacked_views.var(dim=0) + 0.0001)\n",
    "        intra_std_loss = torch.mean(F.relu(self.intra_std_target - intra_std)) #the target variance has to be variable\n",
    "                                                                                 # we can also set it to 1 and decrease the relative strenght of that loss\n",
    "        loss = (\n",
    "            self.sim_coeff * repr_loss\n",
    "            + self.std_coeff * std_loss\n",
    "            + self.cov_coeff * cov_loss\n",
    "            + self.intra_std_coeff * intra_std_loss\n",
    "        )\n",
    "        return 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.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",
    "        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",
    "        if args.device >= 0 and torch.cuda.is_available():\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",
    "    \n",
    "    \n",
    "    def forward(self, x, train_layer):\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",
    "        states = []\n",
    "        for idx, fc in enumerate(self.layers):\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",
    "            if idx == train_layer:\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 += 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 states, loss\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)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate simulations environment + data folders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BASE_PATH = createPath(archi = \"MLP\", 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
}