notebooks/DenseNet.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# Various torch packages\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "# torchvision\n",
    "from torchvision import datasets, transforms\n",
    "\n",
    "# ------------------------\n",
    "# get up one directory \n",
    "import sys, os\n",
    "sys.path.append(os.path.abspath('../'))\n",
    "# ------------------------\n",
    "\n",
    "# custom packages\n",
    "import models.aux_funs as maf\n",
    "import optimizers as op\n",
    "import regularizers as reg\n",
    "import train\n",
    "import math\n",
    "import utils.configuration as cf\n",
    "import utils.datasets as ud\n",
    "from utils.datasets import get_data_set, GaussianSmoothing\n",
    "from models.fully_connected import fully_connected"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------------------------------------------------------------\n",
    "# Fix random seed\n",
    "# -----------------------------------------------------------------------------------\n",
    "random_seed = 2\n",
    "cf.seed_torch(random_seed)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "conf_args = {#\n",
    "    # data specification\n",
    "    'data_file':\"../../Data\", 'train_split':0.95, 'data_set':\"Fashion-MNIST\", 'download':False,\n",
    "    # cuda\n",
    "    'use_cuda':True, 'num_workers':4, 'cuda_device':0, 'pin_memory':True,\n",
    "    #\n",
    "    'epochs':100,\n",
    "    # optimizer\n",
    "    'delta':1.0, 'lr':0.001, 'lamda_0':0.05, 'lamda_1':0.05, 'optim':\"AdaBreg\", 'row_group':True,\n",
    "    'reg':reg.reg_l1_l2, 'beta':0.0,\n",
    "    # model\n",
    "    'model_size':7*[28*28], 'act_fun':torch.nn.ReLU(),\n",
    "    # initialization\n",
    "    'sparse_init':0.03, 'r':[1,5,1],\n",
    "    # misc\n",
    "    'random_seed':random_seed, 'eval_acc':True, 'name':'---', 'super_type':'---'\n",
    "}\n",
    "conf = cf.Conf(**conf_args)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Define DenseNet model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class BasicBlock(nn.Module):\n",
    "    def __init__(self, in_planes, out_planes):\n",
    "        super(BasicBlock, self).__init__()\n",
    "        self.relu = nn.ReLU(inplace=True)\n",
    "        \n",
    "        \n",
    "        self.out_planes = out_planes\n",
    "        self.in_planes = in_planes\n",
    "        \n",
    "        self.bn1 = nn.BatchNorm2d(out_planes)\n",
    "        self.conv1 = nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=1,padding=1)\n",
    "        s = torch.zeros((self.in_planes,))              \n",
    "        \n",
    "    def forward(self, x):\n",
    "        out = self.relu(self.conv1(x))\n",
    "\n",
    "        return torch.cat([x, out], 1)\n",
    "    \n",
    "class DenseBlock(nn.Module):\n",
    "    def __init__(self, num_layers, in_planes, out_planes):\n",
    "        super(DenseBlock, self).__init__()\n",
    "        \n",
    "        layers = []\n",
    "        for i in range(num_layers):\n",
    "            layers.append(BasicBlock(in_planes + i*out_planes, out_planes))\n",
    "        \n",
    "        self.layer = nn.Sequential(*layers)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.layer(x)\n",
    "\n",
    "class LinearBlock(nn.Module):\n",
    "    def __init__(self, in_size, out_size):\n",
    "        super(LinearBlock, self).__init__()\n",
    "        self.linear = nn.Linear(in_size, out_size)\n",
    "        self.act = nn.ReLU()\n",
    "    \n",
    "    def forward(self, x):\n",
    "        x = nn.Flatten()(x[:,-1,:])\n",
    "        return self.act(self.linear(x))\n",
    "\n",
    "class DenseNet(nn.Module):\n",
    "    def __init__(self, depth, planes, num_classes, mean = 0.0, std = 1.0, im_channels = 3, im_size=32):\n",
    "        super(DenseNet, self).__init__()\n",
    "        self.mean = mean\n",
    "        self.std = std\n",
    "        \n",
    "        self.depth = depth\n",
    "        self.planes = planes\n",
    "        self.num_classes = num_classes\n",
    "        \n",
    "        self.conv1 = nn.Conv2d(im_channels, planes, kernel_size=3, stride=1, padding=1, bias=False)\n",
    "        self.block1 = DenseBlock(depth, planes, planes)\n",
    "        \n",
    "        self.trans1 = BasicBlock((depth + 1) * planes, 1)\n",
    "        self.fc = LinearBlock(im_size*im_size, num_classes)\n",
    "    \n",
    "    def forward(self, x):\n",
    "        out = (x-self.mean)/self.std\n",
    "        out = self.conv1(out)\n",
    "        out = self.block1(out)\n",
    "        out = self.trans1(out)\n",
    "        return self.fc(out)     "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# -----------------------------------------------------------------------------------\n",
    "# define the model and an instance of the best model class\n",
    "# -----------------------------------------------------------------------------------\n",
    "model_kwargs = {'mean':conf.data_set_mean, 'std':conf.data_set_std, 'im_size':conf.im_shape[1],\n",
    "                'im_channels':conf.im_shape[0]}    \n",
    "\n",
    "model = DenseNet(5, 12, 10, **model_kwargs)\n",
    "    \n",
    "# sparsify\n",
    "maf.sparse_bias_uniform_(model, 0,conf.r[0])\n",
    "maf.sparse_weight_normal_(model, conf.r[1])\n",
    "maf.sparsify_(model, conf.sparse_init, row_group = conf.row_group)\n",
    "model = model.to(conf.device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------------------------------------------------------------\n",
    "# define the model and an instance of the best model class\n",
    "# -----------------------------------------------------------------------------------\n",
    "model_kwargs = {'mean':conf.data_set_mean, 'std':conf.data_set_std}    \n",
    "\n",
    "def init_weights(conf, model):\n",
    "    # sparsify\n",
    "    maf.sparse_bias_uniform_(model, 0, conf.r[0])\n",
    "    maf.sparse_weight_normal_(model, conf.r[1])\n",
    "    maf.sparse_weight_normal_(model, conf.r[1],ltype=nn.Conv2d)\n",
    "    maf.sparsify_(model, conf.sparse_init, ltype = nn.Conv2d, row_group = conf.row_group)\n",
    "    model = model.to(conf.device)\n",
    "    \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "# -----------------------------------------------------------------------------------\n",
    "# Optimizer\n",
    "# -----------------------------------------------------------------------------------\n",
    "def get_skips(model):\n",
    "    for m in model.modules():\n",
    "        if hasattr(m,'skips'):\n",
    "            yield m.skips\n",
    "        else:\n",
    "            continue\n",
    "\n",
    "def print_skips(model):\n",
    "    for m in model.modules():\n",
    "        if hasattr(m,'skips'):\n",
    "            print((0.001*torch.round(1000*m.skips.data).cpu()))\n",
    "            \n",
    "def skips_to_list(model):\n",
    "    skips = []\n",
    "    for m in model.modules():\n",
    "        if hasattr(m,'skips'):\n",
    "            skips.append(m.skips.data.tolist())\n",
    "    return skips\n",
    "    \n",
    "\n",
    "def init_opt(conf, model):\n",
    "    # Get access to different model parameters\n",
    "    weights_linear = maf.get_weights_linear(model)\n",
    "    weights_conv = maf.get_weights_conv(model)\n",
    "    biases = maf.get_bias(model)\n",
    "    skips = get_skips(model)\n",
    "\n",
    "    # -----------------------------------------------------------------------------------\n",
    "    # Initialize optimizer\n",
    "    # -----------------------------------------------------------------------------------\n",
    "    reg1 = conf.reg(lamda=conf.lamda_0)\n",
    "    reg2 = reg.reg_l1(lamda=conf.lamda_1)\n",
    "    reg3 = reg.reg_l1_l2_conv(lamda=conf.lamda_0)\n",
    "\n",
    "    if conf.optim == \"SGD\":\n",
    "        opt = torch.optim.SGD(model.parameters(), lr=conf.lr, momentum=conf.beta)\n",
    "    elif conf.optim == \"AdaBreg\":\n",
    "        opt = op.AdaBreg([{'params': weights_linear, 'lr' : conf.lr},\n",
    "                           {'params': weights_conv, 'lr' : conf.lr, 'reg' : reg3},\n",
    "                           {'params': biases, 'lr': conf.lr},\n",
    "                           {'params': skips, 'lr':conf.lr, 'reg':reg2}])\n",
    "\n",
    "    # learning rate scheduler\n",
    "    scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(opt, factor=0.5, patience=5,threshold=0.01)\n",
    "    \n",
    "    return opt, scheduler"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "save_params = False\n",
    "if save_params:\n",
    "    conf.write_to_csv()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader, valid_loader, test_loader = ud.get_data_set(conf)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# History and Run Specification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# -----------------------------------------------------------------------------------\n",
    "# initalize history\n",
    "# -----------------------------------------------------------------------------------\n",
    "tracked = ['loss', 'node_sparse']\n",
    "train_hist = {}\n",
    "val_hist = {}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Training"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# -----------------------------------------------------------------------------------\n",
    "# Reinit weigts and the corresponding optimizer\n",
    "# -----------------------------------------------------------------------------------\n",
    "model = init_weights(conf, model)\n",
    "opt, scheduler = init_opt(conf, model)\n",
    "\n",
    "# -----------------------------------------------------------------------------------\n",
    "# train the model\n",
    "# -----------------------------------------------------------------------------------\n",
    "for epoch in range(conf.epochs):\n",
    "    print(25*\"<>\")\n",
    "    print(50*\"|\")\n",
    "    print(25*\"<>\")\n",
    "    print('Epoch:', epoch)\n",
    "\n",
    "    # ------------------------------------------------------------------------\n",
    "    # train step, log the accuracy and loss\n",
    "    # ------------------------------------------------------------------------\n",
    "    train_data = train.train_step(conf, model, opt, train_loader)\n",
    "\n",
    "    # update history\n",
    "    for key in tracked:\n",
    "        if key in train_data:\n",
    "            var_list = train_hist.setdefault(key, [])\n",
    "            var_list.append(train_data[key])        \n",
    "\n",
    "    # ------------------------------------------------------------------------\n",
    "    # validation step\n",
    "    val_data = train.validation_step(conf, model, opt, valid_loader)\n",
    "\n",
    "    print_skips(model)\n",
    "\n",
    "    # update history\n",
    "    for key in tracked:\n",
    "        if key in val_data:\n",
    "            var = val_data[key]\n",
    "            if isinstance(var, list):\n",
    "                for i, var_loc in enumerate(var):\n",
    "                    key_loc = key+\"_\" + str(i)\n",
    "                    var_list = val_hist.setdefault(key_loc, [])\n",
    "                    val_hist[key_loc].append(var_loc)\n",
    "            else:\n",
    "                var_list = val_hist.setdefault(key, [])\n",
    "                var_list.append(var) \n",
    "\n",
    "\n",
    "    scheduler.step(train_data['loss'])\n",
    "    print(\"Learning rate:\",opt.param_groups[0]['lr'])\n",
    "    #best_model(train_data['acc'], val_data['acc'], model=model)"
   ]
  }
 ],
 "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}