Layer-wise SSL - CNN - non-linear projectors - CIFAR - asymetric augmentations.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Layer-wise Self-supervised training of a 4-layers CNN on CIFAR-10 with asymetric data-augmentations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\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 PIL import ImageOps, ImageFilter\n",
    "\n",
    "from utils_conv_asym 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 CNN 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='cifar',\n",
    "    help='Dataset we use for training (default=cifar, 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",
    "    '--channels',\n",
    "    nargs='+',\n",
    "    type=float,\n",
    "    default=[128, 256, 512, 1024],\n",
    "    help='Number of channels per layer (default=(64, 128 (MNIST)))')\n",
    "parser.add_argument(\n",
    "    '--kernel_size',\n",
    "    type=int,\n",
    "    default=5,\n",
    "    help='Size of the convolutional kernel (default: 5)')\n",
    "parser.add_argument(\n",
    "    '--padding',\n",
    "    type=int,\n",
    "    default=2,\n",
    "    help='Padding applied to each convolution (default: 2)')\n",
    "parser.add_argument(\n",
    "    '--nlayers_proj',\n",
    "    type=int,\n",
    "    default=2,\n",
    "    help='Number of layers for each non-linear projector (default: 3)')\n",
    "parser.add_argument(\n",
    "    '--nneurons_proj',\n",
    "    type=int,\n",
    "    default=2048,\n",
    "    help='Number of neuron per layer of each non-linear projector - also used to define the size of the 1x1 convolutional kernel (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",
    "parser.add_argument(\n",
    "    '--dropout',\n",
    "    type=int,\n",
    "    default=0,\n",
    "    help='Apply or not Dropout during training (default: 0)')\n",
    "\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",
    "    '--phflip',\n",
    "    type=float,\n",
    "    default=0.5,\n",
    "    help='Value of the probability for applying a RandomHorizontalFlip (default=0.5)')\n",
    "parser.add_argument(\n",
    "    '--pad',\n",
    "    type=int,\n",
    "    default=4,\n",
    "    help='Padding applied before RandomCrop (default=2)')\n",
    "parser.add_argument(\n",
    "    '--brightness',\n",
    "    type=float,\n",
    "    default=0.4,\n",
    "    help='Brightness value for ColorJittering (default=0.5)')\n",
    "parser.add_argument(\n",
    "    '--contrast',\n",
    "    type=float,\n",
    "    default=0.4,\n",
    "    help='Contrast value for ColorJittering (default=0.5)')\n",
    "parser.add_argument(\n",
    "    '--saturation',\n",
    "    type=float,\n",
    "    default=0.2,\n",
    "    help='Saturation value for ColorJittering (default=0.5)')\n",
    "parser.add_argument(\n",
    "    '--hue',\n",
    "    type=float,\n",
    "    default=0.1,\n",
    "    help='Hue value for ColorJittering (default=0.5)')\n",
    "parser.add_argument(\n",
    "    '--scaleaffine',\n",
    "    nargs='+',\n",
    "    type=float,\n",
    "    default=[0.8, 1.3],\n",
    "    help='Scale parameters for the RandomAffine transform applied (default=(0.7, 1.3))')\n",
    "parser.add_argument(\n",
    "    '--pgrayscale1',\n",
    "    type=float,\n",
    "    default=0.2,\n",
    "    help='Value of the probability for applying a RandomGrayscale (default=0.2)')\n",
    "parser.add_argument(\n",
    "    '--psolarize1',\n",
    "    type=float,\n",
    "    default=0.0,\n",
    "    help='Value of the probability for applying a RandomSolarize (default=0.1)')\n",
    "parser.add_argument(\n",
    "    '--pgaussianblur1',\n",
    "    type=float,\n",
    "    default=1.0,\n",
    "    help='Value of the probability for applying a Gaussiann blur branch 1 (default=1.0)')\n",
    "parser.add_argument(\n",
    "    '--pgrayscale2',\n",
    "    type=float,\n",
    "    default=0.2,\n",
    "    help='Value of the probability for applying a RandomGrayscale (default=0.2)')\n",
    "parser.add_argument(\n",
    "    '--psolarize2',\n",
    "    type=float,\n",
    "    default=0.2,\n",
    "    help='Value of the probability for applying a RandomSolarize (default=0.1)')\n",
    "parser.add_argument(\n",
    "    '--pgaussianblur2',\n",
    "    type=float,\n",
    "    default=0.1,\n",
    "    help='Value of the probability for applying a Gaussiann blur branch 2 (default=0.1)')\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,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=[3,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=[15,15,0,0],\n",
    "    help='Bias value for the sigmoidal parametrization (default=(10, 20, 1, 0))')\n",
    "\n",
    "\n",
    "args, _ = parser.parse_known_args()\n",
    "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_transforms1, base_transforms2, n_views=2):\n",
    "        self.base_transforms1 = base_transforms1 #random transformations branch 1\n",
    "        self.base_transforms2 = base_transforms2 #random transformations branch 2\n",
    "        self.n_views = n_views # number of differents copies with different \n",
    "\n",
    "    def __call__(self, x):\n",
    "        return self.base_transforms1(x), self.base_transforms2(x)\n",
    "    \n",
    "    \n",
    "class GaussianBlur(object):\n",
    "    def __init__(self, p):\n",
    "        self.p = p\n",
    "\n",
    "    def __call__(self, img):\n",
    "        if np.random.rand() < self.p:\n",
    "            sigma = np.random.rand() * 1.9 + 0.1\n",
    "            return img.filter(ImageFilter.GaussianBlur(sigma))\n",
    "        else:\n",
    "            return img"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Dataset creation\n",
    "def get_dataloader_cifar(args):\n",
    "    '''\n",
    "    Function that returns the dataloaders given the current hyperparameters \n",
    "    '''\n",
    "    # random data augmentations for the pre-training stage - branch 1\n",
    "    contrast_transforms1 =  transforms.Compose([transforms.RandomCrop(size=[32,32], padding = args.pad, padding_mode='edge'),\n",
    "                                            torchvision.transforms.RandomHorizontalFlip(p=args.phflip),\n",
    "                                            torchvision.transforms.RandomAffine(degrees=(0, 0), translate=(0.0, 0.0), scale=(args.scaleaffine[0], args.scaleaffine[1])),\n",
    "                                            transforms.ColorJitter(brightness=args.brightness, contrast=args.contrast, saturation=args.saturation, hue=args.hue),\n",
    "                                            torchvision.transforms.RandomApply([transforms.Grayscale(num_output_channels=3)], p=args.pgrayscale1),\n",
    "                                            GaussianBlur(p=args.pgaussianblur1),\n",
    "                                            torchvision.transforms.RandomSolarize(0.0, p=args.psolarize1),\n",
    "                                            transforms.ToTensor(),\n",
    "                                            transforms.Normalize((0.4914,0.4822,0.4465),(0.247,0.243,0.261))])\n",
    "    \n",
    "    # random data augmentations for the pre-training stage\n",
    "    contrast_transforms2 =  transforms.Compose([transforms.RandomCrop(size=[32,32], padding = args.pad, padding_mode='edge'),\n",
    "                                            torchvision.transforms.RandomHorizontalFlip(p=args.phflip),\n",
    "                                            torchvision.transforms.RandomAffine(degrees=(0, 0), translate=(0.0, 0.0), scale=(args.scaleaffine[0], args.scaleaffine[1])),\n",
    "                                            transforms.ColorJitter(brightness=args.brightness, contrast=args.contrast, saturation=args.saturation, hue=args.hue),\n",
    "                                            torchvision.transforms.RandomApply([transforms.Grayscale(num_output_channels=3)], p=args.pgrayscale2),\n",
    "                                            GaussianBlur(p=args.pgaussianblur2),\n",
    "                                            torchvision.transforms.RandomSolarize(0.0, p=args.psolarize2),\n",
    "                                            transforms.ToTensor(),\n",
    "                                            transforms.Normalize((0.4914,0.4822,0.4465),(0.247,0.243,0.261))])\n",
    "\n",
    "\n",
    "    transforms_train_class = transforms.Compose([transforms.RandomCrop(size=[32,32], padding = 1, padding_mode='edge'),\n",
    "                                            transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.1, hue=0.1),\n",
    "                                            torchvision.transforms.RandomHorizontalFlip(p=0.2),  \n",
    "                                            torchvision.transforms.RandomAffine(degrees=(0, 0), translate=(0.0, 0.0), scale=(1, 1.2)),\n",
    "                                            transforms.ToTensor(),\n",
    "                                            transforms.Normalize((0.4914,0.4822,0.4465),(0.247,0.243,0.261))])\n",
    "\n",
    "    # fixed transformation for training testing the linear classifier\n",
    "    transforms_test = transforms.Compose([transforms.ToTensor(),\n",
    "                                        transforms.Normalize((0.4914,0.4822,0.4465),(0.247,0.243,0.261))])\n",
    "\n",
    "    # Train loader for pre-training\n",
    "    train_loader = torch.utils.data.DataLoader(\n",
    "    torchvision.datasets.CIFAR10(root='./data', train=True, download=True,\n",
    "                                transform = ContrastiveTransformations(contrast_transforms1, contrast_transforms2, n_views=args.nviews),\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.CIFAR10(root='./data', train=False, download=True,\n",
    "                                transform = ContrastiveTransformations(contrast_transforms1, contrast_transforms2, n_views=args.nviews),\n",
    "                                target_transform=ReshapeTransformTarget(10)), batch_size = args.batchSize_pretrain, shuffle=True)\n",
    "\n",
    "    # Train loader for the linear classifier\n",
    "    train_loader_classifier = torch.utils.data.DataLoader(torchvision.datasets.CIFAR10(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.CIFAR10(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_cifar(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(data[1].size())\n",
    "plt.figure()\n",
    "plt.imshow(data[0][0][0].view(32,32), cmap = \"gray\")\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(data[1][0][0].view(32,32), cmap = \"gray\")\n",
    "\n",
    "plt.figure()\n",
    "plt.imshow(data[0][0][0].view(32,32) - data[1]\n",
    "           [0][0].view(32,32), 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",
    "- CNN with Relu + BN + MaxPool\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",
    "import torch.nn as nn\n",
    "\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",
    "        if args.dataset == \"mnist\":\n",
    "            self.channels_list = [1] + args.channels\n",
    "            self.size_feature_map = [28] \n",
    "        elif args.dataset == \"cifar\":\n",
    "            self.channels_list = [3] + args.channels\n",
    "            self.size_feature_map = [32] \n",
    "\n",
    "        self.n_views = args.nviews\n",
    "        self.n_layers = args.nlayers\n",
    "        \n",
    "        self.n_neurons_proj = args.nneurons_proj\n",
    "        self.n_layers_proj = args.nlayers_proj\n",
    "\n",
    "        for k in range(self.n_layers):\n",
    "            self.size_feature_map.append(int(((self.size_feature_map[-1]-args.kernel_size+2*args.padding)/1+1)/2))\n",
    "        print(self.size_feature_map)\n",
    "        \n",
    "        kernel_size = args.kernel_size\n",
    "        pad = args.padding\n",
    "        \n",
    "        self.layers = [nn.Conv2d(self.channels_list[k], self.channels_list[k+1], kernel_size, padding = pad, stride=1) for k in range(self.n_layers)]\n",
    "        self.layers = nn.Sequential(*self.layers)\n",
    "        \n",
    "        self.bn = nn.ModuleList([nn.BatchNorm2d(self.channels_list[layer+1]) for layer in range(len(self.channels_list)-1)])\n",
    "                \n",
    "        #conv-base expansion\n",
    "        self.exp_conv = [nn.Conv2d(self.channels_list[k+1], self.n_neurons_proj, 1, padding = 0, stride=1) for k in range(self.n_layers)]\n",
    "        self.exp_conv = nn.Sequential(*self.exp_conv)\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_proj, self.n_neurons_proj, bias = False) for k in range(self.n_layers_proj)]\n",
    "        self.projs1 = nn.Sequential(*self.projs1)\n",
    "\n",
    "        self.projs2 = [nn.Linear(self.n_neurons_proj, self.n_neurons_proj, bias = False) for k in range(self.n_layers_proj)]\n",
    "        self.projs2 = nn.Sequential(*self.projs2)\n",
    "        \n",
    "        self.projs3 = [nn.Linear(self.n_neurons_proj, self.n_neurons_proj, bias = False) for k in range(self.n_layers_proj)]\n",
    "        self.projs3 = nn.Sequential(*self.projs3)\n",
    "\n",
    "        self.projs4 = [nn.Linear(self.n_neurons_proj, self.n_neurons_proj, bias = False) for k in range(self.n_layers_proj)]\n",
    "        self.projs4 = nn.Sequential(*self.projs4)\n",
    "        \n",
    "        self.projs = [self.projs1, self.projs2, self.projs3, self.projs4]\n",
    "        \n",
    "        \n",
    "        self.f = nn.ReLU()\n",
    "        self.dropout = args.dropout\n",
    "        self.dropout_f = torch.nn.Dropout(p=0.2)\n",
    "        self.max_pool = nn.MaxPool2d(2)\n",
    "        self.flatten = nn.Flatten()\n",
    "        \n",
    "        #put model on GPU is available and asked\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, x1, x2, 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",
    "        loss = 0\n",
    "        states = []\n",
    "        for idx, conv in enumerate(self.layers): #we compute the forward pass up to the layer we want to train!\n",
    "            #1. compute forward pass for every layer\n",
    "            x1 = self.bn[idx](self.f(conv(x1)))\n",
    "            x2 = self.bn[idx](self.f(conv(x2)))\n",
    "            \n",
    "            if self.dropout:\n",
    "                x1 = self.dropout_f(x1)\n",
    "                x2 = self.dropout_f(x2)\n",
    "                \n",
    "            x1 = self.max_pool(x1)\n",
    "            x2 = self.max_pool(x2)\n",
    "\n",
    "            # 2. compute the BT loss of that layer - loss function from github - no need a complex code, the matmul operation on the two emebeddings vectors is sufficient as the resulting matrix is the sum of the correlation matrices!\n",
    "            if idx == train_layer:\n",
    "                #1x1 conv + averaged pooling onn each channel\n",
    "                #print(x.size())\n",
    "                y1 = self.exp_conv[idx](x1).sum(-1).sum(-1)\n",
    "                y2 = self.exp_conv[idx](x2).sum(-1).sum(-1)\n",
    "\n",
    "                # feed to a 3 layers MLP (non-linear)\n",
    "                for idx, fc in enumerate(self.projs[idx][:-1]):\n",
    "                    y1 = self.f(fc(y1))\n",
    "                    y2 = self.f(fc(y2))\n",
    "                    \n",
    "                y1 = self.projs[idx][-1](y1)#no Relu on the last layer\n",
    "                y2 = self.projs[idx][-1](y2)\n",
    "                \n",
    "                #y = self.multi_views_fc(y)\n",
    "                #y0, y1 = y[0], y[1]\n",
    "                \n",
    "                loss += self.losses[idx](y1, y2)\n",
    "                break\n",
    "\n",
    "            x1 = x1.detach() # detach to stop the computational graph here\n",
    "            x2 = x2.detach()\n",
    "            \n",
    "        return states, loss\n",
    "    \n",
    "    \n",
    "    def forward_simple(self, x):\n",
    "        '''\n",
    "        Forward pass without computing the loss - \n",
    "        '''\n",
    "        states = []\n",
    "        for idx, conv in enumerate(self.layers):\n",
    "           #1. compute forward pass for every layer\n",
    "            x = self.f(conv(x))\n",
    "            x = self.bn[idx](x)\n",
    "            x = self.max_pool(x)\n",
    "            states.append(x.clone()) #store the states after the max-pool operation!\n",
    "            x = x.detach() # detach to stop the computational graph here\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",
    "        s = x.size()\n",
    "        x = x.view(s[0]*s[1], s[2], s[3], s[4]) #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",
    "        s = x.size()\n",
    "        return x.view(self.n_views, int(s[0]/self.n_views), s[1], s[2], s[3])\n",
    "    \n",
    "    \n",
    "    def multi_views_fc(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",
    "        s = x.size()\n",
    "        return x.view(self.n_views, int(s[0]/self.n_views), -1)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Generate simulations environment + data folders"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "BASE_PATH = createPath(archi = \"CNN\", dataset = \"CIFAR\")"
   ]
  },
  {
   "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
}