gandalfs_main_nb.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "5_X83Fd9u9LB"
   },
   "source": [
    "# Prepair GANDALF Jr."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "id": "dHhrBRO3FNRV"
   },
   "outputs": [],
   "source": [
    "import os, cv2, glob, datetime, sys, glob\n",
    "from PIL import Image\n",
    "import torch, torchvision\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torchvision import datasets, transforms\n",
    "from torchvision.utils import save_image\n",
    "from torch.utils.data import DataLoader\n",
    "from torch.utils.data.dataset import Dataset\n",
    "import numpy as np\n",
    "from PIL import Image\n",
    "from matplotlib.pyplot import imshow, imsave\n",
    "%matplotlib inline\n",
    "\n",
    "# import gandalf tools\n",
    "from gandalfs_tools import *"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "step, starting_epoch, starting_step, inum, n_noise, n_critic = 1, 0, 0, 100, 100, 1\n",
    "ida, sessionid = \"\", \"1\"\n",
    "\n",
    "IMAGE_DIM, normval = (128, 128, 3), (0.5, 0.5, 0.5)\n",
    "DEVICE = torch.device(\"cuda:0\" if torch.cuda.is_available() else \"cpu\")\n",
    "\n",
    "# set local paths\n",
    "SAVE_DIR_ROOT = 'D:\\\\GANDALFSWELT\\\\out\\\\'\n",
    "mempath_root = 'D:\\\\GANDALFSWELT\\\\checkpoints\\\\'\n",
    "\n",
    "INPUT_IMAGES =\"D:\\\\GANDALFSWELT\\\\input\\\\back_in\\\\\"\n",
    "RESIZED_IMAGES = \"D:\\\\GANDALFSWELT\\\\input\\\\resized\\\\\"\n",
    "\n",
    "READY_IMAGES = \"D:\\\\GANDALFSWELT\\\\input\\\\ready\\\\\"\n",
    "READYPLUS_IMAGES =  \"D:\\\\GANDALFSWELT\\\\input\\\\readyplus\\\\\"\n",
    "ORG_IMAGES = \"D:\\\\GANDALFSWELT\\\\input\\\\readyorg\\\\\"\n",
    "GANDALFS_IMAGES = \"D:\\\\GANDALFSWELT\\\\input\\\\backinall\\\\\"\n",
    "\n",
    "transform = transforms.Compose([transforms.Resize((IMAGE_DIM[0],IMAGE_DIM[1])),transforms.ToTensor(),\n",
    "                                transforms.Normalize(mean=normval,std=normval)])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "D:\\GANDALFSWELT\\checkpoints\\1889.pt\n",
      "\t\t\t\t...current rnum: 1889\n"
     ]
    }
   ],
   "source": [
    "# get last memory path\n",
    "mempath = mempath_root + str(max([int(x.replace('.pt','')) for x in os.listdir(mempath_root)])) + '.pt'\n",
    "print(mempath)\n",
    "\n",
    "# set run number\n",
    "rnum = max([int(x.replace('.pt','')) for x in os.listdir(mempath_root)])\n",
    "print(\"\\t\\t\\t\\t...current rnum:\",rnum)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "j-91kkSpLzjX"
   },
   "source": [
    "# Create GANDALF Jr."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 470
    },
    "id": "3nBwFpLkIHDW",
    "outputId": "c295862e-c088-46bc-93ab-42e220a2d603"
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Generator(\n",
       "  (fc): Sequential(\n",
       "    (0): Linear(in_features=100, out_features=32768, bias=True)\n",
       "    (1): ReLU()\n",
       "  )\n",
       "  (conv): Sequential(\n",
       "    (0): Conv2d(512, 512, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False)\n",
       "    (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (2): ReLU()\n",
       "    (3): ConvTranspose2d(512, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "    (4): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (5): ReLU()\n",
       "    (6): ConvTranspose2d(256, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "    (7): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (8): ReLU()\n",
       "    (9): ConvTranspose2d(128, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "    (10): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (11): ReLU()\n",
       "    (12): ConvTranspose2d(128, 3, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "    (13): Tanh()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "D = Discriminator(in_channel=IMAGE_DIM[-1]).to(DEVICE)\n",
    "G = Generator(out_channel=IMAGE_DIM[-1]).to(DEVICE)\n",
    "criterion = nn.BCELoss()\n",
    "D_opt = torch.optim.Adam(D.parameters(), lr=0.001, betas=(0.5, 0.999)) # lr = 0.001\n",
    "G_opt = torch.optim.Adam(G.parameters(), lr=0.001, betas=(0.5, 0.999))\n",
    "G.eval()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "IYIiylnWIFHr"
   },
   "source": [
    "# Reload GANDALF Jr. from memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Discriminator(\n",
       "  (conv): Sequential(\n",
       "    (0): Conv2d(3, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "    (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (2): LeakyReLU(negative_slope=0.2)\n",
       "    (3): Conv2d(512, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "    (4): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (5): LeakyReLU(negative_slope=0.2)\n",
       "    (6): Conv2d(256, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "    (7): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (8): LeakyReLU(negative_slope=0.2)\n",
       "    (9): Conv2d(128, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), bias=False)\n",
       "    (10): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "    (11): LeakyReLU(negative_slope=0.2)\n",
       "    (12): AdaptiveAvgPool2d(output_size=1)\n",
       "  )\n",
       "  (fc): Sequential(\n",
       "    (0): Linear(in_features=128, out_features=1, bias=True)\n",
       "    (1): Sigmoid()\n",
       "  )\n",
       ")"
      ]
     },
     "execution_count": 15,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "checkpoint = torch.load(mempath)\n",
    "D.load_state_dict(checkpoint['D'])\n",
    "G.load_state_dict(checkpoint['G'])\n",
    "D_opt.load_state_dict(checkpoint['d_optim'])\n",
    "G_opt.load_state_dict(checkpoint['g_optim'])        \n",
    "G.train()\n",
    "D.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "ty6yoXL98nyW"
   },
   "source": [
    "# Input Images / Preprocessing\n",
    "\n",
    "#### batch size max.: 89"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "batch_size = 89"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [],
   "source": [
    "# preprocessing step 1: resize\n",
    "preprocess_step1([INPUT_IMAGES+x for x in os.listdir(INPUT_IMAGES)], RESIZED_IMAGES)\n",
    "\n",
    "# preprocessing step 2: convert and save\n",
    "preprocess_step2(READY_IMAGES, RESIZED_IMAGES)\n",
    "preprocess_step2(GANDALFS_IMAGES, RESIZED_IMAGES)\n",
    "preprocess_step2(READYPLUS_IMAGES, RESIZED_IMAGES)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 30,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Num Original Pics:\t 2484 \t Num Batches: 27.91 \n",
      "Num Generated Pics:\t 2541 \t Num Batches: 28.55 \n",
      "\n",
      "Num Total Input Pics:\t 4551 \t Num Batches: 51.13 \n",
      "Num Input Pics Plus:\t 5024 \t Num Batches: 56.45\n"
     ]
    }
   ],
   "source": [
    "def pic_num(path):\n",
    "    return len(os.listdir(path))\n",
    "\n",
    "A_p, O_p, G_p, Ap_p = pic_num(READY_IMAGES), pic_num(ORG_IMAGES), pic_num(GANDALFS_IMAGES), pic_num(READYPLUS_IMAGES)\n",
    "Ap_ev, A_ev, O_ev = int(Ap_p/batch_size), int(A_p/batch_size), int(O_p/batch_size)\n",
    "\n",
    "print('Num Original Pics:\\t',O_p, '\\t Num Batches:', round((O_p/batch_size),2),\n",
    "      '\\nNum Generated Pics:\\t',G_p,'\\t Num Batches:', round((G_p/batch_size),2),\n",
    "      '\\n\\nNum Total Input Pics:\\t',A_p,'\\t Num Batches:', round((A_p/batch_size),2),\n",
    "      '\\nNum Input Pics Plus:\\t',Ap_p,'\\t Num Batches:', round((Ap_p/batch_size),2))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yaMou2kFFNR2"
   },
   "source": [
    "# GANDALF JR starts creating"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [],
   "source": [
    "# settings for output pics & folders:\n",
    "\n",
    "# NOTES:\n",
    "# scale_ev 0.5 gut für video\n",
    "\n",
    "update_savedir = True\n",
    "reset_stepnum = False\n",
    "\n",
    "mixes = 1\n",
    "scale_epoch_num = 1  \n",
    "scale_ev = 0.5"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set No. 1: All Plus Pics\n",
    "**Original and Gandalf Input Pics plus small selection of horizontally flipped G pics**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "num epochs per mix: 5 \t pic ev: 28\n",
      "made folder D:\\GANDALFSWELT\\out\\kater\\1898_scev50_Ap\\\n"
     ]
    }
   ],
   "source": [
    "max_epoch = int(scale_ev*10*scale_epoch_num) \n",
    "ev = int(Ap_ev*scale_ev) \n",
    "print(\"num epochs per mix:\", max_epoch, \"\\t pic ev:\", ev)\n",
    "\n",
    "if update_savedir:\n",
    "    SAVE_DIR = SAVE_DIR_ROOT+str(rnum+1)+\"_scev\"+ str(int(scale_ev*100))+\"_Ap\" +\"\\\\\"\n",
    "    makesavepath(SAVE_DIR)\n",
    "    \n",
    "if reset_stepnum:\n",
    "    step = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = READYPLUS_IMAGES\n",
    "dataset = IMAGES(data_path=data_path,transform=transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "G2b5kteBFNSE",
    "outputId": "11c701e5-e839-4317-8103-69f96344284b",
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Mix 1 of 1\n",
      "\n",
      "\tEpoch: 1/5, D Loss: 0.0018, G Loss: 7.7924\n",
      "\t\t\t\t\t saving D:\\GANDALFSWELT\\out\\kater\\1898_scev50_Ap\\1898_00000.jpg\n",
      "\n",
      "\tEpoch: 1/5, D Loss: 0.0014, G Loss: 6.9597\n",
      "\t\t\t\t\t saving D:\\GANDALFSWELT\\out\\kater\\1898_scev50_Ap\\1898_00001.jpg\n",
      "\n",
      "\tEpoch: 2/5, D Loss: 0.0021, G Loss: 7.6317\n",
      "\t\t\t\t\t saving D:\\GANDALFSWELT\\out\\kater\\1898_scev50_Ap\\1898_00002.jpg\n",
      "\n",
      "\tEpoch: 2/5, D Loss: 0.0029, G Loss: 6.4516\n",
      "\t\t\t\t\t saving D:\\GANDALFSWELT\\out\\kater\\1898_scev50_Ap\\1898_00003.jpg\n"
     ]
    }
   ],
   "source": [
    "inum = 100000\n",
    "rnum += 1\n",
    "\n",
    "for mix in range(mixes):\n",
    "    \n",
    "    print(\"Mix\",mix+1,\"of\",mixes)\n",
    "    data_loader = DataLoader(dataset=dataset,batch_size=batch_size,shuffle=True,drop_last=True,num_workers=0)\n",
    "    D_labels = torch.ones([batch_size, 1]).to(DEVICE)                       # Discriminator Label to real\n",
    "    D_fakes = torch.zeros([batch_size, 1]).to(DEVICE)                       # Discriminator Label to fake\n",
    "    \n",
    "    for epoch in range(max_epoch):\n",
    "        for idx, images in enumerate(data_loader):                          # D  \n",
    "            x = images.to(DEVICE)\n",
    "            x_outputs = D(x)\n",
    "            D_x_loss = criterion(x_outputs, D_labels)\n",
    "            z = torch.randn(batch_size, n_noise).to(DEVICE)                 # G -> D\n",
    "            z_outputs = D(G(z))\n",
    "            D_z_loss = criterion(z_outputs, D_fakes)\n",
    "            D_loss = D_x_loss + D_z_loss\n",
    "            D.zero_grad()                                                   # D loss + opt\n",
    "            D_loss.backward()\n",
    "            D_opt.step()\n",
    "            if step % n_critic == 0:                                        # G loss + opt\n",
    "                z = torch.randn(batch_size, n_noise).to(DEVICE)\n",
    "                z_outputs = D(G(z))\n",
    "                G_loss = criterion(z_outputs, D_labels)\n",
    "                D.zero_grad()\n",
    "                G.zero_grad()\n",
    "                G_loss.backward()\n",
    "                G_opt.step()\n",
    "            if step % ev == 0:                                              # get pic\n",
    "                name = SAVE_DIR +ida+str(rnum)+ \"_\"+str(inum)[1:]+ '.jpg'\n",
    "                print('\\n\\tEpoch: {}/{}, D Loss: {:.4f}, G Loss: {:.4f}'.format(epoch+1,max_epoch,D_loss.item(),G_loss.item()))\n",
    "                G.eval()\n",
    "                img = generate_image(G, 1, n_noise)\n",
    "                imsave(name, img[0],dpi=600)\n",
    "                print(\"\\t\"*5,\"saving\",name)\n",
    "                inum +=1\n",
    "                G.train()\n",
    "            step += 1       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mempath = mempath_root + str(rnum) + \".pt\"\n",
    "torch.save({'D':D.state_dict(),'G':G.state_dict(),'d_optim':D_opt.state_dict(),'g_optim':G_opt.state_dict()},mempath) \n",
    "mempath"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set No. 2: All Pics\n",
    "**Original and Gandalf Input Pics**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_epoch = int(scale_ev*10*scale_epoch_num) \n",
    "ev = int(Ap_ev*scale_ev) \n",
    "print(\"num epochs per mix:\", max_epoch, \"\\t pic ev:\", ev)\n",
    "\n",
    "if update_savedir:\n",
    "    SAVE_DIR = SAVE_DIR_ROOT+str(rnum+1)+\"_scev\"+ str(int(scale_ev*100))+\"_Ap\" +\"\\\\\"\n",
    "    makesavepath(SAVE_DIR)\n",
    "    \n",
    "if reset_stepnum:\n",
    "    step = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = READY_IMAGES\n",
    "dataset = IMAGES(data_path=data_path,transform=transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "G2b5kteBFNSE",
    "outputId": "11c701e5-e839-4317-8103-69f96344284b",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "inum = 100000\n",
    "rnum += 1\n",
    "\n",
    "for mix in range(mixes):\n",
    "    \n",
    "    print(\"Mix\",mix+1,\"of\",mixes)\n",
    "    data_loader = DataLoader(dataset=dataset,batch_size=batch_size,shuffle=True,drop_last=True,num_workers=0)\n",
    "    D_labels = torch.ones([batch_size, 1]).to(DEVICE)                       # Discriminator Label to real\n",
    "    D_fakes = torch.zeros([batch_size, 1]).to(DEVICE)                       # Discriminator Label to fake\n",
    "    \n",
    "    for epoch in range(max_epoch):\n",
    "        for idx, images in enumerate(data_loader):                          # D  \n",
    "            x = images.to(DEVICE)\n",
    "            x_outputs = D(x)\n",
    "            D_x_loss = criterion(x_outputs, D_labels)\n",
    "            z = torch.randn(batch_size, n_noise).to(DEVICE)                 # G -> D\n",
    "            z_outputs = D(G(z))\n",
    "            D_z_loss = criterion(z_outputs, D_fakes)\n",
    "            D_loss = D_x_loss + D_z_loss\n",
    "            D.zero_grad()                                                   # D loss + opt\n",
    "            D_loss.backward()\n",
    "            D_opt.step()\n",
    "            if step % n_critic == 0:                                        # G loss + opt\n",
    "                z = torch.randn(batch_size, n_noise).to(DEVICE)\n",
    "                z_outputs = D(G(z))\n",
    "                G_loss = criterion(z_outputs, D_labels)\n",
    "                D.zero_grad()\n",
    "                G.zero_grad()\n",
    "                G_loss.backward()\n",
    "                G_opt.step()\n",
    "            if step % ev == 0:                                              # get pic\n",
    "                name = SAVE_DIR +ida+str(rnum)+ \"_\"+str(inum)[1:]+ '.jpg'\n",
    "                print('\\n\\tEpoch: {}/{}, D Loss: {:.4f}, G Loss: {:.4f}'.format(epoch+1,max_epoch,D_loss.item(),G_loss.item()))\n",
    "                G.eval()\n",
    "                img = generate_image(G, 1, n_noise)\n",
    "                imsave(name, img[0],dpi=600)\n",
    "                print(\"\\t\"*5,\"saving\",name)\n",
    "                inum +=1\n",
    "                G.train()\n",
    "            step += 1       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mempath = mempath_root + str(rnum) + \".pt\"\n",
    "torch.save({'D':D.state_dict(),'G':G.state_dict(),'d_optim':D_opt.state_dict(),'g_optim':G_opt.state_dict()},mempath) \n",
    "mempath"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set No. 3) Original Pis \n",
    "**Only Original Input Pics**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "max_epoch = int(scale_ev*10*scale_epoch_num) * int((Ap_p/batch_size)/(O_p/batch_size)) # multiply with factorization O/A\n",
    "ev = int(Ap_ev*scale_ev) \n",
    "print(\"num epochs per mix:\", max_epoch, \"\\t pic ev:\", ev)\n",
    "\n",
    "if update_savedir:\n",
    "    SAVE_DIR = SAVE_DIR_ROOT+str(rnum+1)+\"_scev\"+ str(int(scale_ev*100))+\"_Ap\" +\"\\\\\"\n",
    "    makesavepath(SAVE_DIR)\n",
    "    \n",
    "if reset_stepnum:\n",
    "    step = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data_path = ORG_IMAGES\n",
    "dataset = IMAGES(data_path=data_path,transform=transform)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/",
     "height": 1000
    },
    "id": "G2b5kteBFNSE",
    "outputId": "11c701e5-e839-4317-8103-69f96344284b",
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "inum = 100000\n",
    "rnum += 1\n",
    "\n",
    "for mix in range(mixes):\n",
    "    \n",
    "    print(\"Mix\",mix+1,\"of\",mixes)\n",
    "    data_loader = DataLoader(dataset=dataset,batch_size=batch_size,shuffle=True,drop_last=True,num_workers=0)\n",
    "    D_labels = torch.ones([batch_size, 1]).to(DEVICE)                       # Discriminator Label to real\n",
    "    D_fakes = torch.zeros([batch_size, 1]).to(DEVICE)                       # Discriminator Label to fake\n",
    "    \n",
    "    for epoch in range(max_epoch):\n",
    "        for idx, images in enumerate(data_loader):                          # D  \n",
    "            x = images.to(DEVICE)\n",
    "            x_outputs = D(x)\n",
    "            D_x_loss = criterion(x_outputs, D_labels)\n",
    "            z = torch.randn(batch_size, n_noise).to(DEVICE)                 # G -> D\n",
    "            z_outputs = D(G(z))\n",
    "            D_z_loss = criterion(z_outputs, D_fakes)\n",
    "            D_loss = D_x_loss + D_z_loss\n",
    "            D.zero_grad()                                                   # D loss + opt\n",
    "            D_loss.backward()\n",
    "            D_opt.step()\n",
    "            if step % n_critic == 0:                                        # G loss + opt\n",
    "                z = torch.randn(batch_size, n_noise).to(DEVICE)\n",
    "                z_outputs = D(G(z))\n",
    "                G_loss = criterion(z_outputs, D_labels)\n",
    "                D.zero_grad()\n",
    "                G.zero_grad()\n",
    "                G_loss.backward()\n",
    "                G_opt.step()\n",
    "            if step % ev == 0:                                              # get pic\n",
    "                name = SAVE_DIR +ida+str(rnum)+ \"_\"+str(inum)[1:]+ '.jpg'\n",
    "                print('\\n\\tEpoch: {}/{}, D Loss: {:.4f}, G Loss: {:.4f}'.format(epoch+1,max_epoch,D_loss.item(),G_loss.item()))\n",
    "                G.eval()\n",
    "                img = generate_image(G, 1, n_noise)\n",
    "                imsave(name, img[0],dpi=600)\n",
    "                print(\"\\t\"*5,\"saving\",name)\n",
    "                inum +=1\n",
    "                G.train()\n",
    "            step += 1       "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mempath = mempath_root + str(rnum) + \".pt\"\n",
    "torch.save({'D':D.state_dict(),'G':G.state_dict(),'d_optim':D_opt.state_dict(),'g_optim':G_opt.state_dict()},mempath)\n",
    "mempath"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "collapsed_sections": [],
   "name": "Pytorch DCGAN.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "walter",
   "language": "python",
   "name": "walter"
  },
  "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}