train_init_scr.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import ioutils as io\n",
    "import skimage\n",
    "from PIL import Image\n",
    "import glob\n",
    "import os\n",
    "import natsort\n",
    "import numpy as np\n",
    "from skimage.transform import resize\n",
    "import random\n",
    "import cv2\n",
    "from PIL import Image\n",
    "from collections import OrderedDict\n",
    "import random\n",
    "import sys\n",
    "import pickle\n",
    "import open3d as o3d\n",
    "\n",
    "from torchvision import transforms\n",
    "import torch\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "import torch.nn as nn\n",
    "import torchvision.models as models\n",
    "from torch.autograd import Variable\n",
    "\n",
    "sys.path = ['./cyclegan/'] + sys.path\n",
    "\n",
    "from my_networks import scr_net, CreateDiscriminator\n",
    "from models.base_model import BaseModel\n",
    "from models import networks\n",
    "\n",
    "from scr_utils import *\n",
    "\n",
    "import warnings\n",
    "warnings.filterwarnings('ignore')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "norm = 2\n",
    "epochs = 20\n",
    "b_size = 48\n",
    "print_interval = 100\n",
    "# train_real_proba = 1.0\n",
    "date = '0131'\n",
    "NAME = 'scr_cyclegan_b{}_{}_train_his_cyclegan'.format(b_size, date)\n",
    "\n",
    "data_path = '/scratch/zq415/work/pose/data'\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "100000 100000 (100000, 4, 4)\n",
      "1637 1637 (1637, 4, 4)\n"
     ]
    }
   ],
   "source": [
    "render_img_paths, render_scr_paths, render_poses = make_render_dataset(data_path)\n",
    "### set the render_img_paths to histogram matched\n",
    "# render_img_paths = natsort.natsorted(glob.glob('/scratch/zq415/grammar_cor/pose/pose_estimate/train_render_matched'+'/**/*.png', recursive=True))\n",
    "render_img_paths = natsort.natsorted(glob.glob('/scratch/zq415/work/pose/data/rendered_histogram'+'/**/*.png', recursive=True))\n",
    "print(len(render_img_paths), len(render_scr_paths), render_poses.shape)\n",
    "\n",
    "\n",
    "real_val_img_paths, val_scr_label_paths, val_poses = make_render_dataset(data_path, train_flag=False)\n",
    "\n",
    "\n",
    "print(len(real_val_img_paths), len(val_scr_label_paths), val_poses.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "real_val_dataset = dataset_scr(real_val_img_paths, val_poses, val_scr_label_paths,\n",
    "                              transform=transforms.Compose([fix_crop()]))\n",
    "real_val_dataloader = DataLoader(real_val_dataset, batch_size=100,\n",
    "                            shuffle=False, num_workers=4)\n",
    "\n",
    "\n",
    "render_train_dataset = dataset_scr(render_img_paths, render_poses, render_scr_paths,\n",
    "                               transform=transforms.Compose([random_crop()]))\n",
    "render_train_dataloader = DataLoader(render_train_dataset, batch_size=b_size,\n",
    "                            shuffle=True, num_workers=4)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def train(model, device, train_loader, optimizer, criterion, epoch, record):\n",
    "    model.train()\n",
    "    running_loss = 0.0\n",
    "    for i_batch, sample_batched in enumerate(train_loader):\n",
    "        inputs, labels = sample_batched['image'], sample_batched['scr']\n",
    "        inputs, labels = inputs.to(device, dtype=torch.float), labels.to(device, dtype=torch.float)\n",
    "        # zero the parameter gradients\n",
    "        optimizer.zero_grad()\n",
    "        \n",
    "        outputs = model(inputs)[1]\n",
    "        loss = criterion(outputs, torch.nn.functional.interpolate(labels, scale_factor=1/8.0, mode='bilinear'),p=norm)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        running_loss += loss.item()\n",
    "        if (i_batch+1) % print_interval == 0:\n",
    "            out_str = \"epoch {}, batch {}, current loss {}.\\n\".format(epoch+1,i_batch,running_loss/print_interval)\n",
    "            print(out_str)\n",
    "            record.write(out_str)\n",
    "            record.flush()\n",
    "            \n",
    "            img = inputs.permute(0,2,3,1).cpu().detach().numpy()[0,...]\n",
    "            label = labels.permute(0,2,3,1).cpu().detach().numpy()[0,...]\n",
    "            predict = torch.nn.functional.interpolate(outputs,scale_factor=8.0, mode='bilinear').permute(0,2,3,1).cpu().detach().numpy()[0,...]\n",
    "            print(img.shape)\n",
    "            running_loss = 0.0\n",
    "            \n",
    "def test(net, dataloader):\n",
    "    net.eval()\n",
    "    scr_predicts, true_poses = [], []\n",
    "    scr_labels = []\n",
    "\n",
    "    for i_batch, sample_batched in tqdm(enumerate(dataloader)):\n",
    "        inputs, scrs, poses = sample_batched['image'], sample_batched['scr'], sample_batched['pose']\n",
    "        inputs = inputs.to(device, dtype=torch.float)\n",
    "        with torch.no_grad():\n",
    "            scr_predict = net(inputs)[1]\n",
    "\n",
    "        scr_predicts.append(scr_predict.cpu().numpy())\n",
    "        scr_labels.append(torch.nn.functional.interpolate(scrs.cpu(), scale_factor=1/8.0, mode='bilinear').numpy())\n",
    "        true_poses.append(poses.numpy())\n",
    "        \n",
    "    scr_predicts = np.concatenate(scr_predicts)\n",
    "    scr_labels = np.concatenate(scr_labels)\n",
    "    true_poses = np.concatenate(true_poses)\n",
    "    print(scr_predicts.shape, scr_labels.shape, true_poses.shape)\n",
    "    return scr_predicts, scr_labels, true_poses        \n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "using GPU for training\n",
      "cuda:0\n"
     ]
    }
   ],
   "source": [
    "if torch.cuda.is_available():\n",
    "    device = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "    torch.backends.cudnn.benchmark = True\n",
    "    torch.backends.cudnn.enabled = True\n",
    "    print('using GPU for training')\n",
    "else:\n",
    "    device = torch.device('cpu')\n",
    "    print('using CPU for training')\n",
    "print(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "11572803\n"
     ]
    }
   ],
   "source": [
    "#### sc network\n",
    "net = scr_net()\n",
    "    \n",
    "net.to(device)\n",
    "if torch.cuda.device_count() > 1:\n",
    "    print(\"Let's use\", torch.cuda.device_count(), \"GPUs!\")\n",
    "    net = nn.DataParallel(net)\n",
    "    \n",
    "model_total_params = sum(p.numel() for p in net.parameters() if p.requires_grad)\n",
    "print(model_total_params)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(net.parameters(), lr=0.001, weight_decay=0.00001)\n",
    "criterion = maskedPnorm\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "norm: 2, batch_size: 48.\n",
      "\n",
      "epoch 1, batch 99, current loss 2.929746297597885.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 1, batch 199, current loss 1.3059980708360672.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 1, batch 299, current loss 1.0534222757816314.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 1, batch 399, current loss 0.8925641840696334.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 1, batch 499, current loss 0.7500084942579269.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 1, batch 599, current loss 0.609468070268631.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 1, batch 699, current loss 0.5636269581317902.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 1, batch 799, current loss 0.5055500927567482.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 1, batch 899, current loss 0.47024311661720275.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 1, batch 999, current loss 0.6267395177483559.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 1, batch 1099, current loss 0.49938901782035827.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 1, batch 1199, current loss 0.4156216859817505.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 1, batch 1299, current loss 0.412984139919281.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 1, batch 1399, current loss 0.3664962854981422.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 1, batch 1499, current loss 0.3589353629946709.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 1, batch 1599, current loss 0.3633881886303425.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 1, batch 1699, current loss 0.427628476023674.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 1, batch 1799, current loss 0.3818532687425613.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 1, batch 1899, current loss 0.326456238925457.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 1, batch 1999, current loss 0.29733935579657556.\n",
      "\n",
      "(320, 320, 3)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "17it [01:33,  5.50s/it]\n",
      "  0%|          | 0/1637 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1637, 3, 42, 74) (1637, 3, 42, 74) (1637, 4, 4)\n",
      "real_mean_scr_error: 0.635046660900116.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1637/1637 [01:10<00:00, 23.31it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "real_val, r_median: 107.08593896958797, t_median: 9.7109539441329.\n",
      "\n",
      "epoch 2, batch 99, current loss 0.41751931309700013.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 2, batch 199, current loss 0.3067796126008034.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 2, batch 299, current loss 0.29421899244189265.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 2, batch 399, current loss 0.3063937829434872.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 2, batch 499, current loss 0.30761067882180215.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 2, batch 599, current loss 0.2903783640265465.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 2, batch 699, current loss 0.269980860799551.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 2, batch 799, current loss 0.2578193497657776.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 2, batch 899, current loss 0.46974185526371004.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 2, batch 999, current loss 0.45532130151987077.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 2, batch 1099, current loss 0.31641054034233096.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 2, batch 1199, current loss 0.3063148435950279.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 2, batch 1299, current loss 0.2910547052323818.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 2, batch 1399, current loss 0.265308735370636.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 2, batch 1499, current loss 0.2805190445482731.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 2, batch 1599, current loss 0.2636671842634678.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 2, batch 1699, current loss 0.24670111179351806.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 2, batch 1799, current loss 0.23231440290808678.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 2, batch 1899, current loss 0.26980318263173103.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 2, batch 1999, current loss 0.24031519368290902.\n",
      "\n",
      "(320, 320, 3)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "17it [01:22,  4.86s/it]\n",
      "  0%|          | 0/1637 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1637, 3, 42, 74) (1637, 3, 42, 74) (1637, 4, 4)\n",
      "real_mean_scr_error: 0.5896453261375427.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1637/1637 [01:06<00:00, 24.56it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "real_val, r_median: 117.95274406327417, t_median: 8.58408712356746.\n",
      "\n",
      "epoch 3, batch 99, current loss 0.24906323105096817.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 3, batch 199, current loss 0.23136330187320708.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 3, batch 299, current loss 0.2278200300037861.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 3, batch 399, current loss 0.23351360261440277.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 3, batch 499, current loss 0.21509096547961234.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 3, batch 599, current loss 0.22863618552684783.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 3, batch 699, current loss 0.2345196269452572.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 3, batch 799, current loss 0.2929271651804447.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 3, batch 899, current loss 0.2795758725702763.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 3, batch 999, current loss 0.25489909663796423.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 3, batch 1099, current loss 0.21496190384030342.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 3, batch 1199, current loss 0.22460278928279875.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 3, batch 1299, current loss 0.20632360458374024.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 3, batch 1399, current loss 0.2343883466720581.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 3, batch 1499, current loss 0.23505480006337165.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 3, batch 1599, current loss 0.2177892179787159.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 3, batch 1699, current loss 0.2573668107390404.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 3, batch 1799, current loss 0.2193606613576412.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 3, batch 1899, current loss 0.2057410617172718.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 3, batch 1999, current loss 0.22785977616906167.\n",
      "\n",
      "(320, 320, 3)\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "17it [00:57,  3.40s/it]\n",
      "  0%|          | 0/1637 [00:00<?, ?it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "(1637, 3, 42, 74) (1637, 3, 42, 74) (1637, 4, 4)\n",
      "real_mean_scr_error: 0.6277364492416382.\n",
      "\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1637/1637 [01:07<00:00, 24.32it/s]\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "real_val, r_median: 82.96733770109029, t_median: 8.992603107441212.\n",
      "\n",
      "epoch 4, batch 99, current loss 0.21396866723895072.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 4, batch 199, current loss 0.2026227517426014.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 4, batch 299, current loss 0.20354780688881874.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 4, batch 399, current loss 0.22474947839975357.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 4, batch 499, current loss 0.9113996756076813.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 4, batch 599, current loss 0.4280404564738274.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 4, batch 699, current loss 0.3185541374981403.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 4, batch 799, current loss 0.28733678847551347.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 4, batch 899, current loss 0.25420532435178755.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 4, batch 999, current loss 0.23524843096733095.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 4, batch 1099, current loss 0.24590095445513727.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 4, batch 1199, current loss 0.240287850946188.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 4, batch 1299, current loss 0.22643935561180115.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 4, batch 1399, current loss 0.21961886569857597.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 4, batch 1499, current loss 0.2362351283431053.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 4, batch 1599, current loss 0.22153758823871614.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 4, batch 1699, current loss 0.21393025055527687.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 4, batch 1799, current loss 0.21557649448513985.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 4, batch 1899, current loss 0.2002492789924145.\n",
      "\n",
      "(320, 320, 3)\n",
      "epoch 4, batch 1999, current loss 0.20588487163186073.\n",
      "\n",
      "(320, 320, 3)\n"
     ]
    }
   ],
   "source": [
    "record = open('./save_info/'+ NAME +'.txt','w+')\n",
    "\n",
    "basic_str = \"norm: {}, batch_size: {}.\\n\".format(norm, b_size) \n",
    "print(basic_str)\n",
    "record.write(basic_str)\n",
    "record.flush()\n",
    "\n",
    "for epoch in range(epochs):\n",
    "    \n",
    "    train(net, device, render_train_dataloader, optimizer, criterion, epoch, record)\n",
    "    \n",
    "    real_val_scr_predicts, val_scrs, val_true_poses = test(net, real_val_dataloader)\n",
    "    \n",
    "    out_str = 'real_mean_scr_error: {}.\\n'.format(compute_mean_error(val_scrs, real_val_scr_predicts))\n",
    "    print(out_str)\n",
    "    record.write(out_str)\n",
    "    record.flush()\n",
    "    \n",
    "    real_val_median = get_median(real_val_scr_predicts, val_true_poses)\n",
    "    out_str = 'real_val, r_median: {}, t_median: {}.\\n'.format(real_val_median[0], real_val_median[1])\n",
    "    print(out_str)\n",
    "    record.write(out_str)\n",
    "    record.flush()\n",
    "    \n",
    "    torch.save(net.state_dict(), './check_point/scr_b{}_e{}_norm{}_{}_train_his_cyclegan.pth'.format(b_size,epoch,norm,date))\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.6.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}