{"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":"2022-05-14T09:07:01.494777Z","iopub.execute_input":"2022-05-14T09:07:01.495117Z","iopub.status.idle":"2022-05-14T09:07:12.7886Z","shell.execute_reply.started":"2022-05-14T09:07:01.495045Z","shell.execute_reply":"2022-05-14T09:07:12.787773Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"!mkdir images\n!mkdir models","metadata":{"execution":{"iopub.status.busy":"2022-05-14T09:07:12.790861Z","iopub.execute_input":"2022-05-14T09:07:12.791279Z","iopub.status.idle":"2022-05-14T09:07:14.109254Z","shell.execute_reply.started":"2022-05-14T09:07:12.791239Z","shell.execute_reply":"2022-05-14T09:07:14.108368Z"},"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":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2022-05-14T09:07:14.111058Z","iopub.execute_input":"2022-05-14T09:07:14.111325Z","iopub.status.idle":"2022-05-14T09:07:16.497304Z","shell.execute_reply.started":"2022-05-14T09:07:14.111288Z","shell.execute_reply":"2022-05-14T09:07:16.496595Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"PATH = \"../input/gan-getting-started\"\nIMG_SIZE = (256, 256)","metadata":{"execution":{"iopub.status.busy":"2022-05-14T09:07:16.499454Z","iopub.execute_input":"2022-05-14T09:07:16.499691Z","iopub.status.idle":"2022-05-14T09:07:16.504202Z","shell.execute_reply.started":"2022-05-14T09:07:16.499658Z","shell.execute_reply":"2022-05-14T09:07:16.503434Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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":"2022-05-14T09:07:16.505698Z","iopub.execute_input":"2022-05-14T09:07:16.506118Z","iopub.status.idle":"2022-05-14T09:07:16.515179Z","shell.execute_reply.started":"2022-05-14T09:07:16.506077Z","shell.execute_reply":"2022-05-14T09:07:16.514119Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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":"2022-05-14T09:07:16.516615Z","iopub.execute_input":"2022-05-14T09:07:16.516948Z","iopub.status.idle":"2022-05-14T09:07:16.527329Z","shell.execute_reply.started":"2022-05-14T09:07:16.516914Z","shell.execute_reply":"2022-05-14T09:07:16.526674Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import albumentations as A\nfrom albumentations.pytorch import ToTensorV2\n\ntransform = A.Compose([\n    A.Resize(*IMG_SIZE),\n    A.Normalize(mean=[.5]*3, std=[.5]*3),\n    ToTensorV2(),\n])\n\n# aug = A.Compose ([\n#     A.RGBShift(r_shift_limit=10, g_shift_limit=10, b_shift_limit=10, p=.5),\n#     A.HorizontalFlip(p=.5),\n#     A.ShiftScaleRotate(rotate_limit=10),\n#     A.RandomBrightnessContrast(brightness_limit=.2, contrast_limit=.2),\n#     transform\n# ])","metadata":{"execution":{"iopub.status.busy":"2022-05-14T09:07:16.528553Z","iopub.execute_input":"2022-05-14T09:07:16.528898Z","iopub.status.idle":"2022-05-14T09:07:17.902266Z","shell.execute_reply.started":"2022-05-14T09:07:16.528785Z","shell.execute_reply":"2022-05-14T09:07:17.901529Z"},"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":"2022-05-14T09:07:17.90363Z","iopub.execute_input":"2022-05-14T09:07:17.903895Z","iopub.status.idle":"2022-05-14T09:08:02.024702Z","shell.execute_reply.started":"2022-05-14T09:07:17.903861Z","shell.execute_reply":"2022-05-14T09:08:02.024016Z"},"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":"2022-05-14T09:08:02.025997Z","iopub.execute_input":"2022-05-14T09:08:02.026394Z","iopub.status.idle":"2022-05-14T09:08:02.030382Z","shell.execute_reply.started":"2022-05-14T09:08:02.026357Z","shell.execute_reply":"2022-05-14T09:08:02.029708Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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":"2022-05-14T09:08:02.033698Z","iopub.execute_input":"2022-05-14T09:08:02.034176Z","iopub.status.idle":"2022-05-14T09:08:02.04256Z","shell.execute_reply.started":"2022-05-14T09:08:02.034129Z","shell.execute_reply":"2022-05-14T09:08:02.041792Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# dataset_monet.transform = aug\n# print('Augmentated')\n# plot_images_grid([to_img(dataset_monet[i]) for i in range(25)], figsize=(20, 20))\n# dataset_monet.transform = transform\n# print('Original')\n# plot_images_grid([to_img(dataset_monet[i]) for i in range(25)], figsize=(20, 20))","metadata":{"execution":{"iopub.status.busy":"2022-05-14T09:08:02.044694Z","iopub.execute_input":"2022-05-14T09:08:02.045018Z","iopub.status.idle":"2022-05-14T09:08:02.052706Z","shell.execute_reply.started":"2022-05-14T09:08:02.044974Z","shell.execute_reply":"2022-05-14T09:08:02.051927Z"},"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":"2022-05-14T09:08:02.053888Z","iopub.execute_input":"2022-05-14T09:08:02.054224Z","iopub.status.idle":"2022-05-14T09:08:02.06406Z","shell.execute_reply.started":"2022-05-14T09:08:02.054187Z","shell.execute_reply":"2022-05-14T09:08:02.063376Z"},"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":"2022-05-14T09:08:02.066974Z","iopub.execute_input":"2022-05-14T09:08:02.067247Z","iopub.status.idle":"2022-05-14T09:08:02.442483Z","shell.execute_reply.started":"2022-05-14T09:08:02.067192Z","shell.execute_reply":"2022-05-14T09:08:02.44011Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"BATCH_SIZE = 1\nWORKERS = 2","metadata":{"execution":{"iopub.status.busy":"2022-05-14T09:08:02.44356Z","iopub.execute_input":"2022-05-14T09:08:02.443856Z","iopub.status.idle":"2022-05-14T09:08:02.448329Z","shell.execute_reply.started":"2022-05-14T09:08:02.443819Z","shell.execute_reply":"2022-05-14T09:08:02.447654Z"},"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":"2022-05-14T09:08:02.449958Z","iopub.execute_input":"2022-05-14T09:08:02.450484Z","iopub.status.idle":"2022-05-14T09:08:02.457742Z","shell.execute_reply.started":"2022-05-14T09:08:02.450446Z","shell.execute_reply":"2022-05-14T09:08:02.456879Z"},"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),\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),\n    ]\n    return nn.Sequential(*modules)\n    ","metadata":{"execution":{"iopub.status.busy":"2022-05-14T09:08:02.459193Z","iopub.execute_input":"2022-05-14T09:08:02.459812Z","iopub.status.idle":"2022-05-14T09:08:02.470001Z","shell.execute_reply.started":"2022-05-14T09:08:02.459662Z","shell.execute_reply":"2022-05-14T09:08:02.469072Z"},"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)\n    )","metadata":{"execution":{"iopub.status.busy":"2022-05-14T09:08:02.471643Z","iopub.execute_input":"2022-05-14T09:08:02.472198Z","iopub.status.idle":"2022-05-14T09:08:02.479342Z","shell.execute_reply.started":"2022-05-14T09:08:02.47216Z","shell.execute_reply":"2022-05-14T09:08:02.478335Z"},"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)\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    \n    ","metadata":{"execution":{"iopub.status.busy":"2022-05-14T09:08:02.48091Z","iopub.execute_input":"2022-05-14T09:08:02.48158Z","iopub.status.idle":"2022-05-14T09:08:02.491741Z","shell.execute_reply.started":"2022-05-14T09:08:02.481538Z","shell.execute_reply":"2022-05-14T09:08:02.490922Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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)\n    \n    ","metadata":{"execution":{"iopub.status.busy":"2022-05-14T09:08:02.493021Z","iopub.execute_input":"2022-05-14T09:08:02.493275Z","iopub.status.idle":"2022-05-14T09:08:02.505592Z","shell.execute_reply.started":"2022-05-14T09:08:02.493243Z","shell.execute_reply":"2022-05-14T09:08:02.504836Z"},"trusted":true},"execution_count":null,"outputs":[]},{"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\n            out_ch = 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":"2022-05-14T09:08:02.507113Z","iopub.execute_input":"2022-05-14T09:08:02.507437Z","iopub.status.idle":"2022-05-14T09:08:02.517125Z","shell.execute_reply.started":"2022-05-14T09:08:02.507396Z","shell.execute_reply":"2022-05-14T09:08:02.516077Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print(\"Generator\")\nprint(summary(Generator(), (1, 3, *IMG_SIZE)))\nprint()\nprint(\"Discriminator\")\nprint(summary(Discriminator(), (1, 3, *IMG_SIZE)))","metadata":{"execution":{"iopub.status.busy":"2022-05-14T09:08:02.518494Z","iopub.execute_input":"2022-05-14T09:08:02.518888Z","iopub.status.idle":"2022-05-14T09:08:10.789954Z","shell.execute_reply.started":"2022-05-14T09:08:02.518849Z","shell.execute_reply":"2022-05-14T09:08:10.789236Z"},"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":"2022-05-14T09:08:10.791064Z","iopub.execute_input":"2022-05-14T09:08:10.792169Z","iopub.status.idle":"2022-05-14T09:08:10.80006Z","shell.execute_reply.started":"2022-05-14T09:08:10.792129Z","shell.execute_reply":"2022-05-14T09:08:10.799202Z"},"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":"2022-05-14T09:08:10.8017Z","iopub.execute_input":"2022-05-14T09:08:10.802098Z","iopub.status.idle":"2022-05-14T09:08:10.810945Z","shell.execute_reply.started":"2022-05-14T09:08:10.80206Z","shell.execute_reply":"2022-05-14T09:08:10.810203Z"},"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":"2022-05-14T09:08:10.812573Z","iopub.execute_input":"2022-05-14T09:08:10.812896Z","iopub.status.idle":"2022-05-14T09:08:10.823324Z","shell.execute_reply.started":"2022-05-14T09:08:10.812859Z","shell.execute_reply":"2022-05-14T09:08:10.822575Z"},"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":"2022-05-14T09:08:10.82496Z","iopub.execute_input":"2022-05-14T09:08:10.825312Z","iopub.status.idle":"2022-05-14T09:08:10.833436Z","shell.execute_reply.started":"2022-05-14T09:08:10.825276Z","shell.execute_reply":"2022-05-14T09:08:10.832603Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Log ","metadata":{}},{"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":"2022-05-14T09:08:10.834909Z","iopub.execute_input":"2022-05-14T09:08:10.835197Z","iopub.status.idle":"2022-05-14T09:08:10.843405Z","shell.execute_reply.started":"2022-05-14T09:08:10.835162Z","shell.execute_reply":"2022-05-14T09:08:10.842741Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Training","metadata":{}},{"cell_type":"code","source":"def save_checkpoint(state, save_path):\n    torch.save(state, save_path)","metadata":{"execution":{"iopub.status.busy":"2022-05-14T09:08:10.844733Z","iopub.execute_input":"2022-05-14T09:08:10.844982Z","iopub.status.idle":"2022-05-14T09:08:10.851738Z","shell.execute_reply.started":"2022-05-14T09:08:10.844949Z","shell.execute_reply":"2022-05-14T09:08:10.850945Z"},"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":"2022-05-14T09:08:10.856529Z","iopub.execute_input":"2022-05-14T09:08:10.85681Z","iopub.status.idle":"2022-05-14T09:08:10.860836Z","shell.execute_reply.started":"2022-05-14T09:08:10.856784Z","shell.execute_reply":"2022-05-14T09:08:10.860049Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"from itertools import chain","metadata":{"execution":{"iopub.status.busy":"2022-05-14T09:08:10.862561Z","iopub.execute_input":"2022-05-14T09:08:10.862848Z","iopub.status.idle":"2022-05-14T09:08:10.869378Z","shell.execute_reply.started":"2022-05-14T09:08:10.86281Z","shell.execute_reply":"2022-05-14T09:08:10.868524Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"LR = 2e-4\nLAMBDA = 10\nEPOCHS = 40\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":"2022-05-14T09:08:10.870573Z","iopub.execute_input":"2022-05-14T09:08:10.870905Z","iopub.status.idle":"2022-05-14T09:08:10.878648Z","shell.execute_reply.started":"2022-05-14T09:08:10.87087Z","shell.execute_reply":"2022-05-14T09:08:10.878007Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"device = \"cuda\" if torch.cuda.is_available() else \"cpu\"","metadata":{"execution":{"iopub.status.busy":"2022-05-14T09:08:10.880205Z","iopub.execute_input":"2022-05-14T09:08:10.880476Z","iopub.status.idle":"2022-05-14T09:08:10.887635Z","shell.execute_reply.started":"2022-05-14T09:08:10.880431Z","shell.execute_reply":"2022-05-14T09:08:10.886838Z"},"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.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":"2022-05-14T09:08:10.888945Z","iopub.execute_input":"2022-05-14T09:08:10.88931Z","iopub.status.idle":"2022-05-14T09:08:11.346927Z","shell.execute_reply.started":"2022-05-14T09:08:10.889276Z","shell.execute_reply":"2022-05-14T09:08:11.346171Z"},"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":"2022-05-14T09:08:11.348375Z","iopub.execute_input":"2022-05-14T09:08:11.348659Z","iopub.status.idle":"2022-05-14T09:08:11.355469Z","shell.execute_reply.started":"2022-05-14T09:08:11.348623Z","shell.execute_reply":"2022-05-14T09:08:11.354523Z"},"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(\n        dataloader, leave=False, total=len(dataloader))\n    for i, (photo, monet) in enumerate(tqdm_dataloader):\n        photo = photo.to(device)\n        monet = monet.to(device)\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":"2022-05-14T09:08:11.357259Z","iopub.execute_input":"2022-05-14T09:08:11.357565Z"},"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()","metadata":{"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":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"viz_gen_images(gen_images, 10)","metadata":{"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":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"len(os.listdir(\"./images\"))","metadata":{"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import shutil\nshutil.make_archive(\"./images\", 'zip', \"./images\")","metadata":{"trusted":true},"execution_count":null,"outputs":[]}]}