{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a51921f8-2248-4632-ab5a-5a57f3c7a766",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Example of using the advertrain library"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "636f0ed2-8e6e-4a96-b427-4a2dd63a6f3d",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jovyan/Maturation/env-robustml/lib/python3.10/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "from torch import nn, optim\n",
    "from pathlib import Path\n",
    "import data_all\n",
    "from robustML.advertrain.models import ConvNet, ResNet, ConvNetDropblock, ResNetDropblock\n",
    "from robustML.advertrain.training.classical_training import ClassicalTraining\n",
    "from robustML.advertrain.training.adversarial_training import AdversarialTraining\n",
    "from robustML.advertrain.training.autoattack_training import AutoAttackTraining\n",
    "from robustML.advertrain.training.fire_training import FIRETraining\n",
    "from robustML.advertrain.training.trades_training import TRADESTraining\n",
    "from robustML.advertrain.transforms import DataTransformations\n",
    "import sys\n",
    "import os\n",
    "import argparse"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c372c53a-6afe-441b-ac97-68d47209aa43",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set here the path to the directory containing the folder c00 of the example dataset. it can be downloaded\n",
    "# here http://minio-storage.apps.confianceai-public.irtsysx.fr/ml-models/robust-ml-dataset.zip\n",
    "\n",
    "ROOT_DIR = \"/home/jovyan/Maturation/as434_robustness_platform/examples/dataset/c00\" \n",
    "\n",
    "DEVICE = \"cuda:1\"\n",
    "ARCHITECTURE = \"resnetdropblock\"\n",
    "LEARNING_RATE = 1e-4\n",
    "MODELS_PATH = \"results/\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "2b0f8833-4876-4859-8bf5-cce9362e3272",
   "metadata": {},
   "outputs": [],
   "source": [
    "transformer = DataTransformations()\n",
    "train_transform = transformer.get_train_transforms()\n",
    "test_transform = transformer.get_test_transforms()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "ea9994dc-ce26-4383-9833-a1a12daab109",
   "metadata": {},
   "outputs": [],
   "source": [
    "path_data = \"learning_data/\"\n",
    "\n",
    "train_dataloader1, val_dataloader1 = data_all.get_train_dataloader(\n",
    "    ROOT_DIR,\n",
    "    path_data+\"train_dataset_all.csv\",\n",
    "    path_data+\"val_dataset_all.csv\",\n",
    "    batch_size=16,\n",
    "    train_transform=train_transform,\n",
    "    test_transform=test_transform,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "55e7b801-1ecb-4ef9-9f61-7c3472ce6e54",
   "metadata": {},
   "outputs": [],
   "source": [
    "if ARCHITECTURE == \"convnet\":\n",
    "    model = ConvNet(DEVICE)\n",
    "elif ARCHITECTURE == \"resnet\":\n",
    "    model = ResNet(DEVICE)\n",
    "elif ARCHITECTURE == \"convnetdropblock\":\n",
    "    model = ConvNetDropblock(DEVICE)\n",
    "elif ARCHITECTURE == \"resnetdropblock\":\n",
    "    model = ResNetDropblock(DEVICE)\n",
    "else:\n",
    "    raise ValueError(\"Unknown architecture\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3c56e5f0-a31f-48e6-a3a2-fa8a0d198292",
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(model.parameters(), lr=LEARNING_RATE)\n",
    "criterion = nn.CrossEntropyLoss(weight=torch.Tensor([1, 1]).to(DEVICE))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "89a983a6-7377-4e5f-a98f-4510aa632a1e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0\n",
      "GPU device used  cuda:1\n"
     ]
    }
   ],
   "source": [
    "if not torch.cuda.is_available() or DEVICE ==\"cpu\":\n",
    "        print(\"ERROR : Learning should be done on a system with GPU\")\n",
    "        sys.exit()\n",
    "        \n",
    "print(torch.cuda.current_device())\n",
    "print(\"GPU device used \", DEVICE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f00d709e-05ca-414f-a127-77e0700582ef",
   "metadata": {},
   "source": [
    "## Train class"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5b6c5d00-4d0d-4e78-8ee5-b8ec44069799",
   "metadata": {},
   "source": [
    "### ClassicalTraining"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "85a473db-105f-4d56-8d59-0213dabd0d19",
   "metadata": {},
   "outputs": [],
   "source": [
    "METHODE=\"vanilla\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "5b001bca-044a-4a04-904b-aa1a774a3b37",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = ClassicalTraining(model, optimizer, criterion, DEVICE)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "080de670-3045-4c71-925a-3a1eabe548fd",
   "metadata": {},
   "source": [
    "### AdversarialTraining"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "afabc6dd-f091-422e-aead-618eb718c932",
   "metadata": {},
   "outputs": [],
   "source": [
    "METHODE = \"adversarial_training\"\n",
    "EPSILON = 8 / 255"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "62a4139b-fc6b-4bc0-8ae6-9857435dc391",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = AdversarialTraining(model, optimizer, criterion, DEVICE, EPSILON)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56f3ccdb-8d0e-40d7-a602-e4655b71ff14",
   "metadata": {},
   "source": [
    "### AutoAttackTraining"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "e4f27698-5395-4ffa-8ad7-79e4a9788e7f",
   "metadata": {},
   "outputs": [],
   "source": [
    "METHODE = \"autoattack_training\"\n",
    "EPSILON = 8 / 255"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "7064d42d-be56-4ab1-8da3-66c748d03477",
   "metadata": {},
   "outputs": [],
   "source": [
    "trainer = AutoAttackTraining(model, optimizer, criterion, DEVICE, \"ce\", EPSILON)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "72778a7e-ccb1-4cfc-be4d-fc8d9b769d62",
   "metadata": {},
   "source": [
    "### FireTraining"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "8596e48b-bf47-4f45-acfc-708887fde676",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "METHODE = \"fire_training\"\n",
    "EPSILON = 8 / 255\n",
    "BETA = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "44575d87-8ce8-45b6-833c-972b40c11c3d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "trainer = FIRETraining(model, optimizer, DEVICE, EPSILON, BETA)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20e3e283-2aa3-410d-adc3-678dc1f3b186",
   "metadata": {},
   "source": [
    "### TradesTraining"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "de1ca19b-56d0-4fcb-b697-f0a1fe4b5657",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "METHODE = \"trades_training\"\n",
    "EPSILON = 8 / 255\n",
    "BETA = 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "53db02fa-3c6d-421e-aeea-70718611bf08",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "trainer = TRADESTraining(model, optimizer, DEVICE, EPSILON, BETA)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48752917-3cfc-400d-9008-dd2b5827e4bc",
   "metadata": {},
   "source": [
    "## Fit the model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "ed8fe6a1-3890-4926-b76f-9f302b8ed48f",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using 'results/resnetdropblock_trades_training' directory\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jovyan/Maturation/env-robustml/lib/python3.10/site-packages/torch/nn/_reduction.py:42: UserWarning: size_average and reduce args will be deprecated, please use reduction='sum' instead.\n",
      "  warnings.warn(warning.format(ret))\n",
      "Epochs 1/2 : Training: 100%|██████████| 750/750 [02:43<00:00,  4.59it/s]\n",
      "Validation: 100%|██████████| 33/33 [00:08<00:00,  3.93it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 1/2\n",
      "Train Loss:  0.690, Acc:  0.673, Recall:  0.617, Precision:  0.685,F1 Score:  0.649\n",
      "Validation Loss:  0.566, Acc:  0.773, Recall:  0.776, Precision:  0.739,F1 Score:  0.757\n",
      "Validation loss decreased (inf --> 0.566105). Saving model ...\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Epochs 2/2 : Training: 100%|██████████| 750/750 [02:45<00:00,  4.54it/s]\n",
      "Validation: 100%|██████████| 33/33 [00:08<00:00,  4.12it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 2/2\n",
      "Train Loss:  0.538, Acc:  0.797, Recall:  0.795, Precision:  0.799,F1 Score:  0.797\n",
      "Validation Loss:  0.514, Acc:  0.837, Recall:  0.867, Precision:  0.795,F1 Score:  0.829\n",
      "Validation loss decreased (0.566105 --> 0.514262). Saving model ...\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "{'loss': [0.6897706389427185, 0.5378051996231079],\n",
       " 'acc': [0.6733333468437195, 0.79666668176651],\n",
       " 'val_loss': [0.5661050081253052, 0.5142622590065002],\n",
       " 'val_acc': [0.7727272510528564, 0.8371211886405945],\n",
       " 'optimizer': 'Adam (\\nParameter Group 0\\n    amsgrad: False\\n    betas: (0.9, 0.999)\\n    capturable: False\\n    differentiable: False\\n    eps: 1e-08\\n    foreach: None\\n    fused: False\\n    lr: 0.0001\\n    maximize: False\\n    weight_decay: 0\\n)',\n",
       " 'loss_func': 'None',\n",
       " 'epochs': 2,\n",
       " 'patience': 2}"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "path = os.path.join(MODELS_PATH,\n",
    "                    f\"{ARCHITECTURE}\"+\"_\"+f\"{METHODE}\")\n",
    "\n",
    "if not os.path.exists(path):\n",
    "    try:\n",
    "        os.makedirs(path)\n",
    "    except FileExistsError:\n",
    "        pass\n",
    "\n",
    "print(f\"Using '{path}' directory\")\n",
    "\n",
    "trainer.fit(epochs=2, train_dataloader=train_dataloader1, val_dataloader=val_dataloader1, patience=2, checkpoint=path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9341f59d-c40c-430f-a855-9c5e65f78afc",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "env-robust-ml",
   "language": "python",
   "name": "env-robust-ml"
  },
  "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.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}