train.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import skimage\n",
    "import skimage.io\n",
    "import skimage.transform\n",
    "import numpy as np\n",
    "import random\n",
    "import scipy.misc\n",
    "import torch.optim as optim\n",
    "from models import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SEED = 123\n",
    "np.random.seed(SEED)\n",
    "torch.manual_seed(SEED)\n",
    "use_cuda = torch.cuda.is_available()\n",
    "if use_cuda:\n",
    "    torch.cuda.manual_seed(SEED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "IMAGE_SIZE = 64\n",
    "BATCH_SIZE = 32\n",
    "N_EPOCHS = 1000000\n",
    "LEARNING_RATE = 0.002\n",
    "PRINT_EVERY = 100\n",
    "SAVE_IMAGE_EVERY = 100\n",
    "SAVE_MODEL_EVERY = 1000\n",
    "IMAGE_DIR = '/mnt/disk0/kevin1kevin1k/final/images/'\n",
    "MODEL_DIR = '/mnt/disk0/kevin1kevin1k/final/models/'\n",
    "TAGS_PATH = '../hw4/data/tags_clean.csv'\n",
    "FACES_DIR = '../hw4/data/faces/'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "HAIR_REMOVE = (\n",
    "    'pubic hair',\n",
    "    'damage hair',\n",
    "    'short hair',\n",
    "    'long hair',\n",
    ")\n",
    "\n",
    "EYES_REMOVE = (\n",
    "    '11 eyes',\n",
    "    'bicolored eyes',\n",
    ")\n",
    "\n",
    "index2tags = dict()\n",
    "tag2indices = dict()\n",
    "\n",
    "is_hair = lambda tag: tag.endswith(' hair') and tag not in HAIR_REMOVE\n",
    "is_eyes = lambda tag: tag.endswith(' eyes') and tag not in EYES_REMOVE\n",
    "\n",
    "with open(TAGS_PATH) as f:\n",
    "    for line in f:\n",
    "        index, tags = line.strip().split(',')\n",
    "        tag_list = tags.split('\\t')\n",
    "        tags = []\n",
    "        for t in tag_list:\n",
    "            tag = t.split(':')[0].strip()\n",
    "            if is_hair(tag) or is_eyes(tag):\n",
    "                tags.append(tag)\n",
    "                if tag not in tag2indices:\n",
    "                    tag2indices[tag] = []\n",
    "                tag2indices[tag].append(index)\n",
    "        if len(tags) >= 1:\n",
    "            index2tags[index] = tags\n",
    "\n",
    "num_indices = len(index2tags)\n",
    "# print(num_indices)\n",
    "\n",
    "# has_hair = set()\n",
    "# has_eyes = set()\n",
    "# for tag in tag2indices:\n",
    "#     if is_hair(tag):\n",
    "#         print(tag, len(tag2indices[tag]))\n",
    "#         has_hair.update(tag2indices[tag])\n",
    "# for tag in tag2indices:\n",
    "#     if is_eyes(tag):\n",
    "#         print(tag, len(tag2indices[tag]))\n",
    "#         has_eyes.update(tag2indices[tag])\n",
    "\n",
    "# has_both = sorted(list(has_hair & has_eyes))\n",
    "# print(len(has_both))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "tags_list = sorted(list(tag2indices.keys()))\n",
    "num_tags = len(tags_list)\n",
    "def tags_to_vector(tags):\n",
    "    vec = np.zeros((num_tags))\n",
    "    for tag in tags:\n",
    "        index = tags_list.index(tag)\n",
    "        vec[index] = 1.0\n",
    "    return vec"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "images = np.zeros((num_indices, IMAGE_SIZE, IMAGE_SIZE, 3))\n",
    "tag_matrix = np.zeros((num_indices, num_tags))\n",
    "\n",
    "def fill_images_and_tag_matrix():\n",
    "    global images\n",
    "    global tag_matrix\n",
    "\n",
    "    for i, face_index in enumerate(sorted(index2tags.keys())):\n",
    "        # (96, 96, 3)\n",
    "        image = skimage.io.imread(os.path.join(FACES_DIR, face_index + '.jpg'))\n",
    "\n",
    "        # (64, 64, 3)\n",
    "        image_resized = skimage.transform.resize(image, (IMAGE_SIZE, IMAGE_SIZE), mode='constant')\n",
    "\n",
    "        images[i] = image_resized\n",
    "        tag_matrix[i] = tags_to_vector(index2tags[face_index])\n",
    "\n",
    "    images = images.transpose(0, 3, 1, 2)\n",
    "#     print(images.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fill_images_and_tag_matrix()\n",
    "images = images * 2 - 1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def enable_gradients(net):\n",
    "    for p in net.parameters():\n",
    "        p.requires_grad = True\n",
    "\n",
    "def disable_gradients(net):\n",
    "    for p in net.parameters():\n",
    "        p.requires_grad = False"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "enc = Encoder()\n",
    "dec = Decoder(num_tags, image_size=IMAGE_SIZE)\n",
    "dis = Discriminator(num_tags)\n",
    "if use_cuda:\n",
    "    enc = enc.cuda()\n",
    "    dec = dec.cuda()\n",
    "    dis = dis.cuda()\n",
    "\n",
    "betas=(0.5, 0.999)\n",
    "opt_enc_dec = optim.Adam(list(enc.parameters()) + list(dec.parameters()), lr=LEARNING_RATE, betas=betas)\n",
    "opt_dis = optim.Adam(dis.parameters(), lr=LEARNING_RATE, betas=betas)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "for ep in range(N_EPOCHS):\n",
    "    indices = random.sample(list(range(num_indices)), BATCH_SIZE)\n",
    "    X_data = images[indices]\n",
    "    X_tensor = torch.FloatTensor(X_data)\n",
    "    if use_cuda:\n",
    "        X_tensor = X_tensor.cuda()\n",
    "    X = Variable(X_tensor, requires_grad=False)\n",
    "\n",
    "    tags_data = tag_matrix[indices]\n",
    "    tags_tensor = torch.FloatTensor(tags_data)\n",
    "    if use_cuda:\n",
    "        tags_tensor = tags_tensor.cuda()\n",
    "    tags = Variable(tags_tensor, requires_grad=False)\n",
    "    \n",
    "    # update dis\n",
    "    \n",
    "    enable_gradients(dis)\n",
    "    disable_gradients(enc)\n",
    "    \n",
    "    dis.zero_grad()\n",
    "    \n",
    "    EX = enc(X)\n",
    "    P = dis(EX)\n",
    "    bce = nn.BCEWithLogitsLoss()\n",
    "    L_dis = bce(P, tags)\n",
    "    L_dis.backward(retain_graph=True)\n",
    "    opt_dis.step()\n",
    "    \n",
    "    # update enc and dec\n",
    "    \n",
    "    enable_gradients(enc)\n",
    "    enable_gradients(dec)\n",
    "    disable_gradients(dis)\n",
    "    \n",
    "    enc.zero_grad()\n",
    "    dec.zero_grad()\n",
    "    \n",
    "    X_ = dec(EX, tags)\n",
    "    \n",
    "    mse = nn.MSELoss()\n",
    "    bce = nn.BCEWithLogitsLoss()\n",
    "    lambda_E = 0.0001 * min(ep / 500000, 1.0)\n",
    "    L_enc_dec = mse(X_, X) - lambda_E * bce(P, 1 - tags)\n",
    "    L_enc_dec.backward()\n",
    "    opt_enc_dec.step()\n",
    "    \n",
    "    if (ep + 1) % SAVE_IMAGE_EVERY == 0:\n",
    "        X_data = images[3:4]\n",
    "        X_tensor = torch.FloatTensor(X_data)\n",
    "        if use_cuda:\n",
    "            X_tensor = X_tensor.cuda()\n",
    "        X = Variable(X_tensor, requires_grad=False)\n",
    "\n",
    "        tags_data = tag_matrix[3:4]\n",
    "        tags_fixed = tags_to_vector(('black hair', 'green eyes'))\n",
    "        \n",
    "        for i, tags_ in enumerate((tags_data * 0.1, tags_data * 0.5, tags_data, tags_fixed)):\n",
    "            tags_tensor = torch.FloatTensor(tags_)\n",
    "            if use_cuda:\n",
    "                tags_tensor = tags_tensor.cuda()\n",
    "            tags = Variable(tags_tensor, requires_grad=False)\n",
    "\n",
    "            EX = enc(X)\n",
    "            X_ = dec(EX, tags)\n",
    "            img = X_.data.add_(1.0).mul_(0.5).cpu().numpy().transpose(0, 2, 3, 1)[0]\n",
    "            image_path = os.path.join(IMAGE_DIR, 'test_{}_{}.jpg'.format(ep + 1, i + 1))\n",
    "            scipy.misc.imsave(image_path, img)\n",
    "\n",
    "    if (ep + 1) % SAVE_MODEL_EVERY == 0:\n",
    "        enc_path = os.path.join(MODEL_DIR, 'enc_{}.pt'.format(ep + 1))\n",
    "        dec_path = os.path.join(MODEL_DIR, 'dec_{}.pt'.format(ep + 1))\n",
    "        dis_path = os.path.join(MODEL_DIR, 'dis_{}.pt'.format(ep + 1))\n",
    "        torch.save(enc, enc_path)\n",
    "        torch.save(dec, dec_path)\n",
    "        torch.save(dis, dis_path)\n",
    "    \n",
    "    if (ep + 1) % PRINT_EVERY == 0:\n",
    "        print('ep {}, L_dis {:.4f}, L_enc_dec {:.4f}'.format(ep + 1, L_dis.data[0], L_enc_dec.data[0]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# import matplotlib.pyplot as plt\n",
    "# %matplotlib inline\n",
    "\n",
    "# img = images[3].transpose(1, 2, 0)\n",
    "# plt.figure()\n",
    "# plt.imshow(img)"
   ]
  }
 ],
 "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.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}