{"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":"%config Completer.use_jedi = False","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-11-27T05:58:32.294448Z","iopub.execute_input":"2021-11-27T05:58:32.295029Z","iopub.status.idle":"2021-11-27T05:58:32.321054Z","shell.execute_reply.started":"2021-11-27T05:58:32.294927Z","shell.execute_reply":"2021-11-27T05:58:32.320412Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"import os\nimport datetime\nimport glob\nimport time\nimport cv2\nimport itertools\nfrom tqdm.notebook import tqdm\nimport shutil\n\nfrom PIL import Image\n\nimport numpy as np\nimport matplotlib.pyplot as plt\n\nimport torch\nimport torch.nn as nn\nimport torch.optim as optim\nfrom torch.utils.data import DataLoader, Dataset\n\nimport torchvision.transforms as transforms\nfrom torchvision.utils import make_grid","metadata":{"execution":{"iopub.status.busy":"2021-11-27T05:58:33.729532Z","iopub.execute_input":"2021-11-27T05:58:33.730069Z","iopub.status.idle":"2021-11-27T05:58:35.496065Z","shell.execute_reply.started":"2021-11-27T05:58:33.730033Z","shell.execute_reply":"2021-11-27T05:58:35.495306Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Data loader","metadata":{}},{"cell_type":"code","source":"img_path = '../input/gan-getting-started/'\nmonet_path = glob.glob(img_path + 'monet_jpg/*')\nphoto_path = glob.glob(img_path + 'photo_jpg/*')\n\nprint('Dataset')\nprint(f'- monet data : {len(monet_path)}\\n- photo data : {len(photo_path)}')","metadata":{"execution":{"iopub.status.busy":"2021-11-27T05:58:38.005363Z","iopub.execute_input":"2021-11-27T05:58:38.005608Z","iopub.status.idle":"2021-11-27T05:58:38.567016Z","shell.execute_reply.started":"2021-11-27T05:58:38.005582Z","shell.execute_reply":"2021-11-27T05:58:38.566253Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Custom_dataset(Dataset):\n    def __init__(self, img_path : list , transforms = None, mode = 'train'):\n        super().__init__()\n\n        self.path_monet = img_path[0]\n        self.path_photo = img_path[1]\n        self.transforms = transforms\n        self.mode = mode\n        \n    def __getitem__(self, idx):\n        if self.mode == 'train':\n            monet_img = self.path_monet[idx]\n            monet_img = Image.open(monet_img).convert('RGB')\n            monet_img = self.transforms(monet_img)\n            \n            photo_idx = np.random.randint(0, len(self.path_photo))\n            photo_img = self.path_photo[photo_idx]\n            photo_img = Image.open(photo_img).convert('RGB')\n            photo_img = self.transforms(photo_img)\n            \n            return monet_img, photo_img\n    \n    def __len__(self):\n        if self.mode == 'train':\n            return len(self.path_monet)\n        elif self.mode == 'test':\n            return len(self.path_photo)","metadata":{"execution":{"iopub.status.busy":"2021-11-27T05:58:38.568737Z","iopub.execute_input":"2021-11-27T05:58:38.569219Z","iopub.status.idle":"2021-11-27T05:58:38.578163Z","shell.execute_reply.started":"2021-11-27T05:58:38.569180Z","shell.execute_reply":"2021-11-27T05:58:38.577455Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"def show_img(img):\n    img = make_grid(img, nrow = 6).permute([1, 2, 0]).detach().numpy()\n    plt.figure(figsize = (12, 8))\n    plt.imshow(img)\n    plt.show()","metadata":{"execution":{"iopub.status.busy":"2021-11-27T05:58:39.138999Z","iopub.execute_input":"2021-11-27T05:58:39.139425Z","iopub.status.idle":"2021-11-27T05:58:39.144917Z","shell.execute_reply.started":"2021-11-27T05:58:39.139391Z","shell.execute_reply":"2021-11-27T05:58:39.143113Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"img_path = [monet_path, photo_path]\nsample_transform = transforms.Compose([\n    transforms.ToTensor()\n])\nsample = Custom_dataset(img_path = img_path, transforms = sample_transform, mode = 'train')\nsample_loader = DataLoader(sample, batch_size = 6)\nsample_monet, sample_photo = next(iter(sample_loader))","metadata":{"execution":{"iopub.status.busy":"2021-11-27T05:58:39.903828Z","iopub.execute_input":"2021-11-27T05:58:39.904463Z","iopub.status.idle":"2021-11-27T05:58:40.075975Z","shell.execute_reply.started":"2021-11-27T05:58:39.904426Z","shell.execute_reply":"2021-11-27T05:58:40.075253Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print('Monet data')\nshow_img(sample_monet)\n\nprint('Photo data')\nshow_img(sample_photo)","metadata":{"execution":{"iopub.status.busy":"2021-11-27T05:58:40.321224Z","iopub.execute_input":"2021-11-27T05:58:40.321462Z","iopub.status.idle":"2021-11-27T05:58:40.941202Z","shell.execute_reply.started":"2021-11-27T05:58:40.321436Z","shell.execute_reply":"2021-11-27T05:58:40.940567Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Modeling","metadata":{}},{"cell_type":"code","source":"# Define Conv Block\n'''\n1. Conv_up\n2. Conv_down\n3. Residual_block\n'''\n\n# 1.Conv_up\nclass Conv_up(nn.Module):\n    '''\n    convTranspose - instanceNorm - ReLU - (dropout)\n    '''\n    def __init__(self, in_ch, out_ch, kernel_size = 4, stride = 2, \n                 padding = 1, output_padding = 1, drop_out = True):\n        super().__init__()\n        \n        self.convT = nn.ConvTranspose2d(in_ch, out_ch,\n                                       kernel_size = kernel_size,\n                                       stride = stride,\n                                       padding = padding,\n                                       output_padding = output_padding,\n                                       bias = False)\n        self.instance_norm = nn.InstanceNorm2d(out_ch)\n        self.relu = nn.ReLU()\n        self.drop_out = drop_out\n        \n    def forward(self, x):\n        x = self.convT(x)\n        x = self.instance_norm(x)\n        x = self.relu(x)\n        if self.drop_out:\n            x = nn.Dropout2d(0.5)(x)\n\n        return x\n\n# 2. Conv_down\nclass Conv_down(nn.Module):\n    '''\n    Conv2d - instanceNorm - LeakyReLU\n    '''\n    def __init__(self, in_ch, out_ch,\n                 kernel_size = 4,\n                 stride = 2,\n                 padding = 1,\n                 batch_Norm = True):\n        super().__init__()\n        \n        self.conv = nn.Conv2d(in_ch, out_ch,\n                             kernel_size = kernel_size,\n                             stride = stride,\n                             padding = padding,\n                             bias = True)\n        self.instance_norm = nn.InstanceNorm2d(out_ch)\n        self.relu = nn.ReLU()\n        self.batch = batch_Norm\n        \n    def forward(self, x):\n        x = self.conv(x)\n        if self.batch:\n            x = self.instance_norm(x)\n        x = self.relu(x)\n        \n        return x\n    \n# 3. Residual_block\nclass Residual_block(nn.Module):\n    '''\n    Conv2d - InstanceNorm - Relu - Conv2d - InstanceNorm\n    '''\n    def __init__(self, in_ch, out_ch, kernel_size = 3, stride = 1, padding = 1):\n        super().__init__()\n        \n        self.res = nn.Sequential(\n            nn.Conv2d(in_ch, out_ch,\n                     kernel_size = kernel_size,\n                     stride = stride,\n                     padding = padding,\n                     padding_mode = 'reflect',\n                     bias = False),\n            nn.InstanceNorm2d(out_ch),\n            nn.ReLU(),\n            nn.Conv2d(out_ch, out_ch,\n                     kernel_size = kernel_size,\n                     stride = stride,\n                     padding = padding,\n                     padding_mode = 'reflect',\n                     bias = False),\n            nn.InstanceNorm2d(out_ch)\n        )\n    \n    def forward(self, x):\n        return x + self.res(x)","metadata":{"execution":{"iopub.status.busy":"2021-11-27T05:58:42.205066Z","iopub.execute_input":"2021-11-27T05:58:42.205454Z","iopub.status.idle":"2021-11-27T05:58:42.221160Z","shell.execute_reply.started":"2021-11-27T05:58:42.205420Z","shell.execute_reply":"2021-11-27T05:58:42.220492Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Discriminator(nn.Module):\n    '''\n    C64 - C128 - C256 - C512 - C512 - 1\n    '''\n    def __init__(self, n_features = 64):\n        super().__init__()\n        \n        self.main = nn.Sequential(\n            Conv_down(3, n_features, batch_Norm = False),\n            Conv_down(n_features * 1, n_features * 2),\n            Conv_down(n_features * 2, n_features * 4),\n            Conv_down(n_features * 4, n_features * 8, stride = 1),\n            nn.Conv2d(n_features * 8, 1, 4, 1, 1, bias = False)\n        )\n    \n    def forward(self, x):\n        x = self.main(x)\n        \n        return x\n    \ndef test():\n    D = Discriminator()\n    x = torch.randn(1, 3, 256, 256)\n    out = D(x)\n    print(out.shape)\n    print('Discriminator is OK')\n    \ntest()","metadata":{"execution":{"iopub.status.busy":"2021-11-27T05:58:43.021555Z","iopub.execute_input":"2021-11-27T05:58:43.022229Z","iopub.status.idle":"2021-11-27T05:58:43.200729Z","shell.execute_reply.started":"2021-11-27T05:58:43.022192Z","shell.execute_reply":"2021-11-27T05:58:43.199995Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Generator(nn.Module):\n    '''\n    kernel   D64 - D128 - D256 - R256 * n - U128 - U64 - U3\n    filter   7x7 - 3x3  - 3x3  -          -  3x3 - 3x3 - 7x7\n    stride    1     2      2       1          2     2     1\n    '''\n    def __init__(self, n_features = 64, n_res = 9):\n        super().__init__()\n        \n        self.main = nn.Sequential(\n            nn.ReflectionPad2d(3),\n            Conv_down(3, n_features, 7, 1, 0, False),\n            Conv_down(n_features * 1, n_features * 2, 3, 2),\n            Conv_down(n_features * 2, n_features * 4, 3, 2),\n            *[\n                Residual_block(n_features * 4, n_features * 4) for _ in range(n_res)\n            ],\n            Conv_up(n_features * 4, n_features * 2, 3, 2, 1),\n            Conv_up(n_features * 2, n_features * 1, 3, 2, 1),\n            nn.ReflectionPad2d(3),\n            nn.Conv2d(n_features, 3, 7, 1, 0, bias = False),\n            nn.Tanh()   \n        )\n        \n    def forward(self, x):\n        x = self.main(x)\n        \n        return x\n    \ndef test():\n    G = Generator()\n    x = torch.randn(1, 3, 256, 256)\n    out = G(x)\n    print(out.shape)\n    print('Generator is OK')\n    \ntest()","metadata":{"execution":{"iopub.status.busy":"2021-11-27T05:58:43.391072Z","iopub.execute_input":"2021-11-27T05:58:43.391318Z","iopub.status.idle":"2021-11-27T05:58:45.694841Z","shell.execute_reply.started":"2021-11-27T05:58:43.391293Z","shell.execute_reply":"2021-11-27T05:58:45.693357Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# weight initialization \ndef weight_init(m : 'model'):\n    classname = m.__class__.__name__\n    if classname.find('Conv') != -1:\n        nn.init.normal_(m.weight.data, 0.0, 0.02)\n    elif classname.find('Instance') != -1:\n        nn.init.normal_(m.weight.data, 1.0, 0.02)\n        nn.init.constant_(m.bias.data, 0)","metadata":{"execution":{"iopub.status.busy":"2021-11-27T05:58:45.696547Z","iopub.execute_input":"2021-11-27T05:58:45.696877Z","iopub.status.idle":"2021-11-27T05:58:45.703194Z","shell.execute_reply.started":"2021-11-27T05:58:45.696839Z","shell.execute_reply":"2021-11-27T05:58:45.702021Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Train","metadata":{}},{"cell_type":"code","source":"device = 'cuda' if torch.cuda.is_available() else 'cpu'\n\n# Define Generator\n# monet -> photo\nnetG_A = Generator().to(device)\n# photo -> monet\nnetG_B = Generator().to(device)\n\n# Define Discrimator\nnetD_A = Discriminator().to(device)\nnetD_B = Discriminator().to(device)\n\n# weight initialization\nweight_init(netG_A)\nweight_init(netG_B)\nweight_init(netD_A)\nweight_init(netD_B)\n\n# setting\nEPOCHS = 100\nlr = 2e-4\nb1 = 0.5\nb2 = 0.999\n\n# optimizer\nnetG_optim = optim.Adam(itertools.chain(netG_A.parameters(), netG_B.parameters()), lr = lr, betas = (b1, b2))\nnetD_optim = optim.Adam(itertools.chain(netD_A.parameters(), netD_B.parameters()), lr = lr, betas = (b1, b2))","metadata":{"execution":{"iopub.status.busy":"2021-11-27T06:01:01.129704Z","iopub.execute_input":"2021-11-27T06:01:01.130391Z","iopub.status.idle":"2021-11-27T06:01:01.405961Z","shell.execute_reply.started":"2021-11-27T06:01:01.130356Z","shell.execute_reply":"2021-11-27T06:01:01.405241Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Data loader\ntransform = transforms.Compose([\n            transforms.ToTensor(),\n            transforms.Normalize([0.5, 0.5, 0.5], \n                                [0.5, 0.5, 0.5])\n])\n\ntrain_dataset = Custom_dataset([monet_path, photo_path], transforms = transform, mode = 'train')\ntrain_loader = DataLoader(train_dataset, batch_size = 1, shuffle = True)","metadata":{"execution":{"iopub.status.busy":"2021-11-27T05:58:48.569872Z","iopub.execute_input":"2021-11-27T05:58:48.570109Z","iopub.status.idle":"2021-11-27T05:58:48.575081Z","shell.execute_reply.started":"2021-11-27T05:58:48.570077Z","shell.execute_reply":"2021-11-27T05:58:48.574438Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# loss function\n'''\n1. GAN loss - L2\n2. Cycle loss - L1\n3. Identity loss - L1\n'''\n\nGAN_LOSS = nn.MSELoss()\nCycle_LOSS = nn.L1Loss()\nIdentity_LOSS = nn.L1Loss()\n\ndef train_model(EPOCHS = EPOCHS):\n    train_hist = {}\n    train_hist['G_losses'] = []\n    train_hist['D_losses'] = []\n    \n    print('train is starting')\n    \n    for epoch in range(EPOCHS):\n        t = time.time()\n        \n        netG_A.train()\n        netG_B.train()\n        \n        netD_A.train()\n        netD_B.train()\n        \n        G_losses = 0\n        D_losses = 0\n        \n        for A, B in train_loader:\n            # A : monet, B : photo\n            A, B = A.to(device), B.to(device)\n            \n            A2B = netG_A(A) # fake B\n            B2A = netG_B(B) # fake A\n            \n            pred_real_A = netD_A(A)\n            pred_fake_A = netD_A(B2A.detach())\n            \n            # GAN LOSS for Discrimiator\n            D_A_loss = GAN_LOSS(pred_real_A, torch.ones_like(pred_real_A)) +\\\n                        GAN_LOSS(pred_fake_A, torch.zeros_like(pred_fake_A))\n\n            pred_real_B = netD_B(B)\n            pred_fake_B = netD_B(A2B.detach())\n            \n            D_B_loss = GAN_LOSS(pred_real_B, torch.ones_like(pred_real_B)) +\\\n                        GAN_LOSS(pred_fake_B, torch.zeros_like(pred_fake_B))\n            \n            D_loss = (D_A_loss + D_B_loss) / 2\n                    \n            \n            A2B2A = netG_B(A2B) # fake A\n            B2A2B = netG_A(B2A) # fake B\n            \n            pred_fake_A = netD_A(B2A)\n            pred_fake_B = netD_B(A2B)\n            \n            # GAN LOSS for Generator\n            \n            G_GAN_loss = GAN_LOSS(pred_fake_A, torch.ones_like(pred_fake_A)) +\\\n                          GAN_LOSS(pred_fake_B, torch.ones_like(pred_fake_B))\n            \n            # Cycle LOSS\n            G_Cycle_loss = Cycle_LOSS(A2B2A, A) + Cycle_LOSS(B2A2B, B)\n            \n            # Identity LOSS\n            G_identity_loss = Identity_LOSS(netG_A(B), B) + Identity_LOSS(netG_B(A), A)\n            \n            G_loss = G_GAN_loss + 10 * G_Cycle_loss + 5 * G_identity_loss\n               \n            # backward Generator\n            netG_optim.zero_grad()\n            G_loss.backward()\n            netG_optim.step()\n         \n            # backward Discriminator\n            netD_optim.zero_grad()\n            D_loss.backward()\n            netD_optim.step()     \n            \n            D_losses += D_loss / len(train_loader)\n            G_losses += G_loss / len(train_loader)\n            \n        print(f'[{epoch + 1}/{EPOCHS}]\\tD_loss : {D_losses:.6f}\\tG_loss : {G_losses:.6f}\\ttime : {time.time() - t:.3f}s')\n     \n        train_hist['G_losses'].append(G_losses.item())\n        train_hist['D_losses'].append(D_losses.item())\n\n        # save model per 10epochs\n        if (epoch + 1) % 10 == 0:\n            if not os.path.exists('/kaggle/working/model_G'):\n                os.makedirs('/kaggle/working/model_G')\n\n            if not os.path.exists('/kaggle/working/model_D'):\n                os.makedirs('/kaggle/working/model_D')\n\n            # save Generator\n            torch.save(netG_A, '/kaggle/working/model_G/' + f'netG(uNet)_A{epoch + 1}.pt')\n            torch.save(netG_B, '/kaggle/working/model_G/' + f'netG(uNet)_B{epoch + 1}.pt')\n\n            # save Discriminator\n            torch.save(netD_A, '/kaggle/working/model_D/' + f'netD(uNet)_A{epoch + 1}.pt')\n            torch.save(netD_B, '/kaggle/working/model_D/' + f'netD(uNet)_B{epoch + 1}.pt')\n\n            print(f'Model is saved at {epoch + 1}epochs')\n\n\n    return train_hist","metadata":{"execution":{"iopub.status.busy":"2021-11-27T06:01:08.576292Z","iopub.execute_input":"2021-11-27T06:01:08.576541Z","iopub.status.idle":"2021-11-27T06:01:08.594184Z","shell.execute_reply.started":"2021-11-27T06:01:08.576515Z","shell.execute_reply":"2021-11-27T06:01:08.593404Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"train_hist = train_model(EPOCHS)","metadata":{"execution":{"iopub.status.busy":"2021-11-27T06:01:09.069216Z","iopub.execute_input":"2021-11-27T06:01:09.069733Z","iopub.status.idle":"2021-11-27T06:01:41.553579Z","shell.execute_reply.started":"2021-11-27T06:01:09.069697Z","shell.execute_reply":"2021-11-27T06:01:41.551783Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"plt.plot(train_hist['G_losses'], label = 'G loss')\nplt.plot(train_hist['D_losses'], label = 'D loss')\nplt.legend()\nplt.show()","metadata":{"execution":{"iopub.status.busy":"2021-11-27T06:00:46.735501Z","iopub.execute_input":"2021-11-27T06:00:46.736047Z","iopub.status.idle":"2021-11-27T06:00:46.937758Z","shell.execute_reply.started":"2021-11-27T06:00:46.736010Z","shell.execute_reply":"2021-11-27T06:00:46.937087Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"","metadata":{},"execution_count":null,"outputs":[]}]}