train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import cv2\n",
    "from pylab import *\n",
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "import torchvision.transforms as transforms\n",
    "# from torchsummary import summary\n",
    "from dataset import CelebA\n",
    "\n",
    "from torch.autograd import Variable\n",
    "\n",
    "import torch.optim as optim\n",
    "import sys\n",
    "\n",
    "from albumentations import (\n",
    "    HorizontalFlip,\n",
    "    VerticalFlip,\n",
    "    Normalize,\n",
    "    Compose,\n",
    "    PadIfNeeded,\n",
    "    RandomCrop,\n",
    "    Rotate,\n",
    "    Resize\n",
    ")\n",
    "from models import LinkNet34\n",
    "# from fastai.vision import *\n",
    "# from fastai.callbacks.hooks import *\n",
    "# from fastai.utils.mem import *\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "class LossBinary:\n",
    "\n",
    "    def __init__(self, jaccard_weight=0):\n",
    "        self.nll_loss = nn.BCEWithLogitsLoss()\n",
    "        self.jaccard_weight = jaccard_weight\n",
    "\n",
    "    def __call__(self, outputs, targets):\n",
    "        loss = (1 - self.jaccard_weight) * self.nll_loss(outputs, targets)\n",
    "\n",
    "        if self.jaccard_weight:\n",
    "            eps = 1e-15\n",
    "            jaccard_target = (targets == 1).float()\n",
    "            jaccard_output = F.sigmoid(outputs)\n",
    "\n",
    "            intersection = (jaccard_output * jaccard_target).sum()\n",
    "            union = jaccard_output.sum() + jaccard_target.sum()\n",
    "\n",
    "            loss -= self.jaccard_weight * torch.log((intersection + eps) / (union - intersection + eps))\n",
    "        return loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "root = '/media/nasir/Drive1/datasets/celeba_real/'\n",
    "size = 256\n",
    "train_aug= Compose([\n",
    "            HorizontalFlip(p=0.5),\n",
    "            Rotate(15),\n",
    "         ])\n",
    "trainset = CelebA(root = root, train=True, augmentation=train_aug)\n",
    "\n",
    "trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,\n",
    "                                          shuffle=True, num_workers=4)\n",
    "\n",
    "\n",
    "valset = CelebA(root = root, train=False)\n",
    "\n",
    "valset_loader = torch.utils.data.DataLoader(valset, batch_size=4,\n",
    "                                          shuffle=True, num_workers=4)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def freeze(model, l):\n",
    "    total_params = 0\n",
    "    for i, param in enumerate(model.parameters()):\n",
    "        total_params+=1\n",
    "    for i, param in enumerate(model.parameters()):\n",
    "        if i < total_params - l:\n",
    "            param.requires_grad = False\n",
    "        else:\n",
    "            param.requires_grad = True\n",
    "            \n",
    "    return model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "valid_loss_es = []\n",
    "def validate(net, criterion):\n",
    "    net.eval()\n",
    "    valid_loss = 0.0\n",
    "    for i, (imgs, true_masks) in enumerate(valset_loader):\n",
    "        imgs, true_masks = Variable(imgs.cuda()), Variable(true_masks.to(dtype=torch.float, device = device))\n",
    "\n",
    "\n",
    "        masks_pred = net(imgs)\n",
    "        masks_probs_flat = masks_pred.view(-1)\n",
    "        true_masks_flat = true_masks.view(-1)\n",
    "\n",
    "        loss = criterion(masks_probs_flat, true_masks_flat)\n",
    "        valid_loss += loss.item()\n",
    "        \n",
    "    net.train()\n",
    "    valid_loss = valid_loss/i\n",
    "    \n",
    "    valid_loss_es.append(valid_loss)\n",
    "    best_valid_loss = min(valid_loss_es)\n",
    "\n",
    "    print(f'validation loss {round(valid_loss, 4)}, {round(best_valid_loss, 4)}')\n",
    "    if valid_loss <= best_valid_loss:\n",
    "        best_valid_loss = valid_loss\n",
    "        print('saving ... ')\n",
    "        torch.save(net.state_dict(), 'linknet.pth')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, trainloader, epochs, lr=1e-2):\n",
    "    optimizer = optim.SGD(model.parameters(),\n",
    "        lr=lr,\n",
    "        momentum=0.9,\n",
    "        weight_decay=0.0005\n",
    "    )\n",
    "    cosine = optim.lr_scheduler.CosineAnnealingLR(optimizer, 2)\n",
    "    criterion = LossBinary(jaccard_weight=1)\n",
    "#     criterion = nn.BCELoss()\n",
    "    for epoch in range(epochs):\n",
    "        model.train()\n",
    "        epoch_loss = 0\n",
    "        for i, (imgs, true_masks) in enumerate(trainloader):\n",
    "\n",
    "            imgs, true_masks = Variable(imgs.cuda()), Variable(true_masks.to(dtype=torch.float, device = device))\n",
    "\n",
    "\n",
    "            optimizer.zero_grad()\n",
    "            masks_pred = model(imgs)\n",
    "            masks_probs_flat = masks_pred.view(-1)\n",
    "            true_masks_flat = true_masks.view(-1)\n",
    "            loss = criterion(masks_probs_flat, true_masks_flat)\n",
    "            epoch_loss += loss.item()\n",
    "            loss.backward()\n",
    "            optimizer.step()\n",
    "#             if i % 10 == 9:\n",
    "            sys.stdout.write(f'\\rEpoch: {epoch} ---- Loss: {round(epoch_loss/(i+1), 4)}')\n",
    "            sys.stdout.flush()\n",
    "            cosine.step()\n",
    "            \n",
    "        sys.stdout.write('\\n')\n",
    "        validate(model, criterion)\n",
    "        sys.stdout.write('\\n')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = LinkNet34()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "model.train()\n",
    "model.to(device)\n",
    "model.load_state_dict(torch.load('linknet.pth'))\n",
    "\n",
    "summary(model, (3, 256, 256))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "train(model, trainloader, 5, lr=1e-2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train(model, trainloader, 5, lr=1e-3)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "train(model, trainloader, 10, lr=1e-4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "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.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}