FaceTest.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.autograd import Variable\n",
    "from torch.optim.lr_scheduler import MultiStepLR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class G(nn.Module):\n",
    "    def __init__(self, channels):\n",
    "        super(self.__class__, self).__init()\n",
    "        self.channels = channels\n",
    "        self.g = nn.Sequential(\n",
    "                    nn.ConvTranspose2d(self.channels, 512, kernel_size=4, stride=2, padding=1),\n",
    "                    nn.BatchNorm2d(512),\n",
    "                    nn.LeakyReLU(0.2, inplace=True),\n",
    "                    nn.Conv2d(512, 512, kernel_size=1, stride=1, padding=0),\n",
    "                    nn.BatchNorm2d(512),\n",
    "                    nn.LeakyReLU(0.2, inplace=True),\n",
    "        \n",
    "                    nn.ConvTranspose2d(512, 256, kernel_size=4, stride=2, padding=1),\n",
    "                    nn.BatchNorm2d(256),\n",
    "                    nn.LeakyReLU(0.2, inplace=True),\n",
    "                    nn.Conv2d(256, 256, kernel_size=1, stride=1, padding=0),\n",
    "                    nn.BatchNorm2d(256),\n",
    "                    nn.LeakyReLU(0.2, inplace=True),\n",
    "        \n",
    "                    nn.ConvTranspose2d(256, 128, kernel_size=4, stride=2, padding=1),\n",
    "                    nn.BatchNorm2d(128),\n",
    "                    nn.LeakyReLU(0.2, inplace=True),\n",
    "                    nn.Conv2d(128, 128, kernel_size=1, stride=1, padding=0),\n",
    "                    nn.BatchNorm2d(128),\n",
    "                    nn.LeakyReLU(0.2, inplace=True),\n",
    "        \n",
    "                    nn.ConvTranspose2d(128, 64, kernel_size=4, stride=2, padding=1),\n",
    "                    nn.BatchNorm2d(64),\n",
    "                    nn.LeakyReLU(0.2, inplace=True),\n",
    "                    nn.Conv2d(64, 64, kernel_size=1, stride=1, padding=0),\n",
    "                    nn.BatchNorm2d(64),\n",
    "                    nn.LeakyReLU(0.2, inplace=True),\n",
    "        \n",
    "                    nn.ConvTranspose2d(64, 32, kernel_size=4, stride=2, padding=1),\n",
    "                    nn.BatchNorm2d(32),\n",
    "                    nn.LeakyReLU(0.2, inplace=True),\n",
    "                    nn.Conv2d(32, 32, kernel_size=1, stride=1, padding=0),\n",
    "                    nn.BatchNorm2d(32),\n",
    "                    nn.LeakyReLU(0.2, inplace=True),\n",
    "        \n",
    "                    nn.ConvTranspose2d(32, 3, kernel_size=4, stride=2, padding=1),\n",
    "                    nn.Tanh())\n",
    "        \n",
    "    def forward(self, inp):\n",
    "        resh_inp = inp.view(inp.size()[0], 512, 1, 1)\n",
    "        out = self.g(resh_inp)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class D(nn.Module):\n",
    "    def __init__(self, channels, al=0.2):\n",
    "        self.channels = channels\n",
    "        self.al = al\n",
    "        self.d = nn.Sequential(\n",
    "                    nn.Conv2d(3, self.channels, kernel_size=3, stride=2, padding=1),\n",
    "                    nn.BatchNorm2d(self.channels),\n",
    "                    nn.LeakyReLU(self.alpha, inplace=True),\n",
    "        \n",
    "                    nn.Conv2d(self.channels, self.channels*2, kernel_size=3, stride=2, padding=1),\n",
    "                    nn.BatchNorm2d(self.channels),\n",
    "                    nn.LeakyReLU(self.alpha, inplace=True),\n",
    "        \n",
    "                    nn.Conv2d(self.channels*2, self.channels*4, kernel_size=3, stride=2, padding=1),\n",
    "                    nn.BatchNorm2d(self.channels),\n",
    "                    nn.LeakyReLU(self.alpha, inplace=True),\n",
    "        \n",
    "                    nn.Conv2d(self.channels*4, self.channels*2, kernel_size=3, stride=2, padding=1),\n",
    "                    nn.BatchNorm2d(self.channels),\n",
    "                    nn.LeakyReLU(self.alpha, inplace=True),\n",
    "        \n",
    "                    nn.Conv2d(self.channels*2, self.channels, kernel_size=3, stride=2, padding=1),\n",
    "                    nn.BatchNorm2d(self.channels),\n",
    "                    nn.LeakyReLU(self.alpha, inplace=True),\n",
    "                    \n",
    "                    nn.Conv2d(self.channels, 3, kernel_size=3, stride=1))\n",
    "        \n",
    "    def forward(self, inp):\n",
    "        out = self.d(inp)\n",
    "        return out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class ZeroPadBottom(object):\n",
    "    ''' Zero pads batch of image tensor Variables on bottom to given size. Input (B, C, H, W) - padded on H axis. '''\n",
    "    def __init__(self, size, use_gpu=True):\n",
    "        self.size = size\n",
    "        self.use_gpu = use_gpu\n",
    "        \n",
    "    def __call__(self, sample):\n",
    "        B, C, H, W = sample.size()\n",
    "        diff = self.size - H\n",
    "        padding = Variable(torch.zeros(B, C, diff, W), requires_grad=False)\n",
    "        if self.use_gpu:\n",
    "            padding = padding.cuda()\n",
    "        zero_padded = torch.cat((sample, padding), dim=2)\n",
    "        return zero_padded"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "class face_model(object):\n",
    "    def __init__(self, use_gpu=True):\n",
    "        super(face_model, self).__init__(use_gpu)\n",
    "        self.generator_loss_func = None\n",
    "        self.discriminator_loss_func = None\n",
    "        self.gan_loss_func = None\n",
    "        self.generator_smooth_func = None\n",
    "        self.source_val_loader = None\n",
    "        self.source_test_loader = None\n",
    "        self.target_test_loader = None\n",
    "        self.source_train_loader = None\n",
    "        self.target_train_loader = None\n",
    "        self.batch_size = 128\n",
    "        self.lossCE = nn.CrossEntropyLoss()\n",
    "        \n",
    "    def make_loader(self):\n",
    "        '''TO DO'''\n",
    "    \n",
    "    def make_model(self):\n",
    "        \n",
    "        self.model = {}\n",
    "        f = sphere20a(feature = True)\n",
    "        f.load_state_dict(torch.load('./sphere20a_20171020.pth'))\n",
    "        for params in f.parameters():\n",
    "            params.require_grad = False\n",
    "        \n",
    "        self.model['f'] = f\n",
    "        self.model['g'] = D(128, alpha=0.2)\n",
    "        self.model['d'] = G(channels=512)\n",
    "        \n",
    "        if self.use_gpu:\n",
    "            self.model['g'] = self.model['g'].cuda()\n",
    "            self.model['d'] = self.model['d'].cuda()\n",
    "            \n",
    "        self.up = nn.Upsample(size=(96,96), mode='bilinear')\n",
    "        self.pad = ZeroPadBottom(112)\n",
    "        \n",
    "    def make_loss_func(self):\n",
    "        \n",
    "        self.lossCE = nn.CrossEntropyLoss().cuda()\n",
    "        self.lossMSE = nn.MSELoss().cuda()\n",
    "        lab0, lab1, lab2 = (torch.LongTensor(self.batch_size) for i in range(3))\n",
    "        lab0 = Variable(lab0.cuda())\n",
    "        lab1 = Variable(lab1.cuda())\n",
    "        lab2 = Variable(lab2.cuda())\n",
    "        \n",
    "        lab0.data.resize_(self.batch_size).fill_(0)\n",
    "        lab1.data.resize_(self.batch_size).fill_(1)\n",
    "        lab2.data.resize_(self.batch_size).fill_(2)\n",
    "        \n",
    "        self.lab0 = lab0\n",
    "        self.lab1 = lab1\n",
    "        self.lab2 = lab2\n",
    "        \n",
    "        self.make_generator_loss_func()\n",
    "        self.make_discriminator_loss_func()\n",
    "        self.make_smooth_func()\n",
    "        self.make_dist_func_targ_domain()\n",
    "    \n",
    "    def make_opt(self):\n",
    "        \n",
    "        self.generator_opt = optim.Adam(self.model['g'].parameters(), lr = 2e-4, betas=(0.5, 0.999), weight_decay=1e-6)\n",
    "        self.discriminator_opt = optim.Adam(self.model['d'].parameters(), lr = 2e-4, betas=(0.5, 0.999), weight_decay=1e-6)\n",
    "        \n",
    "        self.generator_lr_sche = MultiStepLR(self.generator_opt, milestones=[15000], gamma=0.1)\n",
    "        self.discriminator_lr_sche = MultiStepLR(self.discriminator_opt, milestones=[15000], gamma=0.1)\n",
    "    \n",
    "    def cos_sim(self, x, y):\n",
    "        ab = torch.sum(x*y, dim=1)\n",
    "        a = torch.sqrt(torch.sum(x*x, dim=1))\n",
    "        b = torch.sqrt(torch.sum(y*y, dim=1))\n",
    "        sim = ab/(a*b)\n",
    "        avg_sim = torch.mean(sim)\n",
    "        cos_loss = 1-avg_sim\n",
    "        return cos_loss\n",
    "        \n",
    "    def make_generator_loss_func(self):\n",
    "        \n",
    "        def gloss(s_f, s_g_f, s_d_g, t, t_g, t_d_g, al, be, gam):\n",
    "            l_gang = self.lossCE(s_d_g.squeeze(), self.lab2)+self.lossCE(t_d_g.squeeze(), self.lab2)\n",
    "            l_const = self.cos_sim(s_f.detach(), s_g_f)\n",
    "            ltv = self.smooth_func(s_g)\n",
    "            ltid = self.lossMSE()\n",
    "            \n",
    "            return l_gang + al*l_const + be*ltid + gam*ltv\n",
    "        \n",
    "        self.generator_loss_func = gloss\n",
    "        \n",
    "    def make_discriminator_loss_func(self):\n",
    "        \n",
    "        def dloss(s_d_g, t_d_g, t_d):\n",
    "            return self.lossCE(s_d_g.squeeze(), self.lab0) + self.lossCE(t_d_g.squeeze(), self.lab1) + self.lossCE(t_d.squeeze(), self.lab2)\n",
    "        \n",
    "        self.discriminator_loss_func = dloss\n",
    "        \n",
    "    def train(self, n_epo, **kwargs):\n",
    "        \n",
    "        l = min(len(self.source_train_loader), len(self.target_train_loader))\n",
    "        msimg_count = 0\n",
    "        tot_batch = 0\n",
    "        \n",
    "        for e in range(n_epo):\n",
    "            \n",
    "            source_data_iter = iter(self.source_train_loader)\n",
    "            target_data_iter = iter(self.target_train_loader)\n",
    "            \n",
    "            for i in range(l):\n",
    "                \n",
    "                self.generator_lr_sche.step()\n",
    "                self.discriminator_lr_sche.step()\n",
    "                \n",
    "                msimg_count+=1\n",
    "                \n",
    "                if msimg_count >= len(self.source_train_loader):\n",
    "                    msimg_count = 0\n",
    "                    source_data_iter = iter(self.source_train_loader)\n",
    "                    \n",
    "                source_data = source_data_iter.next()\n",
    "                target_data = target_data_iter.next()\n",
    "                \n",
    "                if self.batch_size != source_data.size(0) or self.batch_size != target_data.size(0):\n",
    "                    continue\n",
    "                \n",
    "                tot_batch+=1\n",
    "                \n",
    "                if self.use_gpu:\n",
    "                    source_data = Variable(source_data.float().cuda())\n",
    "                    target_data = Variable(target_data.float().cuda())\n",
    "                else:\n",
    "                    source_data = Variable(source_data.float())\n",
    "                    target_data = Variable(target_data.float())\n",
    "                    \n",
    "                for param in self.model['d'].parameters():\n",
    "                    param.requires_grad = True\n",
    "                self.model['d'].zero_grad()\n",
    "                \n",
    "                source_data_pad = self.pad(source_data)\n",
    "                s_f = self.model['f'](source_data_pad)\n",
    "                s_g = self.model['g'](s_f)\n",
    "                s_g = self.up(s_g)\n",
    "#                 s_g_detach = s_g.detach()\n",
    "                s_d_g = self.model['d'](s_g)\n",
    "                \n",
    "                target_data_pad = self.pad(target_data)\n",
    "                t_f = self.model['f'](target_data_pad)\n",
    "                t_g = self.model['g'](t_f)\n",
    "                t_g = self.up(t_g)\n",
    "#                 t_g_detach = t_g.detach()\n",
    "                t_d_g = self.model['d'](t_g)\n",
    "                \n",
    "                t_d = self.model['d'](target_data)\n",
    "                \n",
    "                d_loss = discriminator_loss_func(s_d_g, t_d_g, t_d)\n",
    "                d_loss.backward()\n",
    "                self.discriminator_opt.step()\n",
    "                \n",
    "                for param in self.model['d'].parameters():\n",
    "                    param.requires_grad = False\n",
    "                self.model['g'].zero_grad()\n",
    "                \n",
    "                source_data_pad = self.pad(source_data)\n",
    "                s_f = self.model['f'](source_data_pad)\n",
    "                s_g = self.model['g'](s_f)\n",
    "                s_g = self.up(s_g)\n",
    "#                 s_g_detach = s_g.detach()\n",
    "                s_d_g = self.model['d'](s_g)\n",
    "                s_g_pad = self.pad(s_g)\n",
    "                s_g_f = self.model['f'](s_g_pad)\n",
    "                \n",
    "                target_data_pad = self.pad(target_data)\n",
    "                t_f = self.model['f'](target_data_pad)\n",
    "                t_g = self.model['g'](t_f)\n",
    "                t_g = self.up(t_g)\n",
    "#                 t_g_detach = t_g.detach()\n",
    "                t_d_g = self.model['d'](t_g)\n",
    "                \n",
    "                g_loss = generator_loss_func(s_f, s_g_f, s_d_g, t, t_g, t_d_g)\n",
    "                g_loss.backward()\n",
    "                self.generator_opt.step()\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true
   },
   "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.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}