extra_loss_CIFAR_large_8_255.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "button": false,
    "colab": {},
    "colab_type": "code",
    "collapsed": true,
    "deletable": true,
    "id": "UoK0sWEypQ-s",
    "new_sheet": false,
    "run_control": {
     "read_only": false
    }
   },
   "outputs": [],
   "source": [
    "#%matplotlib inline\n",
    "%reload_ext autoreload\n",
    "%autoreload 2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "button": false,
    "colab": {},
    "colab_type": "code",
    "collapsed": true,
    "deletable": true,
    "id": "Osqh5JlNpQ-2",
    "new_sheet": false,
    "run_control": {
     "read_only": false
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "from torch.utils.data import DataLoader\n",
    "import torch.utils.data as data_utils\n",
    "#from utils import epoch, epoch_robust_bound, epoch_calculate_robust_err, Flatten, generate_kappa_schedule_CIFAR, generate_epsilon_schedule_CIFAR"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "button": false,
    "colab_type": "text",
    "deletable": true,
    "id": "iuWv5TZpukKi",
    "new_sheet": false,
    "run_control": {
     "read_only": false
    }
   },
   "source": [
    "# Utils and helper functions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "button": false,
    "colab": {},
    "colab_type": "code",
    "collapsed": true,
    "deletable": true,
    "id": "0kiV7tfS5c5y",
    "new_sheet": false,
    "run_control": {
     "read_only": false
    }
   },
   "outputs": [],
   "source": [
    "class Flatten(nn.Module):\n",
    "    def forward(self, x):\n",
    "        return x.view(x.shape[0], -1) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "button": false,
    "colab": {},
    "colab_type": "code",
    "collapsed": true,
    "deletable": true,
    "id": "1XdcOwIVuv7r",
    "new_sheet": false,
    "run_control": {
     "read_only": false
    }
   },
   "outputs": [],
   "source": [
    "\n",
    "\n",
    "def epoch(loader, model, device, opt=None):\n",
    "    \"\"\"Standard training/evaluation epoch over the dataset\"\"\"\n",
    "    total_loss, total_err = 0.,0.\n",
    "    for X,y in loader:\n",
    "        X,y = X.to(device), y.to(device)\n",
    "        yp,_ = model(X)\n",
    "        loss = nn.CrossEntropyLoss()(yp,y)\n",
    "        if opt:\n",
    "            opt.zero_grad()\n",
    "            loss.backward()\n",
    "            opt.step()\n",
    "        \n",
    "        total_err += (yp.max(dim=1)[1] != y).sum().item()\n",
    "        total_loss += loss.item() * X.shape[0]\n",
    "    return total_err / len(loader.dataset), total_loss / len(loader.dataset)\n",
    "\n",
    "\n",
    "def bound_propagation(model, initial_bound):\n",
    "    l, u = initial_bound\n",
    "    bounds = []\n",
    "    bounds.append(initial_bound)\n",
    "    list_of_layers = list(model.children())\n",
    "    \n",
    "    for i in range(len(list_of_layers)):\n",
    "        layer = list_of_layers[i]\n",
    "        \n",
    "        if isinstance(layer, Flatten):\n",
    "            l_ = Flatten()(l)\n",
    "            u_ = Flatten()(u)\n",
    "\n",
    "        elif isinstance(layer, nn.Linear):\n",
    "            l_ = (layer.weight.clamp(min=0) @ l.t() + layer.weight.clamp(max=0) @ u.t() \n",
    "                  + layer.bias[:,None]).t()\n",
    "            u_ = (layer.weight.clamp(min=0) @ u.t() + layer.weight.clamp(max=0) @ l.t() \n",
    "                  + layer.bias[:,None]).t()\n",
    "            \n",
    "        elif isinstance(layer, nn.Conv2d):\n",
    "            l_ = (nn.functional.conv2d(l, layer.weight.clamp(min=0), bias=None, \n",
    "                                       stride=layer.stride, padding=layer.padding,\n",
    "                                       dilation=layer.dilation, groups=layer.groups) +\n",
    "                  nn.functional.conv2d(u, layer.weight.clamp(max=0), bias=None, \n",
    "                                       stride=layer.stride, padding=layer.padding,\n",
    "                                       dilation=layer.dilation, groups=layer.groups) +\n",
    "                  layer.bias[None,:,None,None])\n",
    "            \n",
    "            u_ = (nn.functional.conv2d(u, layer.weight.clamp(min=0), bias=None, \n",
    "                                       stride=layer.stride, padding=layer.padding,\n",
    "                                       dilation=layer.dilation, groups=layer.groups) +\n",
    "                  nn.functional.conv2d(l, layer.weight.clamp(max=0), bias=None, \n",
    "                                       stride=layer.stride, padding=layer.padding,\n",
    "                                       dilation=layer.dilation, groups=layer.groups) + \n",
    "                  layer.bias[None,:,None,None])\n",
    "            \n",
    "        elif isinstance(layer, nn.ReLU):\n",
    "            l_ = l.clamp(min=0)\n",
    "            u_ = u.clamp(min=0)\n",
    "            \n",
    "        bounds.append((l_, u_))\n",
    "        l,u = l_, u_\n",
    "    return bounds\n",
    "\n",
    "\n",
    "def interval_based_bound(model, c, bounds, idx):\n",
    "    # requires last layer to be linear\n",
    "    cW = c.t() @ model.last_linear.weight\n",
    "    cb = c.t() @ model.last_linear.bias\n",
    "    \n",
    "    l,u = bounds[-2]\n",
    "    return (cW.clamp(min=0) @ l[idx].t() + cW.clamp(max=0) @ u[idx].t() + cb[:,None]).t()\n",
    "\n",
    "\n",
    "def epoch_robust_bound(loader, model, epsilon_schedule, device, kappa_schedule, batch_counter, opt=None):\n",
    "    robust_err = 0\n",
    "    total_robust_loss = 0\n",
    "    total_mse_loss = 0\n",
    "    total_combined_loss = 0\n",
    "    \n",
    "    C = [-torch.eye(10).to(device) for _ in range(10)]\n",
    "    for y0 in range(10):\n",
    "        C[y0][y0,:] += 1\n",
    "    \n",
    "    for i,data in enumerate(loader,0):\n",
    "      \n",
    "        if i>299:  #calculate only for 300 batches\n",
    "        break      \n",
    "        \n",
    "        mse_loss_list = []\n",
    "        lower_bounds = []\n",
    "        upper_bounds = []\n",
    "        \n",
    "        \n",
    "        X,y = data\n",
    "        X,y = X.to(device), y.to(device)\n",
    "        \n",
    "        ###### fit loss calculation ######\n",
    "        yp,_ = model(X)\n",
    "        fit_loss = nn.CrossEntropyLoss()(yp,y)\n",
    "    \n",
    "        ###### robust loss calculation ######\n",
    "        initial_bound = (X - epsilon_schedule[batch_counter], X + epsilon_schedule[batch_counter])\n",
    "        bounds = bound_propagation(model, initial_bound)\n",
    "        robust_loss = 0\n",
    "        for y0 in range(10):\n",
    "            if sum(y==y0) > 0:\n",
    "                lower_bound = interval_based_bound(model, C[y0], bounds, y==y0)\n",
    "                robust_loss += nn.CrossEntropyLoss(reduction='sum')(-lower_bound, y[y==y0]) / X.shape[0]\n",
    "                \n",
    "                robust_err += (lower_bound.min(dim=1)[0] < 0).sum().item() #increment when true label is not winning       \n",
    "        \n",
    "        total_robust_loss += robust_loss.item() * X.shape[0]  \n",
    "        \n",
    "        ##### MSE Loss #####\n",
    "        \n",
    "        #indices_of_layers = [2,4,7,8] #CNN_small\n",
    "        indices_of_layers = [2,4,6,8,10,13,14] #CNN_medium\n",
    "        \n",
    "        \n",
    "        for i in range(len(indices_of_layers)):\n",
    "            lower_bounds.append(Flatten()(bounds[indices_of_layers[i]][0])) #lower bounds \n",
    "            upper_bounds.append(Flatten()(bounds[indices_of_layers[i]][1])) #upper bounds \n",
    "            mse_loss_list.append(nn.MSELoss()(lower_bounds[i], upper_bounds[i]))\n",
    "            \n",
    "        \n",
    "        mse_loss = mse_loss_list[0] + mse_loss_list[1] + mse_loss_list[2] + mse_loss_list[3] + mse_loss_list[4] + mse_loss_list[5] + mse_loss_list[6]\n",
    "        total_mse_loss += mse_loss.item()\n",
    "        \n",
    "        ###### combined losss ######\n",
    "        combined_loss = kappa_schedule[batch_counter]*fit_loss + (1-kappa_schedule[batch_counter])*robust_loss + mse_loss\n",
    "\n",
    "        total_combined_loss += combined_loss.item()\n",
    "        \n",
    "        batch_counter +=1\n",
    "         \n",
    "        if opt:\n",
    "            opt.zero_grad()\n",
    "            combined_loss.backward()\n",
    "            opt.step() \n",
    "        \n",
    "    return robust_err / len(loader.dataset), total_combined_loss / len(loader.dataset), total_mse_loss/ len(loader.dataset)\n",
    "\n",
    "        \n",
    "def epoch_calculate_robust_err (loader, model, epsilon, device):\n",
    "    robust_err = 0.0\n",
    "    \n",
    "    C = [-torch.eye(10).to(device) for _ in range(10)]\n",
    "    for y0 in range(10):\n",
    "        C[y0][y0,:] += 1\n",
    "\n",
    "\n",
    "    for X,y in loader:\n",
    "        X,y = X.to(device), y.to(device)\n",
    "        \n",
    "        initial_bound = (X - epsilon, X + epsilon)\n",
    "        bounds = bound_propagation(model, initial_bound)\n",
    "\n",
    "        for y0 in range(10):\n",
    "            if sum(y==y0) > 0:\n",
    "                lower_bound = interval_based_bound(model, C[y0], bounds, y==y0)                \n",
    "                robust_err += (lower_bound.min(dim=1)[0] < 0).sum().item() #increment when true label is not winning       \n",
    "        \n",
    "    return robust_err / len(loader.dataset)\n",
    "        \n",
    "        \n",
    "\n",
    "\n",
    "def generate_kappa_schedule_MNIST():\n",
    "\n",
    "    kappa_schedule = 2000*[1] # warm-up phase\n",
    "    kappa_value = 1.0\n",
    "    step = 0.5/58000\n",
    "    \n",
    "    for i in range(58000):\n",
    "        kappa_value = kappa_value - step\n",
    "        kappa_schedule.append(kappa_value)\n",
    "    \n",
    "    return kappa_schedule\n",
    "\n",
    "def generate_epsilon_schedule_MNIST(epsilon_train):\n",
    "    \n",
    "    epsilon_schedule = []\n",
    "    step = epsilon_train/10000\n",
    "            \n",
    "    for i in range(10000):\n",
    "        epsilon_schedule.append(i*step) #ramp-up phase\n",
    "    \n",
    "    for i in range(50000):\n",
    "        epsilon_schedule.append(epsilon_train)\n",
    "        \n",
    "    return epsilon_schedule\n",
    "\n",
    "\n",
    "def generate_kappa_schedule_CIFAR():\n",
    "\n",
    "    kappa_schedule = 10000*[1] # warm-up phase\n",
    "    kappa_value = 1.0\n",
    "    step = 0.5/340000\n",
    "    \n",
    "    for i in range(340000):\n",
    "        kappa_value = kappa_value - step\n",
    "        kappa_schedule.append(kappa_value)\n",
    "    \n",
    "    return kappa_schedule\n",
    "\n",
    "def generate_epsilon_schedule_CIFAR(epsilon_train):\n",
    "    \n",
    "    epsilon_schedule = []\n",
    "    step = epsilon_train/150000\n",
    "            \n",
    "    for i in range(150000):\n",
    "        epsilon_schedule.append(i*step) #ramp-up phase\n",
    "    \n",
    "    for i in range(200000):\n",
    "        epsilon_schedule.append(epsilon_train)\n",
    "        \n",
    "    return epsilon_schedule \n",
    "  \n",
    "  \n",
    "def pgd_linf_rand(model, X, y, epsilon, alpha, num_iter, restarts):\n",
    "    \"\"\" Construct PGD adversarial examples on the samples X, with random restarts\"\"\"\n",
    "    max_loss = torch.zeros(y.shape[0]).to(y.device)\n",
    "    max_delta = torch.zeros_like(X)\n",
    "    \n",
    "    for i in range(restarts):\n",
    "        delta = torch.rand_like(X, requires_grad=True)\n",
    "        delta.data = delta.data * 2 * epsilon - epsilon\n",
    "        \n",
    "        for t in range(num_iter):\n",
    "            loss = nn.CrossEntropyLoss()(model(X + delta)[0], y)\n",
    "            loss.backward()\n",
    "            delta.data = (delta + alpha*delta.grad.detach().sign()).clamp(-epsilon,epsilon)\n",
    "            delta.grad.zero_()\n",
    "        \n",
    "        all_loss = nn.CrossEntropyLoss(reduction='none')(model(X+delta)[0],y)\n",
    "        max_delta[all_loss >= max_loss] = delta.detach()[all_loss >= max_loss]\n",
    "        max_loss = torch.max(max_loss, all_loss)\n",
    "        \n",
    "    return max_delta\n",
    "\n",
    "\n",
    "def epoch_adversarial(model, loader, attack, *args):\n",
    "    total_loss, total_err = 0.,0.\n",
    "    for X,y in loader:\n",
    "        X,y = X.to(device), y.to(device)\n",
    "        delta = attack(model, X, y, *args)\n",
    "        yp = model(X+delta)[0]\n",
    "        loss = nn.CrossEntropyLoss()(yp,y)\n",
    "        \n",
    "        total_err += (yp.max(dim=1)[1] != y).sum().item()\n",
    "        total_loss += loss.item() * X.shape[0]\n",
    "    return total_err / len(loader.dataset), total_loss / len(loader.dataset)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "button": false,
    "colab_type": "text",
    "deletable": true,
    "id": "kfGlWGpuCn9t",
    "new_sheet": false,
    "run_control": {
     "read_only": false
    }
   },
   "source": [
    "# Losses"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "button": false,
    "colab": {},
    "colab_type": "code",
    "collapsed": true,
    "deletable": true,
    "id": "z-8Ucp1YCqWc",
    "new_sheet": false,
    "run_control": {
     "read_only": false
    }
   },
   "outputs": [],
   "source": [
    "#Root MSE\n",
    "\n",
    "def RMSELoss(yp,y):\n",
    "    return torch.sqrt(nn.MSELoss()(yp, y)+1e-6) #small epsilon is added to avoid NaN when mse=0  \n",
    "  \n",
    "#Log (Product of squared errors)\n",
    "\n",
    "def Log_Product_Loss(yp,y):\n",
    "  return torch.mean(torch.log((yp-y)**2+1e-6))\n",
    "\n",
    "#Log (Product of absolute errors)\n",
    "\n",
    "def Log_Product_MAE(yp,y):\n",
    "  return torch.sum(torch.log(torch.abs(yp-y)+1e-6))\n",
    "\n",
    "# Mean Absolute Error\n",
    "\n",
    "def MAELoss(yp,y):\n",
    "  return torch.mean(torch.abs((yp-y)))\n",
    "\n",
    "def DiffVolume(yp,y):\n",
    "  vol_lower_bound = torch.sum(torch.log(yp+1e-6))\n",
    "  vol_upper_bound = torch.sum(torch.log(y+1e-6))\n",
    "  diff = vol_upper_bound - vol_lower_bound\n",
    "  return diff\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "button": false,
    "colab_type": "text",
    "deletable": true,
    "id": "q7bW3yTVvMLW",
    "new_sheet": false,
    "run_control": {
     "read_only": false
    }
   },
   "source": [
    "# Loading dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "button": false,
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "collapsed": false,
    "deletable": true,
    "id": "jMEAJ22NpQ-9",
    "new_sheet": false,
    "outputId": "59930c05-dcc0-4661-dddf-5c50dba0dd94",
    "run_control": {
     "read_only": false
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<torch._C.Generator at 0x7f1368434090>"
      ]
     },
     "execution_count": 5,
     "metadata": {
      "tags": []
     },
     "output_type": "execute_result"
    }
   ],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "torch.manual_seed(0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "button": false,
    "colab": {},
    "colab_type": "code",
    "collapsed": true,
    "deletable": true,
    "id": "qqxdD66ymDYN",
    "new_sheet": false,
    "run_control": {
     "read_only": false
    }
   },
   "outputs": [],
   "source": [
    "BATCH_SIZE = 50\n",
    "dataset_path = './cifar10'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "button": false,
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 34
    },
    "colab_type": "code",
    "collapsed": false,
    "deletable": true,
    "id": "ZdtPZTAOmDYR",
    "new_sheet": false,
    "outputId": "d908902b-3e91-43ea-96ad-387c8148d3f5",
    "run_control": {
     "read_only": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using downloaded and verified file: ./svhn/train_32x32.mat\n"
     ]
    }
   ],
   "source": [
    "trainset = datasets.CIFAR10(root=dataset_path, train=True, download=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "button": false,
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 51
    },
    "colab_type": "code",
    "collapsed": false,
    "deletable": true,
    "id": "xMR9U4qRmDYW",
    "new_sheet": false,
    "outputId": "ea65f0d6-05a8-480d-d7be-073322fa2e95",
    "run_control": {
     "read_only": false
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using downloaded and verified file: ./svhn/train_32x32.mat\n",
      "Using downloaded and verified file: ./svhn/test_32x32.mat\n"
     ]
    }
   ],
   "source": [
    "train_mean = trainset.data.mean(axis=(0,1,2))/255  # [0.49139968  0.48215841  0.44653091]\n",
    "train_std = trainset.data.std(axis=(0,1,2))/255  # [0.24703223  0.24348513  0.26158784]\n",
    "\n",
    "transform_train = transforms.Compose([\n",
    "    transforms.RandomCrop(32, padding=4),\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(train_mean, train_std),\n",
    "])\n",
    "\n",
    "transform_test = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize(train_mean, train_std),\n",
    "])\n",
    "kwargs = {'pin_memory': True}\n",
    "\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(datasets.CIFAR10(\n",
    "    root=dataset_path, train=True, download=True,\n",
    "    transform=transform_train),\n",
    "    batch_size=BATCH_SIZE, shuffle=True, **kwargs)\n",
    "\n",
    "test_loader = torch.utils.data.DataLoader(\n",
    "    datasets.CIFAR10(root=dataset_path, train=False, download=True,\n",
    "    transform=transform_test),\n",
    "    batch_size=BATCH_SIZE, shuffle=False, **kwargs)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "button": false,
    "colab_type": "text",
    "deletable": true,
    "id": "8CJe2L46pQ_z",
    "new_sheet": false,
    "run_control": {
     "read_only": false
    }
   },
   "source": [
    "# Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "button": false,
    "colab": {},
    "colab_type": "code",
    "collapsed": true,
    "deletable": true,
    "id": "OHXddnGZIhdl",
    "new_sheet": false,
    "run_control": {
     "read_only": false
    }
   },
   "outputs": [],
   "source": [
    "class CNN_large(torch.nn.Module):\n",
    "    def __init__(self):\n",
    "\n",
    "        super(CNN_large, self).__init__()\n",
    "        \n",
    "        self.conv1 = nn.Conv2d(3, 64, 3, padding=0, stride=1)\n",
    "        self.relu1 = nn.ReLU() \n",
    "        self.conv2 = nn.Conv2d(64, 64, 3, padding=0, stride=1)\n",
    "        self.relu2 = nn.ReLU()\n",
    "        self.conv3 = nn.Conv2d(64, 128, 3, padding=0, stride=2)\n",
    "        self.relu3 = nn.ReLU()\n",
    "        self.conv4 = nn.Conv2d(128, 128, 3, padding=0, stride=1)\n",
    "        self.relu4 = nn.ReLU()\n",
    "        self.conv5 = nn.Conv2d(128, 128, 3, padding=0, stride=1)\n",
    "        self.relu5 = nn.ReLU()\n",
    "        \n",
    "        self.flat = Flatten()\n",
    "        self.linear1 = nn.Linear(128*9*9, 200)\n",
    "        self.relu6 = nn.ReLU()\n",
    "        self.last_linear = nn.Linear(200, 10)                \n",
    "\n",
    "        \n",
    "    def forward(self, x):\n",
    "        \n",
    "        hidden_activations = []\n",
    "        \n",
    "        x = self.conv1(x)\n",
    "        x = self.relu1(x)\n",
    "        \n",
    "        x = self.conv2(x)\n",
    "        x = self.relu2(x)\n",
    "        \n",
    "        x = self.conv3(x)\n",
    "        x = self.relu3(x)\n",
    "        \n",
    "        x = self.conv4(x)\n",
    "        x = self.relu4(x)\n",
    "       \n",
    "        x = self.conv5(x)\n",
    "        x = self.relu5(x)\n",
    "        \n",
    "        x = self.flat(x)\n",
    "\n",
    "        x = self.linear1(x)\n",
    "        x = self.relu6(x)\n",
    "        \n",
    "        out = self.last_linear(x)\n",
    "\n",
    "        return out, hidden_activations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "button": false,
    "colab": {},
    "colab_type": "code",
    "collapsed": true,
    "deletable": true,
    "id": "pczaC9Iwwy3C",
    "new_sheet": false,
    "run_control": {
     "read_only": false
    }
   },
   "outputs": [],
   "source": [
    "model = CNN_large().to(device)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "button": false,
    "colab_type": "text",
    "deletable": true,
    "id": "KBRQTG-zpQ_3",
    "new_sheet": false,
    "run_control": {
     "read_only": false
    }
   },
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "button": false,
    "colab": {},
    "colab_type": "code",
    "collapsed": true,
    "deletable": true,
    "id": "QgFHqVWQpQ_4",
    "new_sheet": false,
    "run_control": {
     "read_only": false
    }
   },
   "outputs": [],
   "source": [
    "opt = optim.Adam(model.parameters(), lr=1e-3)\n",
    "\n",
    "EPSILON = 8/255\n",
    "EPSILON_TRAIN = 8/255\n",
    "epsilon_schedule = generate_epsilon_schedule_CIFAR(EPSILON_TRAIN)\n",
    "kappa_schedule = generate_kappa_schedule_CIFAR()\n",
    "batch_counter = 0\n",
    "  \n",
    "print(\"Training started...\")\n",
    "for t in range(350):\n",
    "    _, combined_loss, mse_loss = epoch_robust_bound(train_loader, model, epsilon_schedule, device, kappa_schedule, batch_counter, opt)\n",
    "\n",
    "    # check loss and accuracy on test set\n",
    "    #test_err, _ = epoch(test_loader, model, device)\n",
    "    #robust_err = epoch_calculate_robust_err(test_loader, model, EPSILON, device)\n",
    "\n",
    "    batch_counter += 1000\n",
    "\n",
    "    if t == 200:  #decrease learning rate\n",
    "        for param_group in opt.param_groups:\n",
    "            param_group[\"lr\"] = 1e-4\n",
    "\n",
    "    if t == 250:  #decrease learning rate after 250 epochs\n",
    "        for param_group in opt.param_groups:\n",
    "            param_group[\"lr\"] = 1e-5\n",
    "\n",
    "    if t == 300:  #decrease learning rate after 300 epochs\n",
    "        for param_group in opt.param_groups:\n",
    "            param_group[\"lr\"] = 1e-6\n",
    "\n",
    "    print(\"Epoch \" + str(t) + \" done\\n\")\n",
    "\n",
    "#SAVING RESULTS\n",
    "with open('results_extra_loss_CIFAR_large_8_255.txt', 'w') as f:\n",
    "\n",
    "  test_err, _ = epoch(test_loader, model, device)\n",
    "  print (\"Test Error Rate: \" + str(test_err) + \"\\n\")\n",
    "  f.write(\"Test Error Rate: \" + str(test_err) + \"\\n\")\n",
    "\n",
    "  pgd_result = epoch_adversarial(model, test_loader, pgd_linf_rand, EPSILON, 1e-2, 200, 10)[0]\n",
    "  print (\"PGD Error Rate: \" + str(pgd_result) + \"\\n\")\n",
    "  f.write(\"PGD Error Rate: \" + str(pgd_result) + \"\\n\")\n",
    "\n",
    "  verified = epoch_calculate_robust_err(test_loader, model, EPSILON, device)\n",
    "  print (\"Verified: \" + str(verified) + \"\\n\")\n",
    "  f.write(\"Verified: \" + str(verified) + \"\\n\")\n",
    "\n",
    "f.close()\n",
    "\n",
    "#SAVING WEIGHTS\n",
    "torch.save(model.state_dict(), 'weights_extra_loss_CIFAR_large_8_255.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 0,
   "metadata": {
    "button": false,
    "colab": {},
    "colab_type": "code",
    "collapsed": true,
    "deletable": true,
    "id": "RUzUmBFFWJcv",
    "new_sheet": false,
    "run_control": {
     "read_only": false
    }
   },
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "extra_loss_CIFAR_large_8_255.ipynb",
   "provenance": [],
   "version": "0.3.2"
  },
  "kernelspec": {
   "display_name": "Python 3",
   "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.7.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}