Domain confusion.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import torchvision.transforms as transforms\n",
    "import torch.utils.data as data\n",
    "import os\n",
    "from PIL import Image\n",
    "import numpy as np\n",
    "import glob\n",
    "import random\n",
    "from models import Encoder, Tnet, _F_, Destructor_domain, Destructor_shape\n",
    "import torch.nn as nn\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "class XDataset(data.Dataset):\n",
    "    def __init__(self, image_root, point_cloud_root, transform=None, use_2048 = True):\n",
    "        self.image_root = image_root\n",
    "        self.point_cloud_root = point_cloud_root\n",
    "        self.transform = transform\n",
    "        self.use_2048 = use_2048\n",
    "        self.pc_list = glob.glob(point_cloud_root+\"*/*/*2048.npy\")\n",
    "        self.syn_list = glob.glob(image_root+\"*/*/rendering/*.png\")\n",
    "        self.N = len(self.syn_list)\n",
    "        self.l = glob.glob('pix3d/*/*.png')\n",
    "        self.n = len(self.l)\n",
    "    \n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        \n",
    "        image = Image.open(self.syn_list[random.randint(0, self.N-1)]).convert('RGB')\n",
    "        real_image = Image.open(self.l[random.randint(0, self.n-1)]).convert('RGB')\n",
    "        \n",
    "        if self.transform is not None:\n",
    "            image = self.transform(image)\n",
    "            real_image = self.transform(real_image)   \n",
    "            \n",
    "        point_cloud = np.load(self.pc_list[index])\n",
    "        \n",
    "        return image, point_cloud, real_image\n",
    "\n",
    "    def __len__(self):\n",
    "        return len(self.pc_list)\n",
    "\n",
    "def get_loader(image_root, point_cloud_root, use_2048, transform, batch_size, shuffle, num_workers):\n",
    "    \"\"\"Returns torch.utils.data.DataLoader for custom X dataset.\"\"\"\n",
    "    \n",
    "    xdataset = XDataset(image_root, point_cloud_root, transform = transform, use_2048 = use_2048)\n",
    "    \n",
    "    data_loader = torch.utils.data.DataLoader(dataset=xdataset, \n",
    "                                              batch_size=batch_size,\n",
    "                                              shuffle=shuffle,\n",
    "                                              num_workers=num_workers)\n",
    "    return data_loader\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Training Configuration\n",
    "image_root = \"/datasets/cs253-wi20-public/ShapeNetRendering/\"\n",
    "point_cloud_root = \"/datasets/cs253-wi20-public/ShapeNet_pointclouds/\"\n",
    "\n",
    "batch_size = 16\n",
    "shuffle = True\n",
    "num_workers = 8\n",
    "use_2048 = True\n",
    "img_size = 256\n",
    "transform = transforms.Compose([transforms.Resize(img_size,interpolation=2),transforms.CenterCrop(img_size),transforms.ToTensor()])\n",
    "data_loader = get_loader(image_root, point_cloud_root, use_2048, transform, batch_size, shuffle, num_workers)\n",
    "def init_weights(m):\n",
    "    if isinstance(m, nn.Conv1d) or isinstance(m, nn.Linear):\n",
    "        nn.init.xavier_uniform(m.weight.data)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "criterion = nn.BCELoss()\n",
    "num_epochs = 1000\n",
    "learning_rate = 1e-5\n",
    "f = _F_().apply(init_weights).cuda()\n",
    "d_domain = Destructor_domain().apply(init_weights).cuda()\n",
    "d_shape = Destructor_shape().apply(init_weights).cuda()\n",
    "optimizerf = torch.optim.Adam(f.parameters(), lr=learning_rate)\n",
    "optimizer_domain = torch.optim.Adam(d_domain.parameters(), lr=learning_rate)\n",
    "optimizer_shape = torch.optim.Adam(d_shape.parameters(), lr=learning_rate)\n",
    "\n",
    "enc = torch.load('encoder-Copy1').cuda()\n",
    "\n",
    "for epoch in range(num_epochs):\n",
    "    loss_f = 0.0\n",
    "    loss_domain = 0.0\n",
    "    loss_shape = 0.0\n",
    "    \n",
    "    for (i_syn, pc_syn, i_real) in data_loader:\n",
    "        b_size = i_syn.shape[0]\n",
    "        if torch.cuda.is_available():\n",
    "            i_syn = i_syn.cuda()\n",
    "            i_real = i_real.cuda()\n",
    "            pc_syn = pc_syn.transpose(1,2).cuda()\n",
    "        \n",
    "        optimizer_domain.zero_grad()\n",
    "        #real =0; syn =1\n",
    "        enc_syn = f(i_syn)\n",
    "        output = d_domain(enc_syn)\n",
    "        errD_syn = criterion(output, torch.full((b_size,), 1, device=i_syn.device))\n",
    "        enc_real = f(i_real)\n",
    "        output1 = d_domain(enc_real)\n",
    "        errD_real = criterion(output1, torch.full((b_size,), 0, device=i_syn.device))\n",
    "        errD = errD_syn + errD_real\n",
    "        errD.backward()\n",
    "        optimizer_domain.step()\n",
    "        \n",
    "        optimizerf.zero_grad()\n",
    "        enc_syn = f(i_syn)\n",
    "        output = d_domain(enc_syn)\n",
    "        errG_syn = criterion(output, torch.full((b_size,), 0, device=i_syn.device))\n",
    "        enc_real = f(i_real)\n",
    "        output1 = d_domain(enc_real)\n",
    "        errG_real = criterion(output1, torch.full((b_size,), 1, device=i_syn.device))\n",
    "        errG = 0.01*(errG_syn + errG_real)\n",
    "        errG.backward()\n",
    "        optimizerf.step()\n",
    "        \n",
    "        optimizer_shape.zero_grad()\n",
    "        enc_syn = f(i_syn)\n",
    "        enc_real = f(i_real)\n",
    "        enc_pc = enc(pc_syn.float())\n",
    "        output = d_shape(enc_pc)\n",
    "        errD_pc = criterion(output, torch.full((b_size,), 1, device=i_syn.device))\n",
    "        output1 = d_shape(enc_syn)\n",
    "        errD_syn1 = criterion(output1, torch.full((b_size,), 0, device=i_syn.device))\n",
    "        output2 = d_shape(enc_real)\n",
    "        errD_real1 = criterion(output2, torch.full((b_size,), 0, device=i_syn.device))\n",
    "        lossD = errD_pc + errD_syn1 + errD_real1\n",
    "        lossD.backward()\n",
    "        optimizer_shape.step()\n",
    "        \n",
    "        optimizerf.zero_grad()\n",
    "        enc_syn = f(i_syn)\n",
    "        enc_real = f(i_real)\n",
    "        enc_pc = enc(pc_syn.float())\n",
    "        output = d_shape(enc_pc)\n",
    "        output1 = d_shape(enc_syn)\n",
    "        output2 = d_shape(enc_real)\n",
    "        errG_pc = criterion(output, torch.full((b_size,), 0, device=i_syn.device))\n",
    "        errG_syn1 = criterion(output1, torch.full((b_size,), 1, device=i_syn.device))\n",
    "        errG_real1 = criterion(output2, torch.full((b_size,), 1, device=i_syn.device))\n",
    "        lossG = 0.01*(errG_pc + errG_syn1 + errG_real1)\n",
    "        lossG.backward()\n",
    "        optimizerf.step()\n",
    "        \n",
    "        loss_f+=errG.data.detach() +lossG.data.detach() \n",
    "        loss_domain+=errD.data.detach() \n",
    "        loss_shape +=lossD.data.detach() \n",
    "        \n",
    "    print('epoch [{}/{}], F_loss:{:.4f}, Domain_loss:{:.4f}, Shape_loss:{:.4f}'.format(epoch + 1, num_epochs, loss_f/b_size, loss_domain/b_size, loss_shape/b_size ))\n",
    "        \n",
    "                 "
   ]
  },
  {
   "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.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}