pytorch.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torchvision.transforms as tf\n",
    "from PIL import Image\n",
    "import torch.autograd as autograd"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "The images files have the form \"ID_2m_0P_xV_yH_z.jpg\" where ID is the ID of the person, 2m is fixed, 0P means head pose of 0 degrees (only head pose used in this notebook)\n",
    "x is the vertical orientation, y is the horizontal orientation and z is either L for left or R for right eye (note that the right eye patch was flipped horizontally).\n",
    "In training the images are grouped as follows:\n",
    "For a given person and a given eye (R or L) all orientations are grouped together. One element of the data set is of the form\n",
    "imgs_r,angles_r,labels,imgs_t,angles_g where imgs_r is considered the \"real\" image with orientation angles_r, or x_r in the paper,\n",
    "imgs_t with orientation angles_g is the image of the same person with different orientation (could be the same image since we go through a double loop) and the label is the ID of the person"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "class MyDataset(torch.utils.data.Dataset):\n",
    "    def __init__(self, dir_path, transform=None):\n",
    "        super().__init__()\n",
    "        self.transform = transform\n",
    "        self.ids=50\n",
    "        self.data_path = dir_path\n",
    "        self.file_names = [f for f in os.listdir(self.data_path)\n",
    "                      if f.endswith('.jpg')]\n",
    "        self.file_dict = dict()\n",
    "        for f_name in self.file_names:\n",
    "            fields = f_name.split('.')[0].split('_')\n",
    "            identity = fields[0]\n",
    "            head_pose = fields[2]\n",
    "            side = fields[-1]\n",
    "            key = '_'.join([identity, head_pose, side])\n",
    "            if key not in self.file_dict.keys():\n",
    "                self.file_dict[key] = []\n",
    "                self.file_dict[key].append(f_name)\n",
    "            else:\n",
    "                self.file_dict[key].append(f_name)\n",
    "        self.train_images = []\n",
    "        self.train_angles_r = []\n",
    "        self.train_labels = []\n",
    "        self.train_images_t = []\n",
    "        self.train_angles_g = []\n",
    "\n",
    "        self.test_images = []\n",
    "        self.test_angles_r = []\n",
    "        self.test_labels = []\n",
    "        self.test_images_t = []\n",
    "        self.test_angles_g = []\n",
    "        self.preprocess()\n",
    "    def preprocess(self):\n",
    "\n",
    "        for key in self.file_dict.keys():\n",
    "\n",
    "            if len(self.file_dict[key]) == 1:\n",
    "                continue\n",
    "\n",
    "            idx = int(key.split('_')[0])\n",
    "            flip = 1\n",
    "            if key.split('_')[-1] == 'R':\n",
    "                flip = -1\n",
    "\n",
    "            for f_r in self.file_dict[key]:\n",
    "\n",
    "                file_path = os.path.join(self.data_path, f_r)\n",
    "\n",
    "                h_angle_r = flip * float(\n",
    "                    f_r.split('_')[-2].split('H')[0]) / 15.0\n",
    "                    \n",
    "                v_angle_r = float(\n",
    "                    f_r.split('_')[-3].split('V')[0]) / 10.0\n",
    "                    \n",
    "\n",
    "                for f_g in self.file_dict[key]:\n",
    "\n",
    "                    file_path_t = os.path.join(self.data_path, f_g)\n",
    "\n",
    "                    h_angle_g = flip * float(\n",
    "                        f_g.split('_')[-2].split('H')[0]) / 15.0\n",
    "                        \n",
    "                    v_angle_g = float(\n",
    "                        f_g.split('_')[-3].split('V')[0]) / 10.0\n",
    "                        \n",
    "\n",
    "                    if idx <= self.ids:\n",
    "                        self.train_images.append(file_path)\n",
    "                        self.train_angles_r.append([h_angle_r, v_angle_r])\n",
    "                        self.train_labels.append(idx - 1)\n",
    "                        self.train_images_t.append(file_path_t)\n",
    "                        self.train_angles_g.append([h_angle_g, v_angle_g])\n",
    "                    else:\n",
    "                        self.test_images.append(file_path)\n",
    "                        self.test_angles_r.append([h_angle_r, v_angle_r])\n",
    "                        self.test_labels.append(idx - 1)\n",
    "                        self.test_images_t.append(file_path_t)\n",
    "                        self.test_angles_g.append([h_angle_g, v_angle_g])\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        return (\n",
    "            self.transform(Image.open(self.train_images[index])),\n",
    "                torch.tensor(self.train_angles_r[index]),\n",
    "                self.train_labels[index],\n",
    "            self.transform(Image.open(self.train_images_t[index])),\n",
    "                torch.tensor(self.train_angles_g[index]))\n",
    "        \n",
    "    def __len__(self):\n",
    "        return len(self.train_images)\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "transform=tf.Compose([tf.ToTensor(),tf.Resize((64,64),antialias=True)])\n",
    "# dataset=MyDataset(dir_path='/home/user/Downloads/dataset/0P',transform=transform)\n",
    "dataset=MyDataset(dir_path='/home/user/Downloads/dataset/0P',transform=transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "train_loader = torch.utils.data.DataLoader(dataset, batch_size=32, shuffle=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([32, 3, 64, 64]) torch.Size([32, 2]) torch.Size([32]) torch.Size([32, 3, 64, 64]) torch.Size([32, 2])\n"
     ]
    }
   ],
   "source": [
    "imgs_r,angles_r,labels,imgs_g,angles_g=next(iter(train_loader))\n",
    "print(imgs_r.shape,angles_r.shape,labels.shape,imgs_g.shape,angles_g.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "device='cuda' if torch.cuda.is_available() else 'cpu'\n",
    "#device='cpu'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "from networks import Discriminator,Generator2\n",
    "# if os.path.isfile('discriminator.pth'):\n",
    "#     discriminator=torch.load('discriminator.pth')\n",
    "#     print('loaded discriminator')\n",
    "# else:\n",
    "#     discriminator=Discriminator()\n",
    "#     print('created discriminator')\n",
    "# if os.path.isfile('generator.pth'):\n",
    "#     generator=torch.load('generator.pth')\n",
    "#     print('loaded generator')\n",
    "# else:\n",
    "#     generator=Generator2()\n",
    "#     print('created generator')\n",
    "\n",
    "discriminator=Discriminator()\n",
    "generator=Generator2()\n",
    "generator=generator.to(device)\n",
    "discriminator=discriminator.to(device)\n",
    "LR = 5e-5\n",
    "beta1=0.5\n",
    "beta2=0.999\n",
    "optimizer_g = torch.optim.Adam(generator.parameters(), LR,betas=(beta1, beta2))\n",
    "optimizer_d = torch.optim.Adam(discriminator.parameters(), LR,betas=(beta1, beta2))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/user/anaconda3/envs/pytorch2/lib/python3.11/site-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.\n",
      "  warnings.warn(\n",
      "/home/user/anaconda3/envs/pytorch2/lib/python3.11/site-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=VGG19_Weights.IMAGENET1K_V1`. You can also use `weights=VGG19_Weights.DEFAULT` to get the most up-to-date weights.\n",
      "  warnings.warn(msg)\n"
     ]
    }
   ],
   "source": [
    "from loss_network import LossNetwork\n",
    "loss_network=LossNetwork()\n",
    "loss_network=loss_network.to(device)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [],
   "source": [
    "from loss import content_style_loss,adv_loss_d,adv_loss_g,gaze_loss_d,gaze_loss_g,reconstruction_loss"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "def generator_step(generator,discriminator,loss_network,imgs_r,imgs_t,angles_r,angles_g):\n",
    "    optimizer_g.zero_grad()\n",
    "    generator.train()\n",
    "    discriminator.eval()\n",
    "    x_g=generator(imgs_r,angles_g)\n",
    "    x_recon=generator(x_g,angles_r)\n",
    "    loss_adv=adv_loss_g(discriminator,imgs_r,x_g)\n",
    "    loss2=content_style_loss(loss_network,x_g,imgs_t)\n",
    "    loss_p=loss2[0]+loss2[1]\n",
    "    loss_gg=gaze_loss_g(discriminator,x_g,angles_g)\n",
    "    loss_recon=reconstruction_loss(generator,imgs_r,x_recon)\n",
    "    loss=loss_adv+100*loss_p+5*loss_gg+50*loss_recon\n",
    "    loss.backward()\n",
    "    optimizer_g.step()\n",
    "    return loss.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "def discriminator_step(generator,discriminator,imgs_r,imgs_t,angles_r,angles_g):\n",
    "    optimizer_d.zero_grad()\n",
    "    generator.eval()\n",
    "    discriminator.train()\n",
    "    x_g=generator(imgs_r,angles_g)\n",
    "    loss1=adv_loss_d(discriminator,imgs_r,x_g)\n",
    "    loss2=gaze_loss_d(discriminator,imgs_r,angles_r)\n",
    "    loss=loss1+5*loss2\n",
    "    loss.backward()\n",
    "    optimizer_d.step()\n",
    "    return loss.item()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import numpy as np\n",
    "def recover_image(img):\n",
    "    img=img.cpu().numpy().transpose(0, 2, 3, 1)*255\n",
    "    return img.astype(np.uint8)\n",
    "def save_images(imgs, filename):\n",
    "    height=recover_image(imgs[0])[0].shape[0]\n",
    "    width=recover_image(imgs[0])[0].shape[1]\n",
    "    total_width=width*len(imgs)\n",
    "    \n",
    "    new_im = Image.new('RGB', (total_width+len(imgs), height))\n",
    "    for i,img in enumerate(imgs):\n",
    "        result = Image.fromarray(recover_image(img)[0])\n",
    "        new_im.paste(result, (i*width+i,0))\n",
    "    new_im.save(filename)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/bin/bash: /home/user/anaconda3/lib/libtinfo.so.6: no version information available (required by /bin/bash)\n"
     ]
    }
   ],
   "source": [
    "!mkdir -p debug"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "-0.31358757615089417 2.9682517051696777\n",
      "-0.2342233806848526 2.3327081203460693\n",
      "-0.2994817793369293 1.940845012664795\n",
      "-0.3117067217826843 1.8902649879455566\n",
      "-0.35777783393859863 1.546557068824768\n",
      "-0.37074172496795654 1.5946787595748901\n",
      "-0.3110412061214447 1.4235047101974487\n",
      "-0.3506832420825958 1.4569635391235352\n",
      "-0.27027928829193115 1.416990876197815\n",
      "-0.3218216896057129 1.3143155574798584\n",
      "-0.24718238413333893 1.4948667287826538\n",
      "-0.2606305778026581 1.148314118385315\n",
      "-0.2549508810043335 1.1774322986602783\n",
      "-0.25507697463035583 1.0714054107666016\n",
      "-0.19905276596546173 1.139740228652954\n",
      "-0.3423738479614258 1.3190399408340454\n",
      "-0.39515891671180725 1.027995228767395\n",
      "-0.29178595542907715 1.0882527828216553\n",
      "-0.26527005434036255 1.256840467453003\n",
      "-0.18343782424926758 1.3145817518234253\n",
      "-0.27738529443740845 0.972407877445221\n",
      "-0.3530397415161133 1.5277353525161743\n",
      "-0.21616946160793304 1.0154327154159546\n",
      "-0.294808030128479 0.9910871386528015\n",
      "-0.1701388657093048 1.01249361038208\n",
      "-0.21900419890880585 0.9458963871002197\n",
      "-0.22769582271575928 0.8741225004196167\n",
      "-0.24936321377754211 0.8527806401252747\n",
      "-0.15437406301498413 0.8539791107177734\n",
      "-0.20587533712387085 0.7850728034973145\n",
      "-0.3145527243614197 0.9203860759735107\n",
      "-0.218935027718544 0.9283149242401123\n",
      "-0.20624658465385437 0.7164262533187866\n",
      "-0.2460257112979889 0.939175546169281\n",
      "-0.20801356434822083 0.8336048722267151\n",
      "-0.24870239198207855 0.8426695466041565\n",
      "-0.2181064635515213 0.8064556121826172\n",
      "-0.13996881246566772 0.7247059345245361\n",
      "-0.1924586296081543 0.8698497414588928\n",
      "-0.1286771148443222 0.7677893042564392\n",
      "-0.3115207552909851 0.8459518551826477\n",
      "-0.17218442261219025 0.6945117115974426\n",
      "-0.1413082331418991 0.889028012752533\n",
      "-0.1981084793806076 0.6468481421470642\n",
      "-0.21922773122787476 0.8876541256904602\n",
      "-0.2528008818626404 0.8339278101921082\n",
      "-0.2000853270292282 0.7747090458869934\n",
      "-0.18398559093475342 0.6826373338699341\n",
      "-0.248726949095726 0.744856059551239\n",
      "-0.17214909195899963 0.6839426159858704\n",
      "-0.19519953429698944 0.7576154470443726\n",
      "-0.25544750690460205 0.8359072804450989\n",
      "-0.20343376696109772 0.716457724571228\n",
      "-0.17487341165542603 0.6601396203041077\n",
      "-0.11138652265071869 0.6958702802658081\n",
      "-0.1899452954530716 0.7297855019569397\n",
      "-0.22113285958766937 0.6233486533164978\n",
      "-0.2714031934738159 0.643134355545044\n",
      "-0.08174295723438263 0.6397721171379089\n",
      "-0.2873768210411072 0.7286661863327026\n",
      "-0.20197442173957825 0.7159639596939087\n",
      "-0.08778241276741028 0.6032082438468933\n",
      "-0.2263650894165039 0.6563352346420288\n",
      "-0.16922613978385925 0.6244162321090698\n",
      "-0.2142571657896042 0.6349068880081177\n",
      "-0.23447218537330627 0.6694350242614746\n",
      "-0.1820065826177597 0.6288104057312012\n",
      "-0.19174857437610626 0.5930564403533936\n",
      "-0.19961565732955933 0.7379372715950012\n",
      "-0.10162173211574554 0.6217471957206726\n",
      "-0.1283857375383377 0.5884793400764465\n",
      "-0.20148228108882904 0.8070382475852966\n",
      "-0.1450543850660324 0.577135443687439\n",
      "-0.16814084351062775 0.6053802967071533\n",
      "-0.16650749742984772 0.6794916987419128\n",
      "-0.1094692274928093 0.6641079187393188\n",
      "-0.1749386340379715 0.5995951890945435\n",
      "-0.12620647251605988 0.5849091410636902\n",
      "-0.14548973739147186 0.591703474521637\n",
      "-0.07356191426515579 0.6142909526824951\n",
      "-0.18101532757282257 0.6388349533081055\n",
      "-0.1277286857366562 0.5913575887680054\n",
      "-0.16815418004989624 0.579506516456604\n",
      "-0.16377589106559753 0.5410429835319519\n",
      "-0.17506864666938782 0.5722120404243469\n",
      "-0.22073368728160858 0.5803152918815613\n",
      "-0.13670718669891357 0.5822076797485352\n",
      "-0.12052033841609955 0.6082519292831421\n",
      "-0.12457777559757233 0.5402642488479614\n",
      "-0.15331418812274933 0.5689026713371277\n",
      "-0.21199610829353333 0.5806924104690552\n",
      "-0.14264771342277527 0.545079231262207\n",
      "-0.16474361717700958 0.5501030087471008\n",
      "-0.18695227801799774 0.6111180782318115\n",
      "-0.17023415863513947 0.5793765187263489\n",
      "-0.1202152818441391 0.5521765947341919\n",
      "-0.22414188086986542 0.5701684355735779\n",
      "-0.2047164887189865 0.5739836096763611\n",
      "-0.12677207589149475 0.6252591013908386\n",
      "-0.16755664348602295 0.6043359041213989\n",
      "-0.1583997905254364 0.5671118497848511\n",
      "-0.16385102272033691 0.5477347373962402\n",
      "-0.11700200289487839 0.5958592891693115\n",
      "-0.13615739345550537 0.54768967628479\n",
      "-0.18961189687252045 0.5318083167076111\n",
      "-0.28849998116493225 0.5469931364059448\n",
      "-0.226436585187912 0.5431898236274719\n",
      "-0.28204429149627686 1.716888427734375\n",
      "-0.1459781974554062 0.5918998122215271\n",
      "-0.18247447907924652 0.5270678997039795\n",
      "-0.12524203956127167 0.646364152431488\n",
      "-0.19915027916431427 0.47503164410591125\n",
      "-0.12613466382026672 0.6336413621902466\n",
      "-0.1677912324666977 0.6100881099700928\n",
      "-0.18271203339099884 0.5327352285385132\n",
      "-0.12068776786327362 0.5656340718269348\n",
      "-0.17093437910079956 0.5324371457099915\n",
      "-0.11619863659143448 0.4968976080417633\n",
      "-0.15109625458717346 0.6003345847129822\n",
      "-0.1048649325966835 0.5437024831771851\n",
      "-0.2734958827495575 0.549656867980957\n",
      "-0.20283864438533783 0.4737273156642914\n",
      "-0.13598135113716125 0.5222673416137695\n",
      "-0.12276109308004379 0.5822096467018127\n",
      "-0.14700160920619965 0.5351520776748657\n",
      "-0.07705244421958923 0.48460277915000916\n",
      "-0.2332165688276291 0.5189864635467529\n",
      "-0.138630673289299 0.5181263089179993\n",
      "-0.1285402923822403 0.5241039395332336\n",
      "-0.2555902898311615 0.4507821798324585\n",
      "-0.12912768125534058 0.45456668734550476\n",
      "-0.09727749228477478 0.5162322521209717\n",
      "-0.18934066593647003 0.5176610350608826\n",
      "-0.09420374035835266 0.4612128734588623\n",
      "-0.1662670522928238 0.45739874243736267\n",
      "-0.23151379823684692 0.5058619976043701\n",
      "-0.16297686100006104 0.4919949173927307\n",
      "-0.18529120087623596 0.5011815428733826\n",
      "-0.11907671391963959 0.48076778650283813\n",
      "-0.07367987185716629 0.4820939004421234\n",
      "-0.05300775542855263 0.4749205410480499\n",
      "-0.1179971843957901 0.5400040745735168\n",
      "-0.1433565318584442 0.45712316036224365\n",
      "-0.28409111499786377 0.5379238724708557\n",
      "-0.17242488265037537 0.48743125796318054\n",
      "-0.15508303046226501 0.4979252815246582\n",
      "-0.16327926516532898 0.4625117778778076\n",
      "-0.11855760216712952 0.5170938372612\n",
      "-0.16486935317516327 0.504707932472229\n",
      "-0.14482781291007996 0.47287729382514954\n",
      "-0.12687696516513824 0.48900339007377625\n",
      "-0.12391192466020584 0.4880266785621643\n",
      "-0.11946869641542435 0.4398445785045624\n",
      "-0.14859451353549957 0.4515717625617981\n",
      "-0.14776840806007385 0.46926605701446533\n",
      "-0.29076245427131653 0.41265881061553955\n",
      "-0.11479546874761581 0.4017811119556427\n",
      "-0.20562557876110077 0.4405595660209656\n",
      "-0.17587292194366455 0.5682826638221741\n",
      "-0.17757375538349152 0.47771981358528137\n",
      "-0.08002414554357529 0.4367806613445282\n",
      "-0.24147315323352814 0.49362754821777344\n",
      "-0.15549612045288086 0.45089906454086304\n",
      "-0.14629223942756653 0.509382963180542\n",
      "-0.13014493882656097 0.46555978059768677\n",
      "-0.1035449206829071 0.45574939250946045\n",
      "-0.16826468706130981 0.4681404232978821\n",
      "-0.09871090203523636 0.4971483647823334\n",
      "-0.11458286643028259 0.45200568437576294\n",
      "-0.16804035007953644 0.4828857183456421\n",
      "-0.15735562145709991 0.4436054527759552\n",
      "-0.11621641367673874 0.4578518271446228\n",
      "-0.0734591856598854 0.4958399534225464\n",
      "-0.17498169839382172 0.46238791942596436\n",
      "-0.151813805103302 0.4680587351322174\n",
      "-0.1384621560573578 0.47524482011795044\n",
      "-0.12444916367530823 0.44676491618156433\n",
      "-0.12349576503038406 0.4405111074447632\n",
      "-0.12513971328735352 1.87566077709198\n",
      "-0.20195288956165314 0.4857090413570404\n",
      "-0.12510265409946442 0.4537241458892822\n",
      "-0.0951482430100441 0.46683698892593384\n",
      "-0.19551891088485718 0.45510196685791016\n",
      "-0.2598256766796112 0.4170663058757782\n",
      "-0.17390204966068268 0.4203660488128662\n",
      "-0.07963618636131287 0.3992425799369812\n",
      "-0.08179520070552826 0.4533257484436035\n",
      "-0.19929569959640503 0.4832574129104614\n",
      "-0.20489034056663513 0.4125809073448181\n",
      "-0.09987645596265793 0.46946650743484497\n",
      "-0.09551835805177689 0.4303191602230072\n",
      "-0.1562868356704712 0.40341275930404663\n",
      "-0.24671757221221924 0.4835060238838196\n",
      "-0.21526485681533813 0.4199439585208893\n",
      "-0.1028798297047615 0.42486098408699036\n",
      "-0.02530686929821968 0.4202936589717865\n",
      "-0.15586797893047333 0.4394378364086151\n",
      "-0.16054591536521912 0.4022200405597687\n",
      "-0.19941122829914093 0.4581797420978546\n",
      "-0.13985952734947205 0.4351818263530731\n",
      "-0.12215138226747513 0.43885889649391174\n",
      "-0.10614408552646637 0.3957228362560272\n",
      "-0.195384681224823 0.41086122393608093\n",
      "-0.2198360711336136 0.43451085686683655\n",
      "-0.11956994980573654 0.4145947992801666\n",
      "-0.13582906126976013 0.4848421812057495\n",
      "-0.07023224234580994 0.4338870048522949\n",
      "-0.20399896800518036 0.44387903809547424\n",
      "-0.06407269835472107 0.40213072299957275\n",
      "-0.08902610838413239 0.4724620580673218\n",
      "-0.14452272653579712 0.4485990107059479\n",
      "-0.21423159539699554 0.4382757246494293\n",
      "-0.15135551989078522 0.43839019536972046\n",
      "-0.13577711582183838 0.4402657449245453\n",
      "-0.05817725881934166 0.4325472414493561\n",
      "-0.1940726786851883 0.45360440015792847\n",
      "-0.16030296683311462 0.4296739399433136\n",
      "-0.2905258238315582 0.4599488079547882\n",
      "-0.17395374178886414 0.4147350788116455\n",
      "-0.15462632477283478 0.38723066449165344\n",
      "-0.08123195916414261 0.4138782024383545\n",
      "-0.1649918556213379 0.39334455132484436\n",
      "-0.05981915816664696 0.4190574586391449\n",
      "-0.1355431079864502 0.421759694814682\n",
      "-0.16613006591796875 0.3575745224952698\n",
      "-0.17655888199806213 0.4251118004322052\n",
      "-0.10783721506595612 0.41189709305763245\n",
      "-0.0951475203037262 0.4602315127849579\n",
      "-0.12564818561077118 0.47463640570640564\n",
      "-0.07697584480047226 0.43795639276504517\n",
      "-0.1503465324640274 0.38279080390930176\n",
      "-0.18968607485294342 0.42054080963134766\n",
      "-0.10655279457569122 0.4700584411621094\n",
      "-0.15666638314723969 0.4051337242126465\n",
      "-0.08266393095254898 0.395797997713089\n",
      "-0.11580090969800949 0.3863868713378906\n",
      "-0.15688030421733856 0.38845187425613403\n",
      "-0.21556410193443298 0.3887428045272827\n",
      "-0.14177826046943665 0.3836651146411896\n",
      "-0.15272806584835052 0.3726993501186371\n",
      "-0.08810441195964813 0.3902832269668579\n",
      "-0.12906593084335327 0.4193870723247528\n",
      "-0.12513937056064606 0.41057318449020386\n",
      "-0.1217050775885582 0.399828165769577\n",
      "-0.19851699471473694 0.3881068825721741\n",
      "-0.14444713294506073 0.4416557252407074\n",
      "-0.1297980546951294 0.39689335227012634\n",
      "-0.13592135906219482 0.4138667583465576\n",
      "-0.18456238508224487 0.5880352854728699\n",
      "-0.07861055433750153 0.4365350008010864\n",
      "-0.12286359071731567 0.38454508781433105\n",
      "-0.05927090346813202 0.3540217876434326\n",
      "-0.16238440573215485 0.4023171365261078\n",
      "-0.19259628653526306 0.3546966314315796\n",
      "-0.16488701105117798 0.38250038027763367\n",
      "-0.12184988707304001 0.37015286087989807\n",
      "-0.14385753870010376 0.39719995856285095\n",
      "-0.16122928261756897 0.38319724798202515\n",
      "-0.2138468325138092 0.456413596868515\n",
      "-0.13896627724170685 0.3930465579032898\n",
      "-0.09969966858625412 0.3954470455646515\n",
      "-0.11499753594398499 0.3819434642791748\n",
      "-0.16178001463413239 0.4151949882507324\n",
      "-0.13609248399734497 0.38448798656463623\n",
      "-0.18253983557224274 0.36962831020355225\n",
      "-0.12635384500026703 0.4009496569633484\n",
      "-0.08930002897977829 0.3957311809062958\n",
      "-0.12050525099039078 0.3690710961818695\n",
      "-0.12582240998744965 0.39531490206718445\n",
      "-0.1260383278131485 0.3835905194282532\n",
      "-0.2016485184431076 0.3967452943325043\n",
      "-0.15959197282791138 0.3880179524421692\n",
      "-0.13878107070922852 0.37590306997299194\n",
      "-0.13203880190849304 0.37882402539253235\n",
      "-0.1478872448205948 0.3869142532348633\n",
      "-0.15756644308567047 0.3934522271156311\n",
      "-0.2290917932987213 0.43751028180122375\n",
      "-0.20167188346385956 0.4022255837917328\n",
      "-0.08621630072593689 0.3815613389015198\n",
      "-0.14921775460243225 0.4199599027633667\n",
      "-0.21555596590042114 0.3794359266757965\n",
      "-0.1912909299135208 0.39111632108688354\n",
      "-0.18378287553787231 0.37706053256988525\n",
      "-0.19352595508098602 0.4295896887779236\n",
      "-0.1374673992395401 0.38605833053588867\n",
      "-0.10933631658554077 0.392194926738739\n",
      "-0.14508651196956635 0.35916373133659363\n",
      "-0.09358289837837219 0.4397376775741577\n",
      "-0.049361877143383026 0.39033129811286926\n",
      "-0.1828298270702362 0.3443601131439209\n",
      "-0.15611843764781952 0.36892762780189514\n",
      "-0.12370798736810684 0.4152161777019501\n",
      "-0.13090822100639343 0.40366506576538086\n",
      "-0.1524939239025116 0.3665318191051483\n",
      "-0.08418712019920349 0.3907308876514435\n",
      "-0.15921658277511597 0.4233814775943756\n",
      "-0.1596483439207077 0.37567034363746643\n",
      "-0.10168664157390594 0.4606538414955139\n",
      "-0.12446609884500504 0.42199188470840454\n",
      "-0.10453282296657562 0.37910449504852295\n"
     ]
    }
   ],
   "source": [
    "\n",
    "epochs=300\n",
    "for epoch in range(epochs):\n",
    "    count=0\n",
    "  \n",
    "    for imgs_r, angles_r, labels, imgs_t, angles_g in train_loader:\n",
    "        count+=1\n",
    "        imgs_r=imgs_r.to(device)\n",
    "        imgs_t=imgs_t.to(device)\n",
    "        angles_r=angles_r.to(device)\n",
    "        angles_g=angles_g.to(device)\n",
    "        l_d=discriminator_step(generator,discriminator,imgs_r,imgs_t,angles_r,angles_g)\n",
    "        if count%5==0:\n",
    "            l_g=generator_step(generator,discriminator,loss_network,imgs_r,imgs_t,angles_r,angles_g)\n",
    "        if count%1000==0:\n",
    "            imgs=[imgs_r]\n",
    "            for h in [-15,-10,-5,0,5,10,15]:\n",
    "                    a=torch.tile(torch.tensor([h/15.,0.]),[32,1])\n",
    "                    a=a.to(device)\n",
    "                    y=generator(imgs_r,a)\n",
    "                    imgs.append(y.detach())\n",
    "            save_images(imgs, \"./debug/{}_{}.png\".format(epoch,count))\n",
    "    print(l_d,l_g)\n",
    "    if epoch%20==0:\n",
    "        torch.save(generator, './generator.pth')\n",
    "        torch.save(discriminator, './discriminator.pth')\n",
    "     \n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([32, 2]) torch.Size([4, 2])\n"
     ]
    }
   ],
   "source": [
    "a=torch.tile(torch.tensor([0.,0.]),[32,1])\n",
    "print(a.size(),angles_r.size())\n",
    "y=generator(imgs_r.to(device),angles_r.to(device))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "tensorflow",
   "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.11.4"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}