notebooks/find_model_mem_use_params.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Using GPUs\n"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import TensorDataset, DataLoader\n",
    "import numpy as np\n",
    "import time\n",
    "from typing import Type, Any, Callable, Union, List, Optional\n",
    "from torch import Tensor \n",
    "from torchsummary import summary\n",
    "import pickle\n",
    "if torch.cuda.is_available():\n",
    "    print(\"Using GPUs\")\n",
    "    device = torch.device(\"cuda\") \n",
    "else:\n",
    "    device = torch.device(\"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Files already downloaded and verified\n",
      "Files already downloaded and verified\n",
      "{'train': 50000, 'test': 5000, 'validation': 5000}\n"
     ]
    }
   ],
   "source": [
    "torch.manual_seed(43)\n",
    "batch_size = 128\n",
    "\n",
    "### for CIFAR 10\n",
    "# stats = ((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))\n",
    "## for CIFAR 100\n",
    "stats = ((0.507, 0.487, 0.441), (0.267, 0.256, 0.276))\n",
    "\n",
    "transform = torchvision.transforms.Compose([\n",
    "    torchvision.transforms.ToTensor(),\n",
    "    torchvision.transforms.Normalize(*stats),\n",
    "    torchvision.transforms.RandomCrop(32, padding=4, padding_mode='constant'),\n",
    "    torchvision.transforms.RandomHorizontalFlip(p=0.5)\n",
    "])\n",
    "\n",
    "train_set = torchvision.datasets.CIFAR100(root=\"data\", train=True, download=True, transform=transform)\n",
    "# train_set, _ = torch.utils.data.random_split(train_set, [1, len(train_set)-1]) # For sanity checking\n",
    "train_size = len(train_set)\n",
    "test_set = torchvision.datasets.CIFAR100(root=\"data\", train=False, download=True, transform=transform)\n",
    "test_set, validation_set = torch.utils.data.random_split(test_set, [5000, 5000])\n",
    "test_size = len(test_set)\n",
    "validation_size = len(validation_set)\n",
    "\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(train_set, batch_size, shuffle=True, num_workers=4, pin_memory=True)\n",
    "test_loader = torch.utils.data.DataLoader(test_set, batch_size, num_workers=4, pin_memory=True)\n",
    "validation_loader = torch.utils.data.DataLoader(validation_set, batch_size, num_workers=4, pin_memory=True)\n",
    "\n",
    "data_loaders = {\"train\": train_loader, \"test\": test_loader, \"validation\": validation_loader}\n",
    "dataset_sizes = {\"train\": train_size, \"test\": test_size, \"validation\": validation_size}\n",
    "print(dataset_sizes)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "from models import ResNet, DenseNet, DSNet\n",
    "\n",
    "#### Train Configurations, based on DSNet and ResNet paper\n",
    "model_n = 8\n",
    "epochs = 100 \n",
    "milestones = [int(epochs*0.5), int(epochs*0.75)]\n",
    "momentum = 0.9\n",
    "weight_decay = 0.0005\n",
    "gamma = 0.1\n",
    "lr = 0.1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "Path(\"./_results\").mkdir(exist_ok=True)\n",
    "\n",
    "f1 = open('./_results/results_1.txt', 'w')\n",
    "f2 = open('./_results/results_2.txt', 'w')\n",
    "\n",
    "# def print(*args):\n",
    "#     for arg in args:\n",
    "#         f1.write(str(arg) + ' ')\n",
    "#     f1.write('\\n')\n",
    "#     f1.flush()\n",
    "\n",
    "def print2(*args):\n",
    "    for arg in args:\n",
    "        f2.write(str(arg) + ' ')\n",
    "    f2.write('\\n')\n",
    "    f2.flush()\n",
    "    \n",
    "def mprint(b):\n",
    "    gb = b / (10 ** 9)\n",
    "    print('{} GB'.format(gb))\n",
    "    \n",
    "def print_n_params(model):\n",
    "    total_params = sum(p.numel() for p in model.parameters())\n",
    "    print(total_params)\n",
    "               "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_experiments = 5\n",
    "model_names = ['DSNet', 'ResNet', 'DenseNet']\n",
    "model_fns = [\n",
    "    lambda: DSNet(model_n, num_classes=100, device=device),\n",
    "    lambda: ResNet(model_n, num_classes=100, device=device),\n",
    "    lambda: DenseNet(growth_rate=16, block_config=(2 * model_n, 2 * model_n, 2 * model_n),\n",
    "                       num_init_features=16, bn_size=2, num_classes=100)\n",
    "]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "DSNet\n",
      "778212\n",
      "ResNet\n",
      "766116\n",
      "DenseNet\n",
      "771228\n"
     ]
    }
   ],
   "source": [
    "# Find params\n",
    "for n, fn in zip(model_names, model_fns):\n",
    "    model = fn()\n",
    "    print(n)\n",
    "    print_n_params(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "m1: 0\n",
      "m2:\n",
      "0.001166848 GB\n",
      "Configuration:  model: ResNet  model_n: 3  batch size: 128  optimizer:SGD  lr: 0.1  epochs: 100\n",
      "----------------------------- Train --------------------------------\n",
      "Epoch 1/100\n",
      "------------------------------\n",
      "before backward:\n",
      "0.196360192 GB\n",
      "after backward:\n",
      "0.005069312 GB\n",
      "before backward:\n",
      "0.19863552 GB\n",
      "after backward:\n",
      "0.005069312 GB\n",
      "before backward:\n",
      "0.19863552 GB\n",
      "after backward:\n",
      "0.005069312 GB\n",
      "before backward:\n",
      "0.19863552 GB\n",
      "after backward:\n",
      "0.005069312 GB\n",
      "before backward:\n",
      "0.19863552 GB\n",
      "after backward:\n",
      "0.005069312 GB\n",
      "before backward:\n",
      "0.19863552 GB\n",
      "after backward:\n",
      "0.005069312 GB\n",
      "before backward:\n",
      "0.19863552 GB\n",
      "after backward:\n",
      "0.005069312 GB\n",
      "before backward:\n",
      "0.19863552 GB\n",
      "after backward:\n",
      "0.005069312 GB\n",
      "before backward:\n",
      "0.19863552 GB\n",
      "after backward:\n",
      "0.005069312 GB\n",
      "before backward:\n",
      "0.19863552 GB\n",
      "after backward:\n",
      "0.005069312 GB\n",
      "before backward:\n",
      "0.19863552 GB\n",
      "after backward:\n",
      "0.005069312 GB\n",
      "before backward:\n",
      "0.19863552 GB\n",
      "after backward:\n",
      "0.005069312 GB\n",
      "before backward:\n",
      "0.19863552 GB\n",
      "after backward:\n",
      "0.005069312 GB\n",
      "before backward:\n",
      "0.19863552 GB\n",
      "after backward:\n",
      "0.005069312 GB\n",
      "before backward:\n",
      "0.19863552 GB\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "\u001b[0;32m<ipython-input-5-b203f867bb73>\u001b[0m in \u001b[0;36m<module>\u001b[0;34m\u001b[0m\n\u001b[1;32m     67\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     68\u001b[0m             \u001b[0;32mif\u001b[0m \u001b[0mphase\u001b[0m \u001b[0;34m==\u001b[0m \u001b[0;34m\"train\"\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m---> 69\u001b[0;31m                 \u001b[0mloss\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m  \u001b[0;31m# compute gradients\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m     70\u001b[0m                 \u001b[0moptimizer\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mstep\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0;34m)\u001b[0m \u001b[0;31m# update weights/biases\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m     71\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/conda/envs/dl/lib/python3.7/site-packages/torch/_tensor.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(self, gradient, retain_graph, create_graph, inputs)\u001b[0m\n\u001b[1;32m    361\u001b[0m                 \u001b[0mcreate_graph\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    362\u001b[0m                 inputs=inputs)\n\u001b[0;32m--> 363\u001b[0;31m         \u001b[0mtorch\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mautograd\u001b[0m\u001b[0;34m.\u001b[0m\u001b[0mbackward\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgradient\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0minputs\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m    364\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    365\u001b[0m     \u001b[0;32mdef\u001b[0m \u001b[0mregister_hook\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mself\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mhook\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m:\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n",
      "\u001b[0;32m/opt/conda/envs/dl/lib/python3.7/site-packages/torch/autograd/__init__.py\u001b[0m in \u001b[0;36mbackward\u001b[0;34m(tensors, grad_tensors, retain_graph, create_graph, grad_variables, inputs)\u001b[0m\n\u001b[1;32m    173\u001b[0m     Variable._execution_engine.run_backward(  # Calls into the C++ engine to run the backward pass\n\u001b[1;32m    174\u001b[0m         \u001b[0mtensors\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mgrad_tensors_\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mretain_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mcreate_graph\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0minputs\u001b[0m\u001b[0;34m,\u001b[0m\u001b[0;34m\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m--> 175\u001b[0;31m         allow_unreachable=True, accumulate_grad=True)  # Calls into the C++ engine to run the backward pass\n\u001b[0m\u001b[1;32m    176\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m    177\u001b[0m def grad(\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "\n",
    "results = {model_name: [{'train': [], 'test': None} for _ in range(n_experiments)] for model_name in model_names}\n",
    "\n",
    "i = 0\n",
    "model_name = model_names[i]\n",
    "model_fn = model_fns[i]\n",
    "\n",
    "m1 = torch.cuda.memory_allocated()\n",
    "print('m1:', m1)\n",
    "\n",
    "model = model_fn()\n",
    "model.to(device)\n",
    "\n",
    "m2 = torch.cuda.memory_allocated()\n",
    "print('m2:')\n",
    "mprint(m2)\n",
    "\n",
    "loss_fn = nn.CrossEntropyLoss()\n",
    "optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)\n",
    "scheduler = optim.lr_scheduler.MultiStepLR(optimizer, milestones=milestones, gamma=gamma)\n",
    "\n",
    "### Train loop + validation/ also test at the end\n",
    "print(\"Configuration: \", \"model:\", model_name, \" model_n:\", model_n, \" batch size:\", batch_size, \n",
    "      \" optimizer:SGD\", \" lr:\", lr, \" epochs:\", epochs)\n",
    "\n",
    "print(\"----------------------------- Train --------------------------------\")\n",
    "for epoch in range(epochs):\n",
    "    start_time = time.time()\n",
    "    print(\"Epoch {}/{}\".format(epoch+1, epochs))\n",
    "    print(\"-\" * 30)\n",
    "\n",
    "\n",
    "    epoch_loss = {\"train\": 0.0, \"validation\": 0.0}\n",
    "    epoch_acc = {\"train\": 0.0, \"validation\": 0.0}\n",
    "\n",
    "    running_loss = {\"train\": 0.0, \"validation\": 0.0}\n",
    "    running_corrects = {\"train\": 0, \"validation\": 0}\n",
    "\n",
    "    for phase in [\"train\", \"validation\"]:\n",
    "        if phase == \"train\":\n",
    "            model.train(True)\n",
    "        else:\n",
    "            model.train(False)\n",
    "\n",
    "        for data in data_loaders[phase]:\n",
    "            inputs, labels = data \n",
    "\n",
    "            inputs = inputs.to(device)\n",
    "            labels = labels.to(device)\n",
    "\n",
    "            optimizer.zero_grad() # clear all gradients\n",
    "\n",
    "            outputs = model(inputs) # batch_size x num_classes\n",
    "            _, preds = torch.max(outputs.data, 1) # values, indices\n",
    "            loss = loss_fn(outputs, labels)\n",
    "            \n",
    "            print('before backward:')\n",
    "            mprint(torch.cuda.memory_allocated())\n",
    "\n",
    "            if phase == \"train\":\n",
    "                loss.backward()  # compute gradients\n",
    "                optimizer.step() # update weights/biases\n",
    "            \n",
    "            print('after backward:')\n",
    "            mprint(torch.cuda.memory_allocated())\n",
    "            \n",
    "            running_loss[phase] += loss.data.item() * inputs.size(0)\n",
    "            running_corrects[phase] += torch.sum(preds == labels.data).item()\n",
    "\n",
    "        epoch_loss[phase] = running_loss[phase] / dataset_sizes[phase]\n",
    "        epoch_acc[phase] =  running_corrects[phase] / dataset_sizes[phase]\n",
    "\n",
    "    # Visualize the loss and accuracy values.\n",
    "    results_dic = {\n",
    "        'time': np.round(time.time()-start_time, 5),\n",
    "        'train_loss': np.round(epoch_loss[\"train\"], 5),\n",
    "        'train_acc': np.round(epoch_acc[\"train\"], 5),\n",
    "        'val_loss': np.round(epoch_loss[\"validation\"], 5),\n",
    "        'val_acc': np.round(epoch_acc[\"validation\"], 5),\n",
    "    }\n",
    "    print(results_dic)\n",
    "    results[model_name][n]['train'].append(results_dic)\n",
    "\n",
    "    scheduler.step()\n",
    "\n",
    "\n",
    "\n",
    "### evaluating the model with test set\n",
    "print(\"----------------------------- Test --------------------------------\")\n",
    "with torch.no_grad():\n",
    "    model.eval()\n",
    "    running_loss = 0\n",
    "    running_corrects = 0\n",
    "\n",
    "    for data in test_loader:\n",
    "        inputs, labels = data \n",
    "\n",
    "        inputs = inputs.to(device)\n",
    "        labels = labels.to(device)\n",
    "\n",
    "        optimizer.zero_grad() # clear all gradients\n",
    "\n",
    "        outputs = model(inputs) # batch_size x num_classes\n",
    "        _, preds = torch.max(outputs.data, 1) # values, indices\n",
    "        loss = loss_fn(outputs, labels)\n",
    "\n",
    "        running_loss += loss.data.item() * inputs.size(0)\n",
    "        running_corrects += torch.sum(preds == labels.data).item()\n",
    "\n",
    "# Visualize the loss and accuracy values.\n",
    "results_dic = {\n",
    "'time': np.round(time.time()-start_time, 5),\n",
    "'test_loss': np.round(running_loss/ dataset_sizes['test'], 5),\n",
    "'test_acc': np.round(running_corrects/ dataset_sizes['test'], 5),\n",
    "}\n",
    "print(results_dic)\n",
    "results[model_name][n]['test'] = results_dic\n",
    "\n",
    "print2('Experiment {}'.format(n))\n",
    "print2('test_acc', np.round(running_corrects/ dataset_sizes['test'], 5))\n",
    "\n",
    "with open('./_results/results_3.pk', 'wb') as f:\n",
    "    pickle.dump(results, f)\n",
    "    f.flush()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open('./_results/conf_bound_8/results_3.pk', 'rb') as f:\n",
    "    l_results = pickle.load(f)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "l_results['DSNet'][0]['train'] "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "l_results['DSNet'][0]['test']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}