{"metadata":{"kernelspec":{"language":"python","display_name":"Python 3","name":"python3"},"language_info":{"pygments_lexer":"ipython3","nbconvert_exporter":"python","version":"3.6.4","file_extension":".py","codemirror_mode":{"name":"ipython","version":3},"name":"python","mimetype":"text/x-python"}},"nbformat_minor":4,"nbformat":4,"cells":[{"cell_type":"code","source":"!pip install torchinfo","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:52:31.530769Z","iopub.execute_input":"2023-06-16T07:52:31.531249Z","iopub.status.idle":"2023-06-16T07:52:47.156345Z","shell.execute_reply.started":"2023-06-16T07:52:31.531215Z","shell.execute_reply":"2023-06-16T07:52:47.154888Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!mkdir images\n!mkdir models","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:52:47.159284Z","iopub.execute_input":"2023-06-16T07:52:47.15981Z","iopub.status.idle":"2023-06-16T07:52:49.351302Z","shell.execute_reply.started":"2023-06-16T07:52:47.159758Z","shell.execute_reply":"2023-06-16T07:52:49.34962Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nfrom datetime import datetime\nimport numpy as np\nimport pandas as pd\nfrom PIL import Image\nimport torch\nfrom torch import nn\nfrom torch.utils.data import Dataset, DataLoader\nfrom tqdm.notebook import tqdm\n\nfrom torchinfo import summary\n\nimport matplotlib.pyplot as plt\nimport seaborn as sns\nfrom mpl_toolkits.axes_grid1 import ImageGrid\nfrom IPython.display import clear_output, display","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:52:49.353336Z","iopub.execute_input":"2023-06-16T07:52:49.354865Z","iopub.status.idle":"2023-06-16T07:52:53.806157Z","shell.execute_reply.started":"2023-06-16T07:52:49.354788Z","shell.execute_reply":"2023-06-16T07:52:53.804852Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"PATH = \"../input/gan-getting-started\"\nIMG_SIZE = (256, 256)","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:52:53.808778Z","iopub.execute_input":"2023-06-16T07:52:53.809401Z","iopub.status.idle":"2023-06-16T07:52:53.815157Z","shell.execute_reply.started":"2023-06-16T07:52:53.809367Z","shell.execute_reply":"2023-06-16T07:52:53.813797Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**preprocessing images**","metadata":{}},{"cell_type":"code","source":"class ImageDataset(Dataset):\n    def __init__(self, dir_path, transform=None):\n        super().__init__()\n        \n        self.transform = transform\n        self.images = []\n        \n        for fn in tqdm(os.listdir(dir_path)):\n            img = Image.open(os.path.join(dir_path, fn))\n            self.images.append(img)\n            \n    def __len__(self):\n        return len(self.images)\n    \n    def __getitem__(self, idx):\n        image = self.images[idx]\n        \n        if self.transform is not None:\n            image = self.transform(image=np.array(image))[\"image\"]\n            \n        return image","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:52:53.81691Z","iopub.execute_input":"2023-06-16T07:52:53.817363Z","iopub.status.idle":"2023-06-16T07:52:53.833412Z","shell.execute_reply.started":"2023-06-16T07:52:53.81732Z","shell.execute_reply":"2023-06-16T07:52:53.832335Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**preprocessing Monet images**","metadata":{}},{"cell_type":"code","source":"class MonetDataset(Dataset):\n    def __init__(self, datasets, augs=[]):\n        super().__init__()\n        \n        self.datasets = datasets\n        self.augs = augs\n        self.init_augs = [dataset.transform for dataset in datasets]\n            \n    def __len__(self):\n        return min([len(dataset) for dataset in self.datasets])\n    \n    def set_train_mode(self):\n        for dataset, aug in zip(self.datasets, self.augs):\n            dataset.transform = aug\n            \n    def set_test_mode(self):\n        for dataset, aug in zip(self.datasets, self.init_augs):\n            dataset.transform = aug\n    \n    def __getitem__(self, idx):\n        _return = []\n        for dataset in self.datasets:\n            idx = np.random.choice(len(dataset))\n            _return.append(dataset[idx])\n            \n        return tuple(_return)","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:52:53.834637Z","iopub.execute_input":"2023-06-16T07:52:53.835447Z","iopub.status.idle":"2023-06-16T07:52:53.850007Z","shell.execute_reply.started":"2023-06-16T07:52:53.835412Z","shell.execute_reply":"2023-06-16T07:52:53.84891Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Augmentation**","metadata":{}},{"cell_type":"code","source":"import albumentations as A\nfrom albumentations.pytorch import ToTensorV2\nfrom albumentations.augmentations.crops.transforms import RandomCrop\n\ntransform = A.Compose([\n    A.Resize(*IMG_SIZE),\n    A.Normalize(mean=[.5]*3, std=[.5]*3),\n    #A.HorizontalFlip(p=0.7),\n    ToTensorV2(),\n])\n\naug = A.Compose([\n   #randomly shifts the color channels of the image.\n    A.RGBShift(r_shift_limit=10, g_shift_limit=10, b_shift_limit=10, p=.5),\n    #n horizontally flips the image with a probability of 0.5\n    A.HorizontalFlip(p=.5),\n    # randomly applies shifting, scaling, and rotation transformations to the image\n    A.ShiftScaleRotate(rotate_limit=10),\n    # randomly adjusts the brightness and contrast of the image\n    A.RandomBrightnessContrast(brightness_limit=.2, contrast_limit=.2),\n    \n    transform\n])\n","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:52:53.851183Z","iopub.execute_input":"2023-06-16T07:52:53.852433Z","iopub.status.idle":"2023-06-16T07:52:56.075446Z","shell.execute_reply.started":"2023-06-16T07:52:53.852389Z","shell.execute_reply":"2023-06-16T07:52:56.074151Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"path_images = os.path.join(PATH, \"photo_jpg\")\npath_monets = os.path.join(PATH, \"monet_jpg\")\ndataset_photos = ImageDataset(path_images, transform)\ndataset_monet = ImageDataset(path_monets, transform)\ndataset = MonetDataset([dataset_photos, dataset_monet])","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:52:56.077243Z","iopub.execute_input":"2023-06-16T07:52:56.078762Z","iopub.status.idle":"2023-06-16T07:54:04.736626Z","shell.execute_reply.started":"2023-06-16T07:52:56.078719Z","shell.execute_reply":"2023-06-16T07:54:04.735438Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from torchvision.transforms import ToPILImage\n\nto_pil = ToPILImage()\nto_img = lambda x: to_pil(x*.5 + .5)","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:04.738611Z","iopub.execute_input":"2023-06-16T07:54:04.7398Z","iopub.status.idle":"2023-06-16T07:54:04.747452Z","shell.execute_reply.started":"2023-06-16T07:54:04.739752Z","shell.execute_reply":"2023-06-16T07:54:04.74573Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Plot Augmented images**","metadata":{}},{"cell_type":"code","source":"def plot_images_grid(imgs, figsize=(10, 10)):\n    n = int(np.ceil(np.sqrt(len(imgs))))\n    fig = plt.figure(figsize=figsize)\n    grid = ImageGrid(fig, 111,\n                     nrows_ncols=(n, n),\n                     axes_pad=0.1,\n                     )\n    for img, ax in zip(imgs, grid):\n        ax.imshow(img)\n    \n    for ax in grid:\n        ax.set_axis_off()\n        \n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:04.752706Z","iopub.execute_input":"2023-06-16T07:54:04.753089Z","iopub.status.idle":"2023-06-16T07:54:04.764763Z","shell.execute_reply.started":"2023-06-16T07:54:04.753059Z","shell.execute_reply":"2023-06-16T07:54:04.763724Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"    dataset_monet.transform = aug\n    print('Augmented')\n    plot_images_grid([to_img(dataset_monet[i]) for i in range(5)], figsize=(12, 12))\n\n    dataset_monet.transform = transform\n    print('Original')\n    plot_images_grid([to_img(dataset_monet[i]) for i in range(5)], figsize=(12, 12))","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:04.766566Z","iopub.execute_input":"2023-06-16T07:54:04.76691Z","iopub.status.idle":"2023-06-16T07:54:07.662476Z","shell.execute_reply.started":"2023-06-16T07:54:04.766883Z","shell.execute_reply":"2023-06-16T07:54:07.661065Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def display_pair(img, monet, gen=None, figsize=(10, 10)):\n    n_rows = 3 if gen is not None else 2\n    fig = plt.figure(figsize=figsize)\n    grid = ImageGrid(fig, 111,\n                     nrows_ncols=(1, n_rows),\n                     axes_pad=0.1,\n                     )\n    grid_iter = iter(grid)\n    \n    ax_img = next(grid_iter)\n    ax_img.imshow(img)\n    ax_img.set_axis_off()\n    ax_img.set_title('Image')\n    \n    ax_monet = next(grid_iter)\n    ax_monet.imshow(monet)\n    ax_monet.set_axis_off()\n    ax_monet.set_title('Monet')\n    \n    if n_rows > 2:\n        ax_gen = next(grid_iter)\n        ax_gen.imshow(gen)\n        ax_gen.set_axis_off()\n        ax_gen.set_title('Generated')\n    \n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:07.664253Z","iopub.execute_input":"2023-06-16T07:54:07.665222Z","iopub.status.idle":"2023-06-16T07:54:07.67586Z","shell.execute_reply.started":"2023-06-16T07:54:07.665178Z","shell.execute_reply":"2023-06-16T07:54:07.67475Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"photo, monet = dataset[0]\ndisplay_pair(to_img(photo), to_img(monet))","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:07.677619Z","iopub.execute_input":"2023-06-16T07:54:07.67839Z","iopub.status.idle":"2023-06-16T07:54:08.260469Z","shell.execute_reply.started":"2023-06-16T07:54:07.678351Z","shell.execute_reply":"2023-06-16T07:54:08.259114Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"BATCH_SIZE = 1\nWORKERS = 2","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:08.262178Z","iopub.execute_input":"2023-06-16T07:54:08.262616Z","iopub.status.idle":"2023-06-16T07:54:08.268091Z","shell.execute_reply.started":"2023-06-16T07:54:08.262579Z","shell.execute_reply":"2023-06-16T07:54:08.266763Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"dataloader = DataLoader(dataset, batch_size=BATCH_SIZE, num_workers=WORKERS)","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:08.269969Z","iopub.execute_input":"2023-06-16T07:54:08.270374Z","iopub.status.idle":"2023-06-16T07:54:08.282699Z","shell.execute_reply.started":"2023-06-16T07:54:08.270339Z","shell.execute_reply":"2023-06-16T07:54:08.281359Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**MODEL**","metadata":{}},{"cell_type":"code","source":"def Upsample(in_ch, out_ch, r=2, dropout_ratio=0.5):\n    modules = list()\n    \n    if dropout_ratio > 0:\n        modules.append(nn.Dropout(dropout_ratio))\n        \n    modules += [\n        nn.PixelShuffle(2),\n        \n        nn.ReflectionPad2d(1),\n        nn.Conv2d(in_ch//r**2, out_ch, 3, 1),\n        nn.InstanceNorm2d(out_ch),\n        nn.LeakyReLU(negative_slope=0.2, inplace=True),#negative_slope=0.2\n        \n        nn.ReflectionPad2d(1),\n        nn.Conv2d(out_ch, out_ch, 3, 1),\n        nn.InstanceNorm2d(out_ch),\n        nn.LeakyReLU(negative_slope=0.2, inplace=True),#negative_slope=0.2\n    ]\n    return nn.Sequential(*modules)","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:08.284541Z","iopub.execute_input":"2023-06-16T07:54:08.285019Z","iopub.status.idle":"2023-06-16T07:54:08.298118Z","shell.execute_reply.started":"2023-06-16T07:54:08.284979Z","shell.execute_reply":"2023-06-16T07:54:08.297008Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def Downsample(in_ch, out_ch, kernel=3, stride=2, padding=1):\n    return nn.Sequential(\n        nn.ReflectionPad2d(padding),\n        nn.Conv2d(in_ch, out_ch, kernel, stride, bias=True),\n        nn.InstanceNorm2d(out_ch),\n        nn.LeakyReLU(negative_slope=0.2, inplace=True)#negative_slope=0.2\n    )","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:08.299577Z","iopub.execute_input":"2023-06-16T07:54:08.300045Z","iopub.status.idle":"2023-06-16T07:54:08.313379Z","shell.execute_reply.started":"2023-06-16T07:54:08.300004Z","shell.execute_reply":"2023-06-16T07:54:08.312094Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Resblock(nn.Module):\n    def ConvBlock(self, in_ch):\n        return nn.Sequential(\n            nn.ReflectionPad2d(1),\n            nn.Conv2d(in_ch, in_ch, 3, 1, bias=True),\n            nn.InstanceNorm2d(in_ch),\n        )\n    \n    def ActivationFunction(self):\n        return nn.LeakyReLU(negative_slope=0.2, inplace=True)#negative_slope=0.2\n    \n    def __init__(self, in_ch, dropout_ratio=0.5):\n        super().__init__()\n        \n        regularization = nn.Dropout(dropout_ratio) \\\n                             if dropout_ratio > 0 else nn.Identity()\n        \n        self.activation = self.ActivationFunction()\n        \n        self.res = nn.Sequential(\n            self.ConvBlock(in_ch),\n            self.activation,\n            regularization,\n            self.ConvBlock(in_ch),\n        )\n\n    def forward(self, x):\n        return self.activation(x + self.res(x))\n","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:08.315355Z","iopub.execute_input":"2023-06-16T07:54:08.315757Z","iopub.status.idle":"2023-06-16T07:54:08.329863Z","shell.execute_reply.started":"2023-06-16T07:54:08.315726Z","shell.execute_reply":"2023-06-16T07:54:08.328766Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Generator**","metadata":{}},{"cell_type":"code","source":"class Generator(nn.Module):\n    def __init__(self, num_res_blocks=6):\n        super().__init__()\n        modules = list()\n        \n        modules += [\n            Downsample(3, 64, kernel=7, stride=1, padding=3),\n            Downsample(64, 128, kernel=3, stride=2, padding=1),\n            Downsample(128, 256, kernel=3, stride=2, padding=1),\n        ]\n        \n        modules += [\n           Resblock(256) for _ in range(num_res_blocks)\n        ]\n        \n        modules += [\n            Upsample(256, 64),\n            Upsample(64, 16),\n        ]\n        \n        modules += [\n            nn.ReflectionPad2d(3),\n            nn.Conv2d(16, 3, 7, 1),\n            nn.Tanh()\n        ]\n\n        self.gen = nn.Sequential(*modules)\n\n    def forward(self, x):\n        return self.gen(x)","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:08.332041Z","iopub.execute_input":"2023-06-16T07:54:08.332692Z","iopub.status.idle":"2023-06-16T07:54:08.350598Z","shell.execute_reply.started":"2023-06-16T07:54:08.332648Z","shell.execute_reply":"2023-06-16T07:54:08.349535Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Discriminator**","metadata":{}},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    def __init__(self, num_layers=4):\n        super().__init__()\n        modules = [\n            Downsample(3, 32, 4, 2),\n        ]\n        \n        for i in range(num_layers-1):\n            in_ch = 32*2**i#32*2***i\n            out_ch = 32*2**(i+1)#32*2***(i+1)\n            modules.append(\n                Downsample(in_ch, out_ch, 4, 1),\n            )\n            \n        modules += [\n            Downsample(out_ch, out_ch, 4, 2),\n            Downsample(out_ch, out_ch, 4, 2),\n            nn.Conv2d(out_ch, 1, kernel_size=4, stride=1, padding=1),\n        ]   \n        \n        self.disc = nn.Sequential(*modules)\n\n    def forward(self, x):\n        return self.disc(x)","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:08.352063Z","iopub.execute_input":"2023-06-16T07:54:08.35276Z","iopub.status.idle":"2023-06-16T07:54:08.362779Z","shell.execute_reply.started":"2023-06-16T07:54:08.352718Z","shell.execute_reply":"2023-06-16T07:54:08.361999Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def init_weights(m):\n    classname = m.__class__.__name__\n    if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):\n        nn.init.normal_(m.weight.data, 0.0, 0.02)\n        if hasattr(m, 'bias') and m.bias is not None:\n            nn.init.constant_(m.bias.data, 0.0)\n    elif classname.find('BatchNorm2d') != -1:\n        nn.init.normal_(m.weight.data, 1.0, 0.02)\n        nn.init.constant_(m.bias.data, 0.0)","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:08.364437Z","iopub.execute_input":"2023-06-16T07:54:08.364876Z","iopub.status.idle":"2023-06-16T07:54:08.381793Z","shell.execute_reply.started":"2023-06-16T07:54:08.364839Z","shell.execute_reply":"2023-06-16T07:54:08.380545Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class sample_fake:\n    def __init__(self, max_imgs=50):\n        self.max_imgs = max_imgs\n        self.imgs = []\n        self.current = 0\n        \n    def __call__(self, img):\n        if self.current < self.max_imgs:\n            return_img = img\n            self.imgs.append(img)\n            self.current += 1\n        else:\n            if np.random.ranf() > 0.5:\n                idx = np.random.choice(self.max_imgs)\n                return_img = self.imgs[idx]\n                self.imgs[idx] = img\n            else:\n                return_img = img\n        return return_img","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:08.383436Z","iopub.execute_input":"2023-06-16T07:54:08.383932Z","iopub.status.idle":"2023-06-16T07:54:08.398145Z","shell.execute_reply.started":"2023-06-16T07:54:08.383898Z","shell.execute_reply":"2023-06-16T07:54:08.396749Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def get_real(size):\n    if LABEL_SMOOTHING == \"SMOOTH\":\n        return torch.FloatTensor(*size).uniform_(.7, 1.2)\n    elif LABEL_SMOOTHING == \"ABS_SMOOTH\":\n        smooth = torch.FloatTensor(1).uniform_(.7, 1.2)\n        return torch.ones(size)*smooth\n    elif LABEL_SMOOTHING is None:\n        return torch.ones(size)\n    \ndef get_fake(size):\n    if LABEL_SMOOTHING == \"SMOOTH\":\n        return torch.FloatTensor(*size).uniform_(0, .3)\n    elif LABEL_SMOOTHING == \"ABS_SMOOTH\":\n        smooth = torch.FloatTensor(1).uniform_(0, .3)\n        return torch.ones(size)*smooth\n    elif LABEL_SMOOTHING is None:\n        return torch.zeros(size)","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:08.399359Z","iopub.execute_input":"2023-06-16T07:54:08.400238Z","iopub.status.idle":"2023-06-16T07:54:08.413521Z","shell.execute_reply.started":"2023-06-16T07:54:08.400201Z","shell.execute_reply":"2023-06-16T07:54:08.412294Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def random_swap(x, y, p=.5):\n    if np.random.rand() >= p:\n        return x, y\n    else:\n        return y, x","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:08.415796Z","iopub.execute_input":"2023-06-16T07:54:08.416265Z","iopub.status.idle":"2023-06-16T07:54:08.43336Z","shell.execute_reply.started":"2023-06-16T07:54:08.416223Z","shell.execute_reply":"2023-06-16T07:54:08.432013Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def plot_loss(loss_g, loss_d, figsize=(10, 10)):\n    sns.set_style(\"darkgrid\")\n    fig, ax = plt.subplots(figsize=figsize)\n    sns.lineplot(x=range(len(loss_g)), y=loss_g, ax=ax, label=\"Generator\", color=\"blue\")\n    sns.lineplot(x=range(len(loss_d)), y=loss_d, ax=ax, label=\"Discriminator\", color=\"red\")\n    ax.set_title(\"Generator and Discriminator loss\")\n    ax.set_xlabel(\"epoch\")\n    ax.set_ylabel(\"loss\")\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:08.435188Z","iopub.execute_input":"2023-06-16T07:54:08.43567Z","iopub.status.idle":"2023-06-16T07:54:08.448326Z","shell.execute_reply.started":"2023-06-16T07:54:08.435626Z","shell.execute_reply":"2023-06-16T07:54:08.446672Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Training CycleGAN**","metadata":{}},{"cell_type":"code","source":"def save_checkpoint(state, save_path):\n    torch.save(state, save_path)","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:08.4497Z","iopub.execute_input":"2023-06-16T07:54:08.450072Z","iopub.status.idle":"2023-06-16T07:54:08.460146Z","shell.execute_reply.started":"2023-06-16T07:54:08.450042Z","shell.execute_reply":"2023-06-16T07:54:08.458985Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def load_checkpoint(path):\n    print(\"Loading model from: %s\" % path)\n    return torch.load(path)","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:08.461709Z","iopub.execute_input":"2023-06-16T07:54:08.462073Z","iopub.status.idle":"2023-06-16T07:54:08.472801Z","shell.execute_reply.started":"2023-06-16T07:54:08.462043Z","shell.execute_reply":"2023-06-16T07:54:08.471888Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from itertools import chain\n\nLR = 2e-4\nLAMBDA = 10\nEPOCHS = 10\nIDENTITY_COEF = .5\nBETA1 = .5\n#   None - just 0/1 \n#   \"SMOOTH\" - every value in iteration in range 0-0.3/0.7-1.2\n#   \"ABS_SMOOTH\" - same as smooth by single value per iteration\n# (e.g. only .8 for real)\nLABEL_SMOOTHING = \"ABS_SMOOTH\"\nDISCRIMINATOR_LABEL_SWAP_P = 0\n\n# Model loading\nSAVE_MODEL_DIR = None\nLOAD_MODEL_PATH = None\n\n# Log\nUPDATE_EVERY_N = 10\nN_SAMPLES = 10","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:08.479406Z","iopub.execute_input":"2023-06-16T07:54:08.479772Z","iopub.status.idle":"2023-06-16T07:54:08.487352Z","shell.execute_reply.started":"2023-06-16T07:54:08.479742Z","shell.execute_reply":"2023-06-16T07:54:08.486294Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = \"cuda\" if torch.cuda.is_available() else \"cpu\"","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:08.489338Z","iopub.execute_input":"2023-06-16T07:54:08.489705Z","iopub.status.idle":"2023-06-16T07:54:08.502297Z","shell.execute_reply.started":"2023-06-16T07:54:08.489676Z","shell.execute_reply":"2023-06-16T07:54:08.501354Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# monet images Generator\nG = Generator(8).apply(init_weights).to(device)\n# photo images Generator\nF = Generator(8).apply(init_weights).to(device)\n# photo image Discriminator\nD_x = Discriminator().apply(init_weights).to(device)\n# monet images Discriminator\nD_y = Discriminator().apply(init_weights).to(device)\n\nmse_loss = nn.L1Loss()#mse_loss = nn.MSELoss()\nl1_loss = nn.L1Loss()\n\noptimizerG = torch.optim.Adam(\n    chain(G.parameters(), F.parameters()), lr=LR, betas=(BETA1, 0.999))\noptimizerD = torch.optim.Adam(\n    chain(D_x.parameters(), D_y.parameters()), lr=LR, betas=(BETA1, 0.999))\n\nfake_samples_monet = sample_fake(50)\nfake_samples_photo = sample_fake(50)\n\nscheduleG = torch.optim.lr_scheduler.StepLR(optimizerG, 5, .9)\nscheduleD = torch.optim.lr_scheduler.StepLR(optimizerD, 5, .9)","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:08.503697Z","iopub.execute_input":"2023-06-16T07:54:08.50413Z","iopub.status.idle":"2023-06-16T07:54:09.088159Z","shell.execute_reply.started":"2023-06-16T07:54:08.504099Z","shell.execute_reply":"2023-06-16T07:54:09.086883Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"start_epoch = 1\nhistory = {\n    \"gen_loss\":[], \"dis_loss\":[]\n}\nif LOAD_MODEL_PATH is not None:\n    state = load_checkpoint(LOAD_MODEL_PATH)\n    start_epoch = state[\"epoch\"]+1\n    history = state[\"history\"]\n    G.load_state_dict(state[\"G\"])\n    F.load_state_dict(state[\"F\"])\n    D_x.load_state_dict(state[\"D_x\"])\n    D_y.load_state_dict(state[\"D_y\"])\n    optimizerG.load_state_dict(state[\"optimizerG\"])\n    optimizerD.load_state_dict(state[\"optimizerD\"])","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:09.089497Z","iopub.execute_input":"2023-06-16T07:54:09.089881Z","iopub.status.idle":"2023-06-16T07:54:09.098004Z","shell.execute_reply.started":"2023-06-16T07:54:09.089849Z","shell.execute_reply":"2023-06-16T07:54:09.096929Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"Start learning on %s...\" % device)\nfor epoch in range(start_epoch, EPOCHS + 1):\n    avgG_loss = 0\n    avgD_loss = 0\n    dataset.set_train_mode()\n    tqdm_dataloader = tqdm(dataloader, leave=False, total=min(len(dataloader), 25))  # Set the total value to a maximum of 50\n    for i, (photo, monet) in enumerate(tqdm_dataloader):\n        if i >= 25:\n            break\n        # Rest of the training code...\n\n        # Generator training\n        optimizerG.zero_grad()\n        \n        fake_monet = G(photo)\n        fake_photo = F(monet)\n        \n        cycle_monet = G(fake_photo)\n        cycle_photo = F(fake_monet)\n        \n        identity_monet = G(monet)\n        identity_photo = F(photo)\n        \n        D_fake_monet = D_y(fake_monet)\n        D_fake_photo = D_x(fake_photo)\n        \n        real = get_real(D_fake_monet.size()).to(device)\n        \n        adversary_loss = mse_loss(D_fake_monet, real) + mse_loss(D_fake_photo, real)\n        cycle_loss = l1_loss(cycle_photo, photo) + l1_loss(cycle_monet, monet)\n        identity_loss = l1_loss(identity_monet, monet) + l1_loss(identity_photo, photo)\n        \n        loss = adversary_loss + LAMBDA*(cycle_loss + IDENTITY_COEF*identity_loss)\n        avgG_loss += loss.item()\n        \n        loss.backward()\n        optimizerG.step()\n        \n        # Discriminator training\n        optimizerD.zero_grad()\n        \n        fake_monet = fake_samples_monet(fake_monet.detach().cpu()).to(device)\n        fake_photo = fake_samples_photo(fake_photo.detach().cpu()).to(device)\n        \n        fake = get_fake(D_fake_monet.size()).to(device)\n        \n        real, fake = random_swap(real, fake, p=DISCRIMINATOR_LABEL_SWAP_P)\n\n        adversary_monet_loss = mse_loss(D_y(monet), real) + mse_loss(D_y(fake_monet), fake)\n        adversary_photo_loss = mse_loss(D_x(photo), real) + mse_loss(D_x(fake_photo), fake)\n        loss = adversary_monet_loss + adversary_photo_loss\n        avgD_loss += loss.item()\n        \n        loss.backward()\n        optimizerD.step()\n        \n        # logs\n        tqdm_dataloader.set_postfix(\n            D_loss=avgD_loss/(i+1),\n            G_loss=avgG_loss/(i+1),\n        )\n        \n    avgG_loss /= len(dataloader)\n    avgD_loss /= len(dataloader)\n    \n    history[\"gen_loss\"].append(avgG_loss)\n    history[\"dis_loss\"].append(avgD_loss)\n    \n    print(\n        \"Epoch %d | Generator Loss = %.4f | Distriminator Loss = %.4f\" \\\n        % (epoch, avgG_loss, avgD_loss)\n    )\n    if SAVE_MODEL_DIR is not None:\n        state = dict(\n            epoch=epoch,\n            history=history,\n            G=G.state_dict(), \n            F=F.state_dict(), \n            D_x=D_x.state_dict(), \n            D_y=D_y.state_dict(), \n            optimizerG=optimizerG.state_dict(), \n            optimizerD=optimizerD.state_dict(),\n        )\n        fn = \"current\"\n        path = os.path.join(SAVE_MODEL_DIR, fn)\n        save_checkpoint(state, path)\n    \n    if epoch%UPDATE_EVERY_N == 0:\n        clear_output(wait=True)\n        plot_loss(history[\"gen_loss\"], history[\"dis_loss\"], (5, 5))\n        dataset.set_test_mode()\n        for idx in range(N_SAMPLES):\n            with torch.no_grad():\n                photo = dataset_photos[idx][None, ...].to(device)\n                display_pair(to_img(photo[0]), to_img(G(photo)[0]))\n    \n    scheduleG.step()\n    scheduleD.step()","metadata":{"execution":{"iopub.status.busy":"2023-06-16T07:54:09.100288Z","iopub.execute_input":"2023-06-16T07:54:09.100729Z","iopub.status.idle":"2023-06-16T09:29:38.094045Z","shell.execute_reply.started":"2023-06-16T07:54:09.100689Z","shell.execute_reply":"2023-06-16T09:29:38.09299Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"**Make submission**","metadata":{}},{"cell_type":"code","source":"gen_images = []\nphoto_dataloader = DataLoader(dataset_photos, batch_size=1, num_workers=WORKERS)\ndataset.set_test_mode()\nwith torch.no_grad():\n    for img in tqdm(photo_dataloader):\n        gen = G(img.to(device))[0]\n        # unnormalize\n        gen_images.append(to_img(gen))\ndataset.set_train_mode()\n","metadata":{"execution":{"iopub.status.busy":"2023-06-16T09:29:38.095496Z","iopub.execute_input":"2023-06-16T09:29:38.096514Z","iopub.status.idle":"2023-06-16T11:14:17.39308Z","shell.execute_reply.started":"2023-06-16T09:29:38.096474Z","shell.execute_reply":"2023-06-16T11:14:17.391019Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def viz_gen_images(gen_images, n_samples=10):\n    for idx in np.random.permutation(len(gen_images))[:n_samples]:\n        print(idx)\n        display(gen_images[idx])","metadata":{"execution":{"iopub.status.busy":"2023-06-16T11:14:17.396284Z","iopub.execute_input":"2023-06-16T11:14:17.397038Z","iopub.status.idle":"2023-06-16T11:14:17.406284Z","shell.execute_reply.started":"2023-06-16T11:14:17.396974Z","shell.execute_reply":"2023-06-16T11:14:17.40492Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"viz_gen_images(gen_images, 10)","metadata":{"execution":{"iopub.status.busy":"2023-06-16T11:14:17.408021Z","iopub.execute_input":"2023-06-16T11:14:17.408479Z","iopub.status.idle":"2023-06-16T11:14:17.662722Z","shell.execute_reply.started":"2023-06-16T11:14:17.40844Z","shell.execute_reply":"2023-06-16T11:14:17.661432Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"for i, gen in enumerate(gen_images):\n    gen.save(\"./images/\" + str(i+1) + \".jpg\")","metadata":{"execution":{"iopub.status.busy":"2023-06-16T11:14:17.664193Z","iopub.execute_input":"2023-06-16T11:14:17.664516Z","iopub.status.idle":"2023-06-16T11:14:21.729427Z","shell.execute_reply.started":"2023-06-16T11:14:17.664489Z","shell.execute_reply":"2023-06-16T11:14:21.728168Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import shutil\nshutil.make_archive(\"./images\", 'zip', \"./images\")","metadata":{"execution":{"iopub.status.busy":"2023-06-16T11:14:21.73101Z","iopub.execute_input":"2023-06-16T11:14:21.731387Z","iopub.status.idle":"2023-06-16T11:14:25.104015Z","shell.execute_reply.started":"2023-06-16T11:14:21.731356Z","shell.execute_reply":"2023-06-16T11:14:25.102751Z"},"trusted":true},"execution_count":null,"outputs":[]}]}