DeepLearning/GAN/CycleGAN/CycleGAN.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## CycleGAN\n",
    "![](Md_img/2023-02-23-11-08-28.png)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import argparse\n",
    "import os \n",
    "import numpy as np \n",
    "import math\n",
    "import itertools\n",
    "import datetime\n",
    "import time\n",
    "\n",
    "import torchvision.transforms as transforms\n",
    "from torchvision.utils import save_image,make_grid\n",
    "\n",
    "from torch.utils.data import DataLoader\n",
    "from torchvision import datasets\n",
    "from torch.autograd import Variable\n",
    "\n",
    "from models import *\n",
    "from datasets import *\n",
    "from utils import *\n",
    "\n",
    "import torch.nn as nn \n",
    "import torch.nn.functional as F\n",
    "import torch"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 1. parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Namespace(b1=0.5, b2=0.999, batch_size=1, channels=3, checkpoint_interval=-1, dataset_name='monet2photo', decay_epoch=100, epoch=0, img_height=256, img_width=256, lambda_cyc=10.0, lambda_id=5.0, lr=0.0002, n_cpu=8, n_epochs=200, n_residual_blocks=9, sample_interval=100)\n"
     ]
    }
   ],
   "source": [
    "parser = argparse.ArgumentParser()\n",
    "parser.add_argument(\"--epoch\", type=int, default=0, help=\"epoch to start training from\")\n",
    "parser.add_argument(\"--n_epochs\", type=int, default=200, help=\"number of epochs of training\")\n",
    "parser.add_argument(\"--dataset_name\", type=str, default=\"monet2photo\", help=\"name of the dataset\")\n",
    "parser.add_argument(\"--batch_size\", type=int, default=1, help=\"size of the batches\")\n",
    "parser.add_argument(\"--lr\", type=float, default=0.0002, help=\"adam: learning rate\")\n",
    "parser.add_argument(\"--b1\", type=float, default=0.5, help=\"adam: decay of first order momentum of gradient\")\n",
    "parser.add_argument(\"--b2\", type=float, default=0.999, help=\"adam: decay of first order momentum of gradient\")\n",
    "parser.add_argument(\"--decay_epoch\", type=int, default=100, help=\"epoch from which to start lr decay\")\n",
    "parser.add_argument(\"--n_cpu\", type=int, default=8, help=\"number of cpu threads to use during batch generation\")\n",
    "parser.add_argument(\"--img_height\", type=int, default=256, help=\"size of image height\")\n",
    "parser.add_argument(\"--img_width\", type=int, default=256, help=\"size of image width\")\n",
    "parser.add_argument(\"--channels\", type=int, default=3, help=\"number of image channels\")\n",
    "parser.add_argument(\"--sample_interval\", type=int, default=100, help=\"interval between saving generator outputs\")\n",
    "parser.add_argument(\"--checkpoint_interval\", type=int, default=-1, help=\"interval between saving model checkpoints\")\n",
    "parser.add_argument(\"--n_residual_blocks\", type=int, default=9, help=\"number of residual blocks in generator\")\n",
    "parser.add_argument(\"--lambda_cyc\", type=float, default=10.0, help=\"cycle loss weight\")\n",
    "parser.add_argument(\"--lambda_id\", type=float, default=5.0, help=\"identity loss weight\")\n",
    "opt = parser.parse_args(args=[]) # 在jupyter中运行需要加args=[]\n",
    "print(opt)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Datasets and Model\n",
    "![](Md_img/2023-02-22-21-34-13.png)\n",
    "\n",
    "**CycleGAN具有一个双判别器结构**\n",
    "* G:从X到Y的生成器(G_AB)\n",
    "* F:从Y到X的生成器(G_BA)\n",
    "* $D_x$:判别生成的X是真实X，还是生成的X\n",
    "* $D_y$:判别生成的Y是真实Y，还是生成的Y"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create sample and checkpoint directories\n",
    "os.makedirs(\"log/images/%s\"%opt.dataset_name,exist_ok=True)\n",
    "os.makedirs(\"log/save_models/%s\"%opt.dataset_name,exist_ok=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Losses\n",
    "criterion_GAN = torch.nn.MSELoss() \n",
    "criterion_cycle = torch.nn.L1Loss() \n",
    "criterion_identity = torch.nn.L1Loss()\n",
    "# device\n",
    "cuda = torch.cuda.is_available()\n",
    "input_shape = (opt.channels,opt.img_height,opt.img_width)\n",
    "opt.checkpoint_interval = 10"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "##############################\n",
    "#           RESNET\n",
    "##############################\n",
    "\n",
    "\n",
    "class ResidualBlock(nn.Module):\n",
    "    def __init__(self, in_features):\n",
    "        super(ResidualBlock, self).__init__()\n",
    "\n",
    "        self.block = nn.Sequential(\n",
    "            nn.ReflectionPad2d(1),\n",
    "            nn.Conv2d(in_features, in_features, 3),\n",
    "            nn.InstanceNorm2d(in_features),\n",
    "            nn.ReLU(inplace=True),\n",
    "            nn.ReflectionPad2d(1),\n",
    "            nn.Conv2d(in_features, in_features, 3),\n",
    "            nn.InstanceNorm2d(in_features),\n",
    "        )\n",
    "\n",
    "    def forward(self, x):\n",
    "        return x + self.block(x)\n",
    "\n",
    "\n",
    "class GeneratorResNet(nn.Module):\n",
    "    def __init__(self, input_shape, num_residual_blocks):\n",
    "        super(GeneratorResNet, self).__init__()\n",
    "\n",
    "        channels = input_shape[0]\n",
    "\n",
    "        # Initial convolution block\n",
    "        out_features = 64\n",
    "        model = [\n",
    "            nn.ReflectionPad2d(channels),\n",
    "            nn.Conv2d(channels, out_features, 7),\n",
    "            nn.InstanceNorm2d(out_features),\n",
    "            nn.ReLU(inplace=True),\n",
    "        ]\n",
    "        in_features = out_features\n",
    "\n",
    "        # Downsampling\n",
    "        for _ in range(2):\n",
    "            out_features *= 2\n",
    "            model += [\n",
    "                nn.Conv2d(in_features, out_features, 3, stride=2, padding=1),\n",
    "                nn.InstanceNorm2d(out_features),\n",
    "                nn.ReLU(inplace=True),\n",
    "            ]\n",
    "            in_features = out_features\n",
    "\n",
    "        # Residual blocks\n",
    "        for _ in range(num_residual_blocks):\n",
    "            model += [ResidualBlock(out_features)]\n",
    "\n",
    "        # Upsampling\n",
    "        for _ in range(2):\n",
    "            out_features //= 2\n",
    "            model += [\n",
    "                nn.Upsample(scale_factor=2),\n",
    "                nn.Conv2d(in_features, out_features, 3, stride=1, padding=1),\n",
    "                nn.InstanceNorm2d(out_features),\n",
    "                nn.ReLU(inplace=True),\n",
    "            ]\n",
    "            in_features = out_features\n",
    "\n",
    "        # Output layer\n",
    "        model += [nn.ReflectionPad2d(channels), nn.Conv2d(out_features, channels, 7), nn.Tanh()]\n",
    "\n",
    "        self.model = nn.Sequential(*model)\n",
    "\n",
    "    def forward(self, x):\n",
    "        return self.model(x)\n",
    "\n",
    "\n",
    "##############################\n",
    "#        Discriminator\n",
    "##############################\n",
    "\n",
    "\n",
    "class Discriminator(nn.Module):\n",
    "    def __init__(self, input_shape):\n",
    "        super(Discriminator, self).__init__()\n",
    "\n",
    "        channels, height, width = input_shape\n",
    "\n",
    "        # Calculate output shape of image discriminator (PatchGAN)\n",
    "        self.output_shape = (1, height // 2 ** 4, width // 2 ** 4)\n",
    "\n",
    "        def discriminator_block(in_filters, out_filters, normalize=True):\n",
    "            \"\"\"Returns downsampling layers of each discriminator block\"\"\"\n",
    "            layers = [nn.Conv2d(in_filters, out_filters, 4, stride=2, padding=1)]\n",
    "            if normalize:\n",
    "                layers.append(nn.InstanceNorm2d(out_filters))\n",
    "            layers.append(nn.LeakyReLU(0.2, inplace=True))\n",
    "            return layers\n",
    "\n",
    "        self.model = nn.Sequential(\n",
    "            *discriminator_block(channels, 64, normalize=False),\n",
    "            *discriminator_block(64, 128),\n",
    "            *discriminator_block(128, 256),\n",
    "            *discriminator_block(256, 512),\n",
    "            nn.ZeroPad2d((1, 0, 1, 0)),\n",
    "            nn.Conv2d(512, 1, 4, padding=1)\n",
    "        )\n",
    "\n",
    "    def forward(self, img):\n",
    "        return self.model(img)\n",
    "\n",
    "def weights_init_normal(m):\n",
    "    classname = m.__class__.__name__\n",
    "    if classname.find(\"Conv\") != -1:\n",
    "        torch.nn.init.normal_(m.weight.data,0.0,0.02)\n",
    "        if hasattr(m,\"bias\") and m.bias is not None:\n",
    "            torch.nn.init.constant_(m.bias.data,0.0)\n",
    "    elif classname.find(\"BatchNorm2d\") != -1:\n",
    "        torch.nn.init.normal_(m.weight.data,1.0,0.02)\n",
    "        torch.nn.init.constant_(m.bias.data,0.0)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(GeneratorResNet(\n",
       "   (model): Sequential(\n",
       "     (0): ReflectionPad2d((3, 3, 3, 3))\n",
       "     (1): Conv2d(3, 64, kernel_size=(7, 7), stride=(1, 1))\n",
       "     (2): InstanceNorm2d(64, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "     (3): ReLU(inplace=True)\n",
       "     (4): Conv2d(64, 128, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "     (5): InstanceNorm2d(128, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "     (6): ReLU(inplace=True)\n",
       "     (7): Conv2d(128, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))\n",
       "     (8): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "     (9): ReLU(inplace=True)\n",
       "     (10): ResidualBlock(\n",
       "       (block): Sequential(\n",
       "         (0): ReflectionPad2d((1, 1, 1, 1))\n",
       "         (1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))\n",
       "         (2): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "         (3): ReLU(inplace=True)\n",
       "         (4): ReflectionPad2d((1, 1, 1, 1))\n",
       "         (5): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))\n",
       "         (6): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "       )\n",
       "     )\n",
       "     (11): ResidualBlock(\n",
       "       (block): Sequential(\n",
       "         (0): ReflectionPad2d((1, 1, 1, 1))\n",
       "         (1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))\n",
       "         (2): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "         (3): ReLU(inplace=True)\n",
       "         (4): ReflectionPad2d((1, 1, 1, 1))\n",
       "         (5): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))\n",
       "         (6): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "       )\n",
       "     )\n",
       "     (12): ResidualBlock(\n",
       "       (block): Sequential(\n",
       "         (0): ReflectionPad2d((1, 1, 1, 1))\n",
       "         (1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))\n",
       "         (2): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "         (3): ReLU(inplace=True)\n",
       "         (4): ReflectionPad2d((1, 1, 1, 1))\n",
       "         (5): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))\n",
       "         (6): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "       )\n",
       "     )\n",
       "     (13): ResidualBlock(\n",
       "       (block): Sequential(\n",
       "         (0): ReflectionPad2d((1, 1, 1, 1))\n",
       "         (1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))\n",
       "         (2): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "         (3): ReLU(inplace=True)\n",
       "         (4): ReflectionPad2d((1, 1, 1, 1))\n",
       "         (5): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))\n",
       "         (6): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "       )\n",
       "     )\n",
       "     (14): ResidualBlock(\n",
       "       (block): Sequential(\n",
       "         (0): ReflectionPad2d((1, 1, 1, 1))\n",
       "         (1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))\n",
       "         (2): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "         (3): ReLU(inplace=True)\n",
       "         (4): ReflectionPad2d((1, 1, 1, 1))\n",
       "         (5): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))\n",
       "         (6): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "       )\n",
       "     )\n",
       "     (15): ResidualBlock(\n",
       "       (block): Sequential(\n",
       "         (0): ReflectionPad2d((1, 1, 1, 1))\n",
       "         (1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))\n",
       "         (2): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "         (3): ReLU(inplace=True)\n",
       "         (4): ReflectionPad2d((1, 1, 1, 1))\n",
       "         (5): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))\n",
       "         (6): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "       )\n",
       "     )\n",
       "     (16): ResidualBlock(\n",
       "       (block): Sequential(\n",
       "         (0): ReflectionPad2d((1, 1, 1, 1))\n",
       "         (1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))\n",
       "         (2): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "         (3): ReLU(inplace=True)\n",
       "         (4): ReflectionPad2d((1, 1, 1, 1))\n",
       "         (5): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))\n",
       "         (6): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "       )\n",
       "     )\n",
       "     (17): ResidualBlock(\n",
       "       (block): Sequential(\n",
       "         (0): ReflectionPad2d((1, 1, 1, 1))\n",
       "         (1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))\n",
       "         (2): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "         (3): ReLU(inplace=True)\n",
       "         (4): ReflectionPad2d((1, 1, 1, 1))\n",
       "         (5): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))\n",
       "         (6): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "       )\n",
       "     )\n",
       "     (18): ResidualBlock(\n",
       "       (block): Sequential(\n",
       "         (0): ReflectionPad2d((1, 1, 1, 1))\n",
       "         (1): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))\n",
       "         (2): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "         (3): ReLU(inplace=True)\n",
       "         (4): ReflectionPad2d((1, 1, 1, 1))\n",
       "         (5): Conv2d(256, 256, kernel_size=(3, 3), stride=(1, 1))\n",
       "         (6): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "       )\n",
       "     )\n",
       "     (19): Upsample(scale_factor=2.0, mode=nearest)\n",
       "     (20): Conv2d(256, 128, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "     (21): InstanceNorm2d(128, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "     (22): ReLU(inplace=True)\n",
       "     (23): Upsample(scale_factor=2.0, mode=nearest)\n",
       "     (24): Conv2d(128, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))\n",
       "     (25): InstanceNorm2d(64, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "     (26): ReLU(inplace=True)\n",
       "     (27): ReflectionPad2d((3, 3, 3, 3))\n",
       "     (28): Conv2d(64, 3, kernel_size=(7, 7), stride=(1, 1))\n",
       "     (29): Tanh()\n",
       "   )\n",
       " ),\n",
       " Discriminator(\n",
       "   (model): Sequential(\n",
       "     (0): Conv2d(3, 64, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "     (1): LeakyReLU(negative_slope=0.2, inplace=True)\n",
       "     (2): Conv2d(64, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "     (3): InstanceNorm2d(128, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "     (4): LeakyReLU(negative_slope=0.2, inplace=True)\n",
       "     (5): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "     (6): InstanceNorm2d(256, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "     (7): LeakyReLU(negative_slope=0.2, inplace=True)\n",
       "     (8): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "     (9): InstanceNorm2d(512, eps=1e-05, momentum=0.1, affine=False, track_running_stats=False)\n",
       "     (10): LeakyReLU(negative_slope=0.2, inplace=True)\n",
       "     (11): ZeroPad2d((1, 0, 1, 0))\n",
       "     (12): Conv2d(512, 1, kernel_size=(4, 4), stride=(1, 1), padding=(1, 1))\n",
       "   )\n",
       " ))"
      ]
     },
     "execution_count": 13,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# initialize generator and discriminator\n",
    "G_AB = GeneratorResNet(input_shape,opt.n_residual_blocks)\n",
    "G_BA = GeneratorResNet(input_shape,opt.n_residual_blocks)\n",
    "D_A = Discriminator(input_shape)\n",
    "D_B = Discriminator(input_shape)\n",
    "\n",
    "if cuda:\n",
    "    G_AB = G_AB.cuda()\n",
    "    G_BA = G_BA.cuda()\n",
    "    D_A = D_A.cuda()\n",
    "    D_B = D_B.cuda()\n",
    "    criterion_GAN.cuda()\n",
    "    criterion_cycle.cuda()\n",
    "    criterion_identity.cuda()\n",
    "\n",
    "if opt.epoch != 0:\n",
    "    # Load pretrained models\n",
    "    G_AB.load_state_dict(torch.load(\"log/saved_models/%s/G_AB_%d.pth\" % (opt.dataset_name, opt.epoch)))\n",
    "    G_BA.load_state_dict(torch.load(\"log/saved_models/%s/G_BA_%d.pth\" % (opt.dataset_name, opt.epoch)))\n",
    "    D_A.load_state_dict(torch.load(\"log/saved_models/%s/D_A_%d.pth\" % (opt.dataset_name, opt.epoch)))\n",
    "    D_B.load_state_dict(torch.load(\"log/saved_models/%s/D_B_%d.pth\" % (opt.dataset_name, opt.epoch)))\n",
    "else:\n",
    "    # Initialize weights\n",
    "    G_AB.apply(weights_init_normal)\n",
    "    G_BA.apply(weights_init_normal)\n",
    "    D_A.apply(weights_init_normal)\n",
    "    D_B.apply(weights_init_normal)\n",
    "\n",
    "G_AB,D_A"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "# optimizers\n",
    "optimizer_G = torch.optim.Adam(\n",
    "    itertools.chain(G_AB.parameters(),G_BA.parameters()),\n",
    "    lr = opt.lr,betas=(opt.b1,opt.b2)\n",
    ")\n",
    "optimizer_D_A = torch.optim.Adam(\n",
    "    D_A.parameters(),\n",
    "    lr = opt.lr,betas=(opt.b1,opt.b2)\n",
    ")\n",
    "optimizer_D_B = torch.optim.Adam(\n",
    "    D_B.parameters(),\n",
    "    betas=(opt.b1,opt.b2)\n",
    ")\n",
    "\n",
    "class LambdaLR:\n",
    "    def __init__(self, n_epochs, offset, decay_start_epoch):\n",
    "        assert (n_epochs - decay_start_epoch) > 0, \"Decay must start before the training session ends!\"\n",
    "        self.n_epochs = n_epochs\n",
    "        self.offset = offset\n",
    "        self.decay_start_epoch = decay_start_epoch\n",
    "\n",
    "    def step(self, epoch):\n",
    "        return 1.0 - max(0, epoch + self.offset - self.decay_start_epoch) / (self.n_epochs - self.decay_start_epoch)\n",
    "\n",
    "\n",
    "# Learning rate update schedulers\n",
    "lr_scheduler_G = torch.optim.lr_scheduler.LambdaLR(\n",
    "    optimizer_G, lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step\n",
    ")\n",
    "lr_scheduler_D_A = torch.optim.lr_scheduler.LambdaLR(\n",
    "    optimizer_D_A, lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step\n",
    ")\n",
    "lr_scheduler_D_B = torch.optim.lr_scheduler.LambdaLR(\n",
    "    optimizer_D_B, lr_lambda=LambdaLR(opt.n_epochs, opt.epoch, opt.decay_epoch).step\n",
    ")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "class ReplayBuffer:\n",
    "    \"\"\"\n",
    "        样本收集再采样过程\n",
    "        push_and_pop:生成器生成的fake图片还要经过另一个生成器,生成cycle图片\n",
    "        所以通过buffer函数寄存fake图片,用于判别器的更新?\n",
    "    \"\"\"\n",
    "    def __init__(self, max_size=50):\n",
    "        assert max_size > 0, \"Empty buffer or trying to create a black hole. Be careful.\"\n",
    "        self.max_size = max_size\n",
    "        self.data = []\n",
    "\n",
    "    def push_and_pop(self, data):\n",
    "        to_return = []\n",
    "        for element in data.data:\n",
    "            element = torch.unsqueeze(element, 0)\n",
    "            if len(self.data) < self.max_size:\n",
    "                self.data.append(element)\n",
    "                to_return.append(element)\n",
    "            else:\n",
    "                if random.uniform(0, 1) > 0.5:\n",
    "                    i = random.randint(0, self.max_size - 1)\n",
    "                    to_return.append(self.data[i].clone())\n",
    "                    self.data[i] = element\n",
    "                else:\n",
    "                    to_return.append(element)\n",
    "        return Variable(torch.cat(to_return))\n",
    "\n",
    "# Buffers of previously generated samples\n",
    "fake_A_buffer = ReplayBuffer()\n",
    "fake_B_buffer = ReplayBuffer()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "def to_rgb(image):\n",
    "    rgb_image = Image.new(\"RGB\", image.size)\n",
    "    rgb_image.paste(image)\n",
    "    return rgb_image\n",
    "\n",
    "\n",
    "class ImageDataset(Dataset):\n",
    "    def __init__(self, root, transforms_=None, unaligned=False, mode=\"train\"):\n",
    "        self.transform = transforms.Compose(transforms_)\n",
    "        self.unaligned = unaligned\n",
    "\n",
    "        self.files_A = sorted(glob.glob(os.path.join(root, \"%s/A\" % mode) + \"/*.*\"))\n",
    "        self.files_B = sorted(glob.glob(os.path.join(root, \"%s/B\" % mode) + \"/*.*\"))\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        image_A = Image.open(self.files_A[index % len(self.files_A)])\n",
    "\n",
    "        if self.unaligned:\n",
    "            image_B = Image.open(self.files_B[random.randint(0, len(self.files_B) - 1)])\n",
    "        else:\n",
    "            image_B = Image.open(self.files_B[index % len(self.files_B)])\n",
    "\n",
    "        # Convert grayscale images to rgb\n",
    "        if image_A.mode != \"RGB\":\n",
    "            image_A = to_rgb(image_A)\n",
    "        if image_B.mode != \"RGB\":\n",
    "            image_B = to_rgb(image_B)\n",
    "\n",
    "        item_A = self.transform(image_A)\n",
    "        item_B = self.transform(image_B)\n",
    "        return {\"A\": item_A, \"B\": item_B}\n",
    "\n",
    "    def __len__(self):\n",
    "        return max(len(self.files_A), len(self.files_B))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/tmp/ipykernel_2509787/2981494205.py:3: DeprecationWarning: BICUBIC is deprecated and will be removed in Pillow 10 (2023-07-01). Use Resampling.BICUBIC instead.\n",
      "  transforms.Resize(int(opt.img_height * 1.12), Image.BICUBIC),\n",
      "/home/liufeng/anaconda3/envs/pytorch/lib/python3.8/site-packages/torchvision/transforms/transforms.py:332: UserWarning: Argument 'interpolation' of type int is deprecated since 0.13 and will be removed in 0.15. Please use InterpolationMode enum.\n",
      "  warnings.warn(\n"
     ]
    }
   ],
   "source": [
    "# Image transformations\n",
    "transforms_ = [\n",
    "    transforms.Resize(int(opt.img_height * 1.12), Image.BICUBIC),\n",
    "    transforms.RandomCrop((opt.img_height, opt.img_width)),\n",
    "    transforms.RandomHorizontalFlip(),\n",
    "    transforms.ToTensor(),\n",
    "    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),\n",
    "]\n",
    "\n",
    "# Training data loader\n",
    "dataloader = DataLoader(\n",
    "    ImageDataset(\"../../data/%s\" % opt.dataset_name, transforms_=transforms_, unaligned=True),\n",
    "    batch_size=opt.batch_size,\n",
    "    shuffle=True,\n",
    "    num_workers=opt.n_cpu,\n",
    ")\n",
    "\n",
    "# Test data loader\n",
    "val_dataloader = DataLoader(\n",
    "    ImageDataset(\"../../data/%s\" % opt.dataset_name, transforms_=transforms_, unaligned=True, mode=\"test\"),\n",
    "    batch_size=5,\n",
    "    shuffle=True,\n",
    "    num_workers=1,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "Tensor = torch.cuda.FloatTensor if cuda else torch.Tensor\n",
    "def sample_images(batches_done):\n",
    "    \"\"\"Saves a generated sample from the test set\"\"\"\n",
    "    imgs = next(iter(val_dataloader))\n",
    "    G_AB.eval()\n",
    "    G_BA.eval()\n",
    "    real_A = Variable(imgs[\"A\"].type(Tensor))\n",
    "    fake_B = G_AB(real_A)\n",
    "    real_B = Variable(imgs[\"B\"].type(Tensor))\n",
    "    fake_A = G_BA(real_B)\n",
    "    # Arange images along x-axis\n",
    "    real_A = make_grid(real_A, nrow=5, normalize=True)\n",
    "    real_B = make_grid(real_B, nrow=5, normalize=True)\n",
    "    fake_A = make_grid(fake_A, nrow=5, normalize=True)\n",
    "    fake_B = make_grid(fake_B, nrow=5, normalize=True)\n",
    "    # Arange images along y-axis\n",
    "    image_grid = torch.cat((real_A, fake_B, real_B, fake_A), 1)\n",
    "    save_image(image_grid, \"log/images/%s/%s.png\" % (opt.dataset_name, batches_done), normalize=False)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[Epoch 0/200] [Batch 11/6287] [D loss: 0.098257] [G loss: 1.786578, adv: 0.471492, cycle: 0.091518, identity: 0.079981] ETA: 3 days, 15:45:25.917277"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[0;32mIn [24], line 28\u001b[0m\n\u001b[1;32m     25\u001b[0m loss_identity \u001b[39m=\u001b[39m (loss_id_A\u001b[39m+\u001b[39mloss_id_B)\u001b[39m/\u001b[39m\u001b[39m2\u001b[39m\n\u001b[1;32m     27\u001b[0m \u001b[39m# GAN loss\u001b[39;00m\n\u001b[0;32m---> 28\u001b[0m fake_B \u001b[39m=\u001b[39m G_AB(real_A)\n\u001b[1;32m     29\u001b[0m loss_GAN_AB \u001b[39m=\u001b[39m criterion_GAN(D_B(fake_B), valid)\n\u001b[1;32m     30\u001b[0m fake_A \u001b[39m=\u001b[39m G_BA(real_B)\n",
      "File \u001b[0;32m~/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/nn/modules/module.py:1130\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1126\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1127\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1128\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1131\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1132\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[0;32m/media/liufeng/a0b205ec-bfb3-473f-a6f0-0680c5da64ba/project/MachineLearning_DeepLearning/04_GAN/PyTorch-GAN/Example_LF/CycleGAN/models.py:87\u001b[0m, in \u001b[0;36mGeneratorResNet.forward\u001b[0;34m(self, x)\u001b[0m\n\u001b[1;32m     86\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mforward\u001b[39m(\u001b[39mself\u001b[39m, x):\n\u001b[0;32m---> 87\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mmodel(x)\n",
      "File \u001b[0;32m~/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/nn/modules/module.py:1130\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1126\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1127\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1128\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1131\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1132\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[0;32m~/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/nn/modules/container.py:139\u001b[0m, in \u001b[0;36mSequential.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m    137\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mforward\u001b[39m(\u001b[39mself\u001b[39m, \u001b[39minput\u001b[39m):\n\u001b[1;32m    138\u001b[0m     \u001b[39mfor\u001b[39;00m module \u001b[39min\u001b[39;00m \u001b[39mself\u001b[39m:\n\u001b[0;32m--> 139\u001b[0m         \u001b[39minput\u001b[39m \u001b[39m=\u001b[39m module(\u001b[39minput\u001b[39;49m)\n\u001b[1;32m    140\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39minput\u001b[39m\n",
      "File \u001b[0;32m~/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/nn/modules/module.py:1130\u001b[0m, in \u001b[0;36mModule._call_impl\u001b[0;34m(self, *input, **kwargs)\u001b[0m\n\u001b[1;32m   1126\u001b[0m \u001b[39m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[1;32m   1127\u001b[0m \u001b[39m# this function, and just call forward.\u001b[39;00m\n\u001b[1;32m   1128\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39mnot\u001b[39;00m (\u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_backward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_hooks \u001b[39mor\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_forward_pre_hooks \u001b[39mor\u001b[39;00m _global_backward_hooks\n\u001b[1;32m   1129\u001b[0m         \u001b[39mor\u001b[39;00m _global_forward_hooks \u001b[39mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[0;32m-> 1130\u001b[0m     \u001b[39mreturn\u001b[39;00m forward_call(\u001b[39m*\u001b[39;49m\u001b[39minput\u001b[39;49m, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m   1131\u001b[0m \u001b[39m# Do not call functions when jit is used\u001b[39;00m\n\u001b[1;32m   1132\u001b[0m full_backward_hooks, non_full_backward_hooks \u001b[39m=\u001b[39m [], []\n",
      "File \u001b[0;32m~/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/nn/modules/instancenorm.py:72\u001b[0m, in \u001b[0;36m_InstanceNorm.forward\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m     69\u001b[0m \u001b[39mif\u001b[39;00m \u001b[39minput\u001b[39m\u001b[39m.\u001b[39mdim() \u001b[39m==\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_get_no_batch_dim():\n\u001b[1;32m     70\u001b[0m     \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39m_handle_no_batch_input(\u001b[39minput\u001b[39m)\n\u001b[0;32m---> 72\u001b[0m \u001b[39mreturn\u001b[39;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49m_apply_instance_norm(\u001b[39minput\u001b[39;49m)\n",
      "File \u001b[0;32m~/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/nn/modules/instancenorm.py:32\u001b[0m, in \u001b[0;36m_InstanceNorm._apply_instance_norm\u001b[0;34m(self, input)\u001b[0m\n\u001b[1;32m     31\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39m_apply_instance_norm\u001b[39m(\u001b[39mself\u001b[39m, \u001b[39minput\u001b[39m):\n\u001b[0;32m---> 32\u001b[0m     \u001b[39mreturn\u001b[39;00m F\u001b[39m.\u001b[39;49minstance_norm(\n\u001b[1;32m     33\u001b[0m         \u001b[39minput\u001b[39;49m, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mrunning_mean, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mrunning_var, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mweight, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mbias,\n\u001b[1;32m     34\u001b[0m         \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mtraining \u001b[39mor\u001b[39;49;00m \u001b[39mnot\u001b[39;49;00m \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mtrack_running_stats, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49mmomentum, \u001b[39mself\u001b[39;49m\u001b[39m.\u001b[39;49meps)\n",
      "File \u001b[0;32m~/anaconda3/envs/pytorch/lib/python3.8/site-packages/torch/nn/functional.py:2483\u001b[0m, in \u001b[0;36minstance_norm\u001b[0;34m(input, running_mean, running_var, weight, bias, use_input_stats, momentum, eps)\u001b[0m\n\u001b[1;32m   2481\u001b[0m \u001b[39mif\u001b[39;00m use_input_stats:\n\u001b[1;32m   2482\u001b[0m     _verify_spatial_size(\u001b[39minput\u001b[39m\u001b[39m.\u001b[39msize())\n\u001b[0;32m-> 2483\u001b[0m \u001b[39mreturn\u001b[39;00m torch\u001b[39m.\u001b[39;49minstance_norm(\n\u001b[1;32m   2484\u001b[0m     \u001b[39minput\u001b[39;49m, weight, bias, running_mean, running_var, use_input_stats, momentum, eps, torch\u001b[39m.\u001b[39;49mbackends\u001b[39m.\u001b[39;49mcudnn\u001b[39m.\u001b[39;49menabled\n\u001b[1;32m   2485\u001b[0m )\n",
      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "# ----------\n",
    "#  Training\n",
    "# ----------\n",
    "prev_time = time.time()\n",
    "for epoch in range(opt.epoch,opt.n_epochs):\n",
    "    for i,batch in enumerate(dataloader):\n",
    "        # set model input \n",
    "        real_A = Variable(batch[\"A\"].type(Tensor))\n",
    "        real_B = Variable(batch[\"B\"].type(Tensor))\n",
    "        \n",
    "        # Adversarial ground truths\n",
    "        valid = Variable(Tensor(np.ones((real_A.size(0),*D_A.output_shape))),requires_grad=False)\n",
    "        fake =  Variable(Tensor(np.zeros((real_A.size(0),*D_A.output_shape))),requires_grad=False)\n",
    "        \n",
    "        # ------------------\n",
    "        #  Train Generators\n",
    "        # ------------------\n",
    "        G_AB.train()\n",
    "        G_BA.train()\n",
    "        optimizer_G.zero_grad()\n",
    "        \n",
    "        # identity loss \n",
    "        loss_id_A = criterion_identity(G_BA(real_A),real_A)\n",
    "        loss_id_B = criterion_identity(G_AB(real_B),real_B)\n",
    "        loss_identity = (loss_id_A+loss_id_B)/2\n",
    "        \n",
    "        # GAN loss\n",
    "        fake_B = G_AB(real_A)\n",
    "        loss_GAN_AB = criterion_GAN(D_B(fake_B), valid)\n",
    "        fake_A = G_BA(real_B)\n",
    "        loss_GAN_BA = criterion_GAN(D_A(fake_A), valid)\n",
    "\n",
    "        loss_GAN = (loss_GAN_AB + loss_GAN_BA) / 2\n",
    "        \n",
    "        # Cycle loss\n",
    "        recov_A = G_BA(fake_B)\n",
    "        loss_cycle_A = criterion_cycle(recov_A, real_A)\n",
    "        recov_B = G_AB(fake_A)\n",
    "        loss_cycle_B = criterion_cycle(recov_B, real_B)\n",
    "\n",
    "        loss_cycle = (loss_cycle_A + loss_cycle_B) / 2\n",
    "        # Total loss\n",
    "        loss_G = loss_GAN + opt.lambda_cyc * loss_cycle + opt.lambda_id * loss_identity\n",
    "\n",
    "        loss_G.backward()\n",
    "        optimizer_G.step()\n",
    "        \n",
    "        # -----------------------\n",
    "        #  Train Discriminator A\n",
    "        # -----------------------\n",
    "\n",
    "        optimizer_D_A.zero_grad()\n",
    "\n",
    "        # Real loss\n",
    "        loss_real = criterion_GAN(D_A(real_A), valid)\n",
    "        # Fake loss (on batch of previously generated samples)\n",
    "        fake_A_ = fake_A_buffer.push_and_pop(fake_A)\n",
    "        loss_fake = criterion_GAN(D_A(fake_A_.detach()), fake)\n",
    "        # Total loss\n",
    "        loss_D_A = (loss_real + loss_fake) / 2\n",
    "\n",
    "        loss_D_A.backward()\n",
    "        optimizer_D_A.step()\n",
    "\n",
    "        # -----------------------\n",
    "        #  Train Discriminator B\n",
    "        # -----------------------\n",
    "\n",
    "        optimizer_D_B.zero_grad()\n",
    "\n",
    "        # Real loss\n",
    "        loss_real = criterion_GAN(D_B(real_B), valid)\n",
    "        # Fake loss (on batch of previously generated samples)\n",
    "        fake_B_ = fake_B_buffer.push_and_pop(fake_B)\n",
    "        loss_fake = criterion_GAN(D_B(fake_B_.detach()), fake)\n",
    "        # Total loss\n",
    "        loss_D_B = (loss_real + loss_fake) / 2\n",
    "\n",
    "        loss_D_B.backward()\n",
    "        optimizer_D_B.step()\n",
    "\n",
    "        loss_D = (loss_D_A + loss_D_B) / 2\n",
    "\n",
    "        # --------------\n",
    "        #  Log Progress\n",
    "        # --------------\n",
    "\n",
    "        # Determine approximate time left\n",
    "        batches_done = epoch * len(dataloader) + i\n",
    "        batches_left = opt.n_epochs * len(dataloader) - batches_done\n",
    "        time_left = datetime.timedelta(seconds=batches_left * (time.time() - prev_time))\n",
    "        prev_time = time.time()\n",
    "        \n",
    "        # Print log\n",
    "        sys.stdout.write(\n",
    "            \"\\r[Epoch %d/%d] [Batch %d/%d] [D loss: %f] [G loss: %f, adv: %f, cycle: %f, identity: %f] ETA: %s\"\n",
    "            % (\n",
    "                epoch,\n",
    "                opt.n_epochs,\n",
    "                i,\n",
    "                len(dataloader),\n",
    "                loss_D.item(),\n",
    "                loss_G.item(),\n",
    "                loss_GAN.item(),\n",
    "                loss_cycle.item(),\n",
    "                loss_identity.item(),\n",
    "                time_left,\n",
    "            )\n",
    "        )\n",
    "\n",
    "        # If at sample interval save image\n",
    "        if batches_done % opt.sample_interval == 0:\n",
    "            sample_images(batches_done)\n",
    "\n",
    "    # Update learning rates\n",
    "    lr_scheduler_G.step()\n",
    "    lr_scheduler_D_A.step()\n",
    "    lr_scheduler_D_B.step()\n",
    "\n",
    "    if opt.checkpoint_interval != -1 and epoch % opt.checkpoint_interval == 0:\n",
    "        # Save model checkpoints\n",
    "        torch.save(G_AB.state_dict(), \"log/save_models/%s/G_AB_%d.pth\" % (opt.dataset_name, epoch))\n",
    "        torch.save(G_BA.state_dict(), \"log/save_models/%s/G_BA_%d.pth\" % (opt.dataset_name, epoch))\n",
    "        torch.save(D_A.state_dict(), \"log/save_models/%s/D_A_%d.pth\" % (opt.dataset_name, epoch))\n",
    "        torch.save(D_B.state_dict(), \"log/save_models/%s/D_B_%d.pth\" % (opt.dataset_name, epoch))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.save(G_AB.state_dict(), \"log/save_models/%s/G_AB_%d.pth\" % (opt.dataset_name, epoch))\n",
    "torch.save(G_BA.state_dict(), \"log/save_models/%s/G_BA_%d.pth\" % (opt.dataset_name, epoch))\n",
    "torch.save(D_A.state_dict(), \"log/save_models/%s/D_A_%d.pth\" % (opt.dataset_name, epoch))\n",
    "torch.save(D_B.state_dict(), \"log/save_models/%s/D_B_%d.pth\" % (opt.dataset_name, epoch))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_list = os.listdir(\"./log/images/monet2photo/\")\n",
    "for file in file_list:\n",
    "    if not int(file.replace(\".png\",\"\"))%1000==0:\n",
    "        file_path = os.path.join(\"./log/images/monet2photo/\",file)\n",
    "        os.remove(file_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "pytorch",
   "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.8.13 (default, Mar 28 2022, 11:38:47) \n[GCC 7.5.0]"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "971ee1da7255ecde3216791eeee73f2ee8d2956ba6b6da13cdfb2a534ec0c35f"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}