demos/Train and Evaluation.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "bc4944a7",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "97ebc899",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "os.environ[\"CUDA_VISIBLE_DEVICES\"] = \"0\" # Possible GPUs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "997c6d64",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import sys; sys.path.insert(0, '..')\n",
    "import mair"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "35ec77bb",
   "metadata": {},
   "source": [
    "### Settings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5e747405",
   "metadata": {},
   "outputs": [],
   "source": [
    "PATH = \"./models/\"\n",
    "NAME = \"Sample\"\n",
    "SAVE_PATH = PATH + NAME\n",
    "\n",
    "MODEL_NAME = \"ResNet18\"\n",
    "DATA = \"CIFAR10\"\n",
    "MEAN = [0.4914, 0.4822, 0.4465]\n",
    "STD = [0.2023, 0.1994, 0.2010]\n",
    "\n",
    "N_VALIDATION = 1000\n",
    "N_CLASSES = 10\n",
    "\n",
    "EPOCH = 200\n",
    "\n",
    "EPS = 8/255\n",
    "ALPHA = 2/255\n",
    "STEPS = 10"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "39d49f77",
   "metadata": {},
   "source": [
    "### Load Data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "336c1032",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torchvision.datasets as dsets\n",
    "from torch.utils.data import DataLoader\n",
    "import torchvision.transforms as transforms\n",
    "transform = transforms.Compose([\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.4914, 0.4822, 0.4465), \n",
    "                         (0.2023, 0.1994, 0.2010))\n",
    "])\n",
    "\n",
    "MEAN = [0.4914, 0.4822, 0.4465]\n",
    "STD = [0.2023, 0.1994, 0.2010]\n",
    "\n",
    "train_data = dsets.CIFAR10(root='./data', \n",
    "                           train=True,\n",
    "                           download=True, \n",
    "                           transform=transform)\n",
    "\n",
    "test_data  = dsets.CIFAR10(root='./data', \n",
    "                           train=False,\n",
    "                           download=True, \n",
    "                           transform=transform)\n",
    "\n",
    "batch_size = 128\n",
    "train_loader = DataLoader(train_data, \n",
    "                          batch_size=batch_size,\n",
    "                          shuffle=True)\n",
    "\n",
    "test_loader = DataLoader(test_data, \n",
    "                         batch_size=batch_size,\n",
    "                         shuffle=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ea4593c3",
   "metadata": {},
   "source": [
    "### Load Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9806f202",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "model = mair.utils.load_model(model_name=MODEL_NAME, \n",
    "                              n_classes=N_CLASSES).cuda() # Load model\n",
    "rmodel = mair.RobModel(model, n_classes=N_CLASSES, \n",
    "                       normalization_used={'mean':MEAN, 'std':STD}).cuda()\n",
    "\n",
    "# For multi-GPUs\n",
    "# rmodel = rmodel.set_parallel() "
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f462a38c",
   "metadata": {},
   "source": [
    "### Start Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29a7cd48",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from mair.defenses import AT\n",
    "trainer = AT(rmodel, eps=EPS, alpha=ALPHA, steps=STEPS)\n",
    "trainer.record_rob(train_loader, test_loader, eps=EPS, alpha=2/255, steps=10, std=0.1, \n",
    "                   n_train_limit=N_VALIDATION, n_val_limit=N_VALIDATION)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8f36adff",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer.setup(optimizer=\"SGD(lr=0.1, momentum=0.9, weight_decay=0.0005)\",\n",
    "              scheduler=\"Step(milestones=[100, 150], gamma=0.1)\", \n",
    "              scheduler_type=\"Epoch\",\n",
    "              minimizer=None,\n",
    "              n_epochs=EPOCH, n_iters=len(train_loader)\n",
    "             )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1888254e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "trainer.fit(train_loader=train_loader, n_epochs=EPOCH,\n",
    "            save_path=SAVE_PATH, save_best={\"Clean(Val)\":\"HBO\", \"PGD(Val)\":\"HB\"},\n",
    "            save_type=None, save_overwrite=True, record_type=\"Epoch\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87d1c11a",
   "metadata": {},
   "source": [
    "### Evaluation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7197a5ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "rmodel.load_dict('./models/CIFAR10_ResNet18_AT(eps=8, alpha=2, steps=10)/best.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e8b9d903",
   "metadata": {},
   "outputs": [],
   "source": [
    "rmodel.eval_accuracy(test_loader)  # clean accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b9db063c",
   "metadata": {},
   "outputs": [],
   "source": [
    "rmodel.eval_rob_accuracy_gn(test_loader)  # gaussian noise accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9941645d",
   "metadata": {},
   "outputs": [],
   "source": [
    "rmodel.eval_rob_accuracy_fgsm(test_loader, eps)  # FGSM accuracy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "27fbc7b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "rmodel.eval_rob_accuracy_pgd(test_loader, eps, alpha, steps)  # PGD accuracy"
   ]
  }
 ],
 "metadata": {
  "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.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}