{"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":"markdown","source":"# Реализация CycleGAN на архитектуре U-NET для генератора и PatchGAN для дискриминатора","metadata":{"id":"_54dGD_-Z6As"}},{"cell_type":"markdown","source":"## Imports","metadata":{"id":"8j1utLEFw-3T"}},{"cell_type":"code","source":"import torch\nimport torchvision as tv\nfrom torch.utils.data import Dataset, DataLoader\nfrom torchvision.utils import make_grid\nimport torchvision.utils as vutils\nimport torch.backends.cudnn as cudnn\n\nimport os\nimport numpy as np\nfrom PIL import Image\nimport itertools\nimport random\nimport glob\nimport time\nimport errno\nimport shutil\nimport sys\n\nfrom matplotlib import pyplot as plt\nfrom math import sqrt\n%matplotlib inline","metadata":{"id":"ns7EF1sqxBez","execution":{"iopub.status.busy":"2022-02-24T11:06:22.034768Z","iopub.execute_input":"2022-02-24T11:06:22.035095Z","iopub.status.idle":"2022-02-24T11:06:22.044707Z","shell.execute_reply.started":"2022-02-24T11:06:22.03506Z","shell.execute_reply":"2022-02-24T11:06:22.043962Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Глобальные параметры","metadata":{"id":"8QTfx7J0rds3"}},{"cell_type":"code","source":"# GPU / CPU\ndevice = torch.device('cpu')\n\nif torch.cuda.is_available():\n    device = torch.device('cuda')\n\ndevice","metadata":{"id":"2ScDvyvA-OPH","outputId":"7cf8e3bb-f684-4765-bac6-0e114af45f18","execution":{"iopub.status.busy":"2022-02-24T11:06:22.046372Z","iopub.execute_input":"2022-02-24T11:06:22.046765Z","iopub.status.idle":"2022-02-24T11:06:22.057155Z","shell.execute_reply.started":"2022-02-24T11:06:22.046728Z","shell.execute_reply":"2022-02-24T11:06:22.056422Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"PARAMS = {\n    'data_dir': '../input/gan-getting-started', # Directory for dataset\n    'datasetA': 'monet_jpg', # Dataset A\n    'datasetB': 'photo_jpg', # Dataset B\n    'out_dir': 'output', # Directory for output\n    'epochs': 200, # number of epochs\n    'batch_size': 5, # size of batches in training\n    'test_batch_size': 4, # size of batches in inference\n    'lr': 0.0002, # learning rate\n    'img_size': 256, # size of images\n    'channels': 3, # number of image channels\n    'num_blocks': 9, # number of residual blocks\n    'log_interval': 20, # interval between logging and image sampling\n    'seed': 1, # random seed\n}","metadata":{"id":"5Yi3RdEmrFNk","execution":{"iopub.status.busy":"2022-02-24T11:06:22.058407Z","iopub.execute_input":"2022-02-24T11:06:22.058899Z","iopub.status.idle":"2022-02-24T11:06:22.065868Z","shell.execute_reply.started":"2022-02-24T11:06:22.058864Z","shell.execute_reply":"2022-02-24T11:06:22.065055Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Utils","metadata":{"id":"jRiqhoyJXX0N"}},{"cell_type":"code","source":"def clear_folder(folder_path):\n    \"\"\"Clear all contents recursively if the folder exists.\n    Create the folder if it has been accidently deleted.\n    \"\"\"\n    create_folder(folder_path)\n    for the_file in os.listdir(folder_path):\n        _file_path = os.path.join(folder_path, the_file)\n        try:\n            if os.path.isfile(_file_path):\n                os.unlink(_file_path)\n            elif os.path.isdir(_file_path):\n                shutil.rmtree(_file_path)\n        except OSError as _e:\n            print(_e)\n\ndef create_folder(folder_path):\n    \"\"\"Create a folder if it does not exist.\n    \"\"\"\n    try:\n        os.makedirs(folder_path)\n    except OSError as _e:\n        if _e.errno != errno.EEXIST:\n            raise\n","metadata":{"id":"n6w3ypOvXaei","execution":{"iopub.status.busy":"2022-02-24T11:06:22.067182Z","iopub.execute_input":"2022-02-24T11:06:22.067585Z","iopub.status.idle":"2022-02-24T11:06:22.077664Z","shell.execute_reply.started":"2022-02-24T11:06:22.067528Z","shell.execute_reply":"2022-02-24T11:06:22.076909Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## ImageBuffer\n\nAs suggested in the paper, we update the discriminators by randomly picking an\nimage from the history of generated images, rather than the fake samples in real-time.\nThe history of generated images is maintained by the ImageBuffer class, which is defined\nas follows","metadata":{"id":"p-PvjeTDg7TH"}},{"cell_type":"code","source":"class ImageBuffer(object):\n    def __init__(self, depth=50):\n        self.depth = depth\n        self.buffer = []\n\n    def update(self, image):\n        # print(f'Image shape = {image.shape}')\n        if len(self.buffer) == self.depth:\n            i = random.randint(0, self.depth-1)\n            self.buffer[i] = image\n        else:\n            self.buffer.append(image)\n        if random.uniform(0,1) > 0.5:\n            i = random.randint(0, len(self.buffer)-1)\n            return self.buffer[i]\n        else:\n            return image","metadata":{"id":"ludZoo1Ag-sL","execution":{"iopub.status.busy":"2022-02-24T11:06:22.080016Z","iopub.execute_input":"2022-02-24T11:06:22.080275Z","iopub.status.idle":"2022-02-24T11:06:22.087301Z","shell.execute_reply.started":"2022-02-24T11:06:22.080241Z","shell.execute_reply":"2022-02-24T11:06:22.086297Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Images Dataset\nA custom dataset reader that picks up unpaired images from separate\nfolders.","metadata":{"id":"dWkkC-WVemj9"}},{"cell_type":"code","source":"class ImageDataset(Dataset):\n    def __init__(self, transform=None, unaligned=True, \n                 batch_size=0, max_size=0):\n        dir_A = f'{PARAMS[\"data_dir\"]}/{PARAMS[\"datasetA\"]}'\n        dir_B = f'{PARAMS[\"data_dir\"]}/{PARAMS[\"datasetB\"]}'\n        self.transform = tv.transforms.Compose(transform)\n        self.unaligned = unaligned\n        self.files_A = sorted(glob.glob(dir_A + '/*.*'))\n        self.files_B = sorted(glob.glob(dir_B + '/*.*'))\n        self.batch_size = batch_size\n        self.max_size = max_size \n\n    def __getitem__(self, index):\n        # Image.open(self.files_A[index % len(self.files_A)])\n        item_A = self.transform(Image.open(self.files_A[index % len(self.files_A)]))\n\n        if self.unaligned:\n            item_B = self.transform(Image.open(self.files_B[random.randint(0, len(self.files_B) - 1)]))\n        else:\n            item_B = self.transform(Image.open(self.files_B[index % len(self.files_B)]))\n\n        return {'datasetA': item_A, 'datasetB': item_B}\n\n    def __len__(self):\n        res = max(len(self.files_A), len(self.files_B))\n        if self.max_size > 0:\n            res = min(res, self.max_size)\n        # Обрезка по размеру батча (если задано), чтобы размеры сходились\n        if self.batch_size > 1:\n            res = res - res % self.batch_size\n        return res","metadata":{"id":"9KdvRVvghvBP","execution":{"iopub.status.busy":"2022-02-24T11:06:22.089009Z","iopub.execute_input":"2022-02-24T11:06:22.089499Z","iopub.status.idle":"2022-02-24T11:06:22.101614Z","shell.execute_reply.started":"2022-02-24T11:06:22.089464Z","shell.execute_reply":"2022-02-24T11:06:22.100867Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Image Transform","metadata":{"id":"JiI203HDwP-J"}},{"cell_type":"code","source":"transform = [\n#     tv.transforms.Resize(int(PARAMS['img_size']*1.12), Image.BICUBIC),\n#     tv.transforms.RandomCrop((PARAMS['img_size'], PARAMS['img_size'])),\n    tv.transforms.RandomHorizontalFlip(),\n    tv.transforms.ToTensor(),\n    tv.transforms.Normalize((0.5,0.5,0.5), (0.5,0.5,0.5))\n    ]\n","metadata":{"id":"sX7mXn2L-OPX","outputId":"af1127e8-7413-4de7-cd69-4dcfb68133d1","execution":{"iopub.status.busy":"2022-02-24T11:06:22.102856Z","iopub.execute_input":"2022-02-24T11:06:22.103651Z","iopub.status.idle":"2022-02-24T11:06:22.112444Z","shell.execute_reply.started":"2022-02-24T11:06:22.103614Z","shell.execute_reply":"2022-02-24T11:06:22.111705Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## DataLoaders","metadata":{"id":"uR3c3RNWxgOE"}},{"cell_type":"code","source":"train_loader = DataLoader(\n    ImageDataset(transform=transform, \n                 batch_size = PARAMS['batch_size'],\n                 max_size=350\n                ),\n    batch_size = PARAMS['batch_size'],\n    shuffle = True,\n    num_workers = 2,\n    pin_memory=True\n)\n\ntest_loader = DataLoader(\n    ImageDataset(transform=transform,  \n                 batch_size = PARAMS['batch_size'],\n                 max_size=300\n                ),\n    batch_size = PARAMS['test_batch_size'],\n    shuffle = True,\n    num_workers = 2,\n)","metadata":{"id":"TKfEJX4u-OPa","execution":{"iopub.status.busy":"2022-02-24T11:06:22.114358Z","iopub.execute_input":"2022-02-24T11:06:22.114672Z","iopub.status.idle":"2022-02-24T11:06:22.336734Z","shell.execute_reply.started":"2022-02-24T11:06:22.114574Z","shell.execute_reply":"2022-02-24T11:06:22.336097Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Визуализация изображений","metadata":{"id":"B8Thupntpwta"}},{"cell_type":"code","source":"# load a batch\ndata = next(iter(train_loader))\nimgs_x = data['datasetA']\nimgs_y = data['datasetB']\n\n# visualize batch\nfig, axis = plt.subplots(2, 2, figsize=(10, 10))\n\nfor i in range(2):\n    axis[0, i].axis('off')\n    axis[0, i].imshow(np.transpose((imgs_x[i]+1)/2, (1, 2, 0)))\n    axis[0, i].set_title(f'Monet {i+1}')\n    axis[1, i].axis('off')\n    axis[1, i].imshow(np.transpose((imgs_y[i]+1)/2, (1, 2, 0)))\n    axis[1, i].set_title(f'Photo {i+1}')","metadata":{"id":"SCI1WuVY-OPd","outputId":"a096f1dc-b7a3-403e-85aa-8167fc00072f","execution":{"iopub.status.busy":"2022-02-24T11:06:22.338665Z","iopub.execute_input":"2022-02-24T11:06:22.339138Z","iopub.status.idle":"2022-02-24T11:06:25.996817Z","shell.execute_reply.started":"2022-02-24T11:06:22.339102Z","shell.execute_reply":"2022-02-24T11:06:25.996056Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Generator","metadata":{"id":"mp6naB4qSymH"}},{"cell_type":"code","source":"class ResidualBlock(torch.nn.Module):\n    def __init__(self, channels):\n        super(ResidualBlock, self).__init__()\n\n        block = [torch.nn.ReflectionPad2d(1),\n                 torch.nn.Conv2d(channels, channels, 3),\n                 torch.nn.InstanceNorm2d(channels),\n                 torch.nn.ReLU(inplace=True),\n                 torch.nn.ReflectionPad2d(1),\n                 torch.nn.Conv2d(channels, channels, 3),\n                 torch.nn.InstanceNorm2d(channels)]\n        self.block = torch.nn.Sequential(*block)\n\n    def forward(self, x):\n        return x + self.block(x)","metadata":{"id":"AAONeOpG0J_h","execution":{"iopub.status.busy":"2022-02-24T11:06:25.998686Z","iopub.execute_input":"2022-02-24T11:06:25.998909Z","iopub.status.idle":"2022-02-24T11:06:26.007614Z","shell.execute_reply.started":"2022-02-24T11:06:25.998877Z","shell.execute_reply":"2022-02-24T11:06:26.00666Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"class Generator(torch.nn.Module):\n    def __init__(self, channels, num_blocks=9):\n        super(Generator, self).__init__()\n        self.channels = channels\n\n        model = [torch.nn.ReflectionPad2d(3)]\n        model += self._create_layer(self.channels, 64, 7, stride=1, padding=0, transposed=False)\n        # downsampling\n        model += self._create_layer(64, 128, 3, stride=2, padding=1, transposed=False)\n        model += self._create_layer(128, 256, 3, stride=2, padding=1, transposed=False)\n        # residual blocks\n        model += [ResidualBlock(256) for _ in range(num_blocks)]\n        # upsampling\n        model += self._create_layer(256, 128, 3, stride=2, padding=1, transposed=True)\n        model += self._create_layer(128, 64, 3, stride=2, padding=1, transposed=True)\n        # output\n        model += [torch.nn.ReflectionPad2d(3),\n                  torch.nn.Conv2d(64, self.channels, 7),\n                  torch.nn.Tanh()]\n\n        self.model = torch.nn.Sequential(*model)\n\n    def _create_layer(self, size_in, size_out, kernel_size, stride=2, padding=1, transposed=False):\n        layers = []\n        if transposed:\n            layers.append(torch.nn.ConvTranspose2d(size_in, size_out, kernel_size, stride=stride, padding=padding, output_padding=1))\n        else:\n            layers.append(torch.nn.Conv2d(size_in, size_out, kernel_size, stride=stride, padding=padding))\n        layers.append(torch.nn.InstanceNorm2d(size_out))\n        layers.append(torch.nn.ReLU(inplace=True))\n        return layers\n\n    def forward(self, x):\n        return self.model(x)\n","metadata":{"id":"7JEnOGAs0iI1","execution":{"iopub.status.busy":"2022-02-24T11:06:26.008995Z","iopub.execute_input":"2022-02-24T11:06:26.009575Z","iopub.status.idle":"2022-02-24T11:06:26.026166Z","shell.execute_reply.started":"2022-02-24T11:06:26.009523Z","shell.execute_reply":"2022-02-24T11:06:26.025374Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Discriminator\n\nPatchGAN architecture","metadata":{"id":"tpdtHcqfSuGW"}},{"cell_type":"code","source":"class Discriminator(torch.nn.Module):\n    def __init__(self, channels=3):\n        super(Discriminator, self).__init__()\n        self.channels = channels\n\n        self.model = torch.nn.Sequential(\n            *self._create_layer(self.channels, 64, 2, normalize=False),\n            *self._create_layer(64, 128, 2),\n            *self._create_layer(128, 256, 2),\n            *self._create_layer(256, 512, 1),\n            torch.nn.Conv2d(512, 1, 4, stride=1, padding=1)\n        )\n\n    def _create_layer(self, size_in, size_out, stride, normalize=True):\n        layers = [torch.nn.Conv2d(size_in, size_out, 4, stride=stride, padding=1)]\n        if normalize:\n            layers.append(torch.nn.InstanceNorm2d(size_out))\n        layers.append(torch.nn.LeakyReLU(0.2, inplace=True))\n        return layers\n\n    def forward(self, x):\n        return self.model(x)","metadata":{"id":"OklcUydJd_GE","execution":{"iopub.status.busy":"2022-02-24T11:06:26.027659Z","iopub.execute_input":"2022-02-24T11:06:26.028285Z","iopub.status.idle":"2022-02-24T11:06:26.040712Z","shell.execute_reply.started":"2022-02-24T11:06:26.02825Z","shell.execute_reply":"2022-02-24T11:06:26.039744Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Model training and evaluation","metadata":{"id":"M4Ihil561_YH"}},{"cell_type":"code","source":"# Function to initialize the weights\ndef _weights_init(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    elif classname.find('BatchNorm') != -1:\n        torch.nn.init.normal_(m.weight.data, 1.0, 0.02)\n        torch.nn.init.constant_(m.bias.data, 0.0)","metadata":{"id":"IIVBUu-K5TXt","execution":{"iopub.status.busy":"2022-02-24T11:06:26.043034Z","iopub.execute_input":"2022-02-24T11:06:26.043344Z","iopub.status.idle":"2022-02-24T11:06:26.052327Z","shell.execute_reply.started":"2022-02-24T11:06:26.043305Z","shell.execute_reply":"2022-02-24T11:06:26.051438Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"# Trainer\nclass Model(object):\n    def __init__(self,\n                 device,\n                 data_loader,\n                 test_data_loader,\n                 channels,\n                 img_size,\n                 num_blocks):\n        self.name = 'CycleGAN'\n        self.device = device\n        self.data_loader = data_loader\n        self.test_data_loader = test_data_loader\n        self.channels = channels\n        self.img_size = img_size\n        self.num_blocks = num_blocks\n        self.netG_AB = Generator(self.channels, self.num_blocks)\n        self.netG_AB.apply(_weights_init)\n        self.netG_AB.to(self.device)\n        self.netG_BA = Generator(self.channels, self.num_blocks)\n        self.netG_BA.apply(_weights_init)\n        self.netG_BA.to(self.device)\n        self.netD_A = Discriminator(self.channels)\n        self.netD_A.apply(_weights_init)\n        self.netD_A.to(self.device)\n        self.netD_B = Discriminator(self.channels)\n        self.netD_B.apply(_weights_init)\n        self.netD_B.to(self.device)\n\n        self.optim_G = None\n        self.optim_D_A = None\n        self.optim_D_B = None\n\n        self.loss_adv = torch.nn.MSELoss()\n        self.loss_cyc = torch.nn.L1Loss()\n        self.loss_iden = torch.nn.L1Loss()\n\n    @property\n    def generator_AB(self):\n        return self.netG_AB\n\n    @property\n    def generator_BA(self):\n        return self.netG_BA\n\n    @property\n    def discriminator_A(self):\n        return self.netD_A\n\n    @property\n    def discriminator_B(self):\n        return self.netD_B\n\n    def create_optim(self, lr, alpha=0.5, beta=0.999):\n        self.optim_G = torch.optim.Adam(itertools.chain(self.netG_AB.parameters(),\n                                                        self.netG_BA.parameters()),\n                                        lr=lr,\n                                        betas=(alpha, beta))\n        self.optim_D_A = torch.optim.Adam(self.netD_A.parameters(),\n                                          lr=lr/2,\n                                          betas=(alpha, beta))\n        self.optim_D_B = torch.optim.Adam(self.netD_B.parameters(),\n                                          lr=lr/2,\n                                          betas=(alpha, beta))\n\n    def train(self,\n              epochs,\n              log_interval=100,\n              out_dir='',\n              verbose=True):\n        self.netG_AB.train()\n        self.netG_BA.train()\n        self.netD_A.train()\n        self.netD_B.train()\n        lambda_cyc = 10\n        lambda_iden = 5\n        real_label = torch.ones((self.data_loader.batch_size, 1, 30, 30), device=self.device)\n        fake_label = torch.zeros((self.data_loader.batch_size, 1, 30, 30), device=self.device)\n        image_buffer_A = ImageBuffer()\n        image_buffer_B = ImageBuffer()\n        total_time = time.time()\n\n        # Scheduler для уменьшения шага скорости обучения\n        decay_epoch = 110\n\n        lambda_func = lambda epoch: 1 - max(0, epoch-decay_epoch)/(epochs-decay_epoch)\n        lr_scheduler_G = torch.optim.lr_scheduler.LambdaLR(self.optim_G, lr_lambda=lambda_func)\n        lr_scheduler_D_A = torch.optim.lr_scheduler.LambdaLR(self.optim_D_A, lr_lambda=lambda_func)\n        lr_scheduler_D_B = torch.optim.lr_scheduler.LambdaLR(self.optim_D_B, lr_lambda=lambda_func)\n#         lr_scheduler_G = torch.optim.lr_scheduler.StepLR(self.optim_G, step_size=1, gamma=0.9)\n#         lr_scheduler_D_A = torch.optim.lr_scheduler.StepLR(self.optim_D_A, step_size=1, gamma=0.9)\n#         lr_scheduler_D_B = torch.optim.lr_scheduler.StepLR(self.optim_D_B, step_size=1, gamma=0.9)\n\n        for epoch in range(epochs):\n            batch_time = time.time()\n            print(f'Epoch: {epoch}, lr = {lr_scheduler_G.get_last_lr()}')\n            for batch_idx, data in enumerate(self.data_loader):\n                real_A = data['datasetA'].to(self.device)\n                real_B = data['datasetB'].to(self.device)\n\n                # Train G\n                self.optim_G.zero_grad()\n\n                # adversarial loss\n                fake_B = self.netG_AB(real_A)\n                _loss_adv_AB = self.loss_adv(self.netD_B(fake_B), real_label)\n                fake_A = self.netG_BA(real_B)\n                _loss_adv_BA = self.loss_adv(self.netD_A(fake_A), real_label)\n                adv_loss = (_loss_adv_AB + _loss_adv_BA) / 2\n\n                # cycle loss\n                recov_A = self.netG_BA(fake_B)\n                _loss_cyc_A = self.loss_cyc(recov_A, real_A)\n                recov_B = self.netG_AB(fake_A)\n                _loss_cyc_B = self.loss_cyc(recov_B, real_B)\n                cycle_loss = (_loss_cyc_A + _loss_cyc_B) / 2\n\n                # identity loss\n                _loss_iden_A = self.loss_iden(self.netG_BA(real_A), real_A)\n                _loss_iden_B = self.loss_iden(self.netG_AB(real_B), real_B)\n                iden_loss = (_loss_iden_A + _loss_iden_B) / 2\n\n                g_loss = adv_loss + lambda_cyc * cycle_loss + lambda_iden * iden_loss\n                g_loss.backward()\n                self.optim_G.step()\n\n                # Train D_A\n                self.optim_D_A.zero_grad()\n\n                _loss_real = self.loss_adv(self.netD_A(real_A), real_label)\n#                 fake_A = image_buffer_A.update(fake_A)\n                _loss_fake = self.loss_adv(self.netD_A(fake_A.detach()), fake_label)\n                d_loss_A = (_loss_real + _loss_fake) / 2\n\n                d_loss_A.backward()\n                self.optim_D_A.step()\n\n                # Train D_B\n                self.optim_D_B.zero_grad()\n\n                _loss_real = self.loss_adv(self.netD_B(real_B), real_label)\n#                 fake_B = image_buffer_B.update(fake_B)\n                _loss_fake = self.loss_adv(self.netD_B(fake_B.detach()), fake_label)\n                d_loss_B = (_loss_real + _loss_fake) / 2\n\n                d_loss_B.backward()\n                self.optim_D_B.step()\n\n                d_loss = (d_loss_A + d_loss_B) / 2\n\n                if verbose and batch_idx % log_interval == 0 and batch_idx > 0:\n                    print('Epoch {} [{}/{}] loss_D: {:.4f} loss_G: {:.4f} time: {:.2f}'.format(\n                          epoch, batch_idx, len(self.data_loader),\n                          d_loss.mean().item(),\n                          g_loss.mean().item(),\n                          time.time() - batch_time))\n#                     with torch.no_grad():\n#                         imgs = next(iter(self.test_data_loader))\n#                         _real_A = imgs['datasetA'].to(self.device)\n#                         _fake_B = self.netG_AB(_real_A)\n#                         _real_B = imgs['datasetB'].to(self.device)\n#                         _fake_A = self.netG_BA(_real_B)\n#                         viz_sample = torch.cat((_real_A, _fake_B, _real_B, _fake_A), 0)\n#                         vutils.save_image(viz_sample,\n#                                           os.path.join(out_dir, 'samples_{}_{}.png'.format(epoch, batch_idx)),\n#                                           nrow=self.test_data_loader.batch_size,\n#                                           normalize=True)\n                    batch_time = time.time()\n            lr_scheduler_G.step()\n            lr_scheduler_D_A.step()\n            lr_scheduler_D_B.step()\n\n#             self.save_to(path=out_dir, name=self.name, verbose=False)\n        if verbose:\n            print('Total train time: {:.2f}'.format(time.time() - total_time))\n\n    def eval(self,\n             batch_size=None):\n        self.netG_AB.eval()\n        self.netG_BA.eval()\n        self.netD_A.eval()\n        self.netD_B.eval()\n        if batch_size is None:\n            batch_size = self.test_data_loader.batch_size\n\n        with torch.no_grad():\n            for batch_idx, data in enumerate(self.test_data_loader):\n                _real_A = data['testA'].to(self.device)\n                _fake_B = self.netG_AB(_real_A)\n                _real_B = data['testB'].to(self.device)\n                _fake_A = self.netG_BA(_real_B)\n                viz_sample = torch.cat((_real_A, _fake_B, _real_B, _fake_A), 0)\n                vutils.save_image(viz_sample,\n                                  'img_{}.png'.format(batch_idx),\n                                  nrow=batch_size,\n                                  normalize=True)\n\n    def save_to(self,\n                path='',\n                name=None,\n                verbose=True):\n        if name is None:\n            name = self.name\n        if verbose:\n            print('\\nSaving models to {}_G_AB.pt and such ...'.format(name))\n        torch.save(self.netG_AB.state_dict(), os.path.join(path, '{}_G_AB.pt'.format(name)))\n        torch.save(self.netG_BA.state_dict(), os.path.join(path, '{}_G_BA.pt'.format(name)))\n        torch.save(self.netD_A.state_dict(), os.path.join(path, '{}_D_A.pt'.format(name)))\n        torch.save(self.netD_B.state_dict(), os.path.join(path, '{}_D_B.pt'.format(name)))\n\n    def load_from(self,\n                  path='',\n                  name=None,\n                  verbose=True):\n        if name is None:\n            name = self.name\n        if verbose:\n            print('\\nLoading models from {}_G_AB.pt and such ...'.format(name))\n        ckpt_G_AB = torch.load(os.path.join(path, '{}_G_AB.pt'.format(name)))\n        if isinstance(ckpt_G_AB, dict) and 'state_dict' in ckpt_G_AB:\n            self.netG_AB.load_state_dict(ckpt_G_AB['state_dict'], strict=True)\n        else:\n            self.netG_AB.load_state_dict(ckpt_G_AB, strict=True)\n        ckpt_G_BA = torch.load(os.path.join(path, '{}_G_BA.pt'.format(name)))\n        if isinstance(ckpt_G_BA, dict) and 'state_dict' in ckpt_G_BA:\n            self.netG_BA.load_state_dict(ckpt_G_BA['state_dict'], strict=True)\n        else:\n            self.netG_BA.load_state_dict(ckpt_G_BA, strict=True)\n        ckpt_D_A = torch.load(os.path.join(path, '{}_D_A.pt'.format(name)))\n        if isinstance(ckpt_D_A, dict) and 'state_dict' in ckpt_D_A:\n            self.netD_A.load_state_dict(ckpt_D_A['state_dict'], strict=True)\n        else:\n            self.netD_A.load_state_dict(ckpt_D_A, strict=True)\n        ckpt_D_B = torch.load(os.path.join(path, '{}_D_B.pt'.format(name)))\n        if isinstance(ckpt_D_B, dict) and 'state_dict' in ckpt_D_B:\n            self.netD_B.load_state_dict(ckpt_D_B['state_dict'], strict=True)\n        else:\n            self.netD_B.load_state_dict(ckpt_D_B, strict=True)\n","metadata":{"id":"l-RRVzfj6N9d","execution":{"iopub.status.busy":"2022-02-24T11:06:26.054323Z","iopub.execute_input":"2022-02-24T11:06:26.054743Z","iopub.status.idle":"2022-02-24T11:06:26.099402Z","shell.execute_reply.started":"2022-02-24T11:06:26.054707Z","shell.execute_reply":"2022-02-24T11:06:26.098696Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Running model","metadata":{"id":"CG2FBG2xWbJJ"}},{"cell_type":"code","source":"PARAMS['cuda'] = torch.cuda.is_available()\n\nif PARAMS['seed'] is not None:\n    torch.manual_seed(PARAMS['seed'])\n    if PARAMS['cuda']:\n        torch.cuda.manual_seed(PARAMS['seed'])\n    np.random.seed(PARAMS['seed'])\n\ncudnn.benchmark = True\n\nclear_folder(PARAMS['out_dir'])\n\nlog_file = os.path.join(PARAMS['out_dir'], 'log.txt')\nprint(\"Logging to {}\\n\".format(log_file))\n# sys.stdout = StdOut(log_file)\n\nprint(\"PyTorch version: {}\".format(torch.__version__))\nprint(\"CUDA version: {}\\n\".format(torch.version.cuda))\n\nfor key in PARAMS.keys():\n    print(f'{key} = {PARAMS[key]}')","metadata":{"id":"TRrzOdn0-OP5","outputId":"257e8053-d5dc-447f-a629-1c327b789c36","execution":{"iopub.status.busy":"2022-02-24T11:06:26.101787Z","iopub.execute_input":"2022-02-24T11:06:26.102254Z","iopub.status.idle":"2022-02-24T11:06:26.115296Z","shell.execute_reply.started":"2022-02-24T11:06:26.10222Z","shell.execute_reply":"2022-02-24T11:06:26.114446Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"print('Creating model...\\n')\nmodel = Model(device, train_loader, test_loader, PARAMS['channels'], PARAMS['img_size'], PARAMS['num_blocks'])\nmodel.create_optim(PARAMS['lr'])","metadata":{"id":"u1JKPyTedLKg","outputId":"84aa63c4-b169-47ea-df03-be5f7a6e3513","execution":{"iopub.status.busy":"2022-02-24T11:06:26.116546Z","iopub.execute_input":"2022-02-24T11:06:26.116954Z","iopub.status.idle":"2022-02-24T11:06:26.604425Z","shell.execute_reply.started":"2022-02-24T11:06:26.116917Z","shell.execute_reply":"2022-02-24T11:06:26.603679Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%%time\n# Train\nmodel.train(PARAMS['epochs'], PARAMS['log_interval'], PARAMS['out_dir'], True)\n\n# model.save_to('')","metadata":{"id":"fQzavOGbYf3l","outputId":"d0844c90-6f43-4b52-c34a-420f7e6c6483","execution":{"iopub.status.busy":"2022-02-24T11:06:45.62815Z","iopub.execute_input":"2022-02-24T11:06:45.628933Z","iopub.status.idle":"2022-02-24T11:09:53.818816Z","shell.execute_reply.started":"2022-02-24T11:06:45.628898Z","shell.execute_reply":"2022-02-24T11:09:53.817944Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Визуализация результатов","metadata":{"id":"_zXrGabzYGGR"}},{"cell_type":"code","source":"from torch.functional import Tensor\ndef sample_images(real_A, real_B, model, figside=5):\n    assert real_A.size() == real_B.size(), 'The image size for two domains must be the same'\n    \n    netG = model.generator_AB\n    netF = model.generator_BA\n\n    netG.eval()\n    netF.eval()\n    \n    real_A = real_A.type(Tensor).cuda()\n    fake_B = netG(real_A)\n    real_B = real_B.type(Tensor).cuda()\n    fake_A = netF(real_B)\n    \n    nrows = real_A.size(0)\n    real_A = make_grid(real_A, nrow=nrows, normalize=True)\n    fake_B = make_grid(fake_B, nrow=nrows, normalize=True)\n    real_B = make_grid(real_B, nrow=nrows, normalize=True)\n    fake_A = make_grid(fake_A, nrow=nrows, normalize=True)\n    \n    image_grid = torch.cat((real_A, fake_B, real_B, fake_A), 1).cpu().permute(1, 2, 0)\n    \n    plt.figure(figsize=(figside*nrows, figside*4))\n    plt.imshow(image_grid)\n    plt.axis('off')\n    plt.show()","metadata":{"id":"4_AstPkOcis4","execution":{"iopub.status.busy":"2022-02-22T13:25:40.013128Z","iopub.execute_input":"2022-02-22T13:25:40.013873Z","iopub.status.idle":"2022-02-22T13:25:40.023497Z","shell.execute_reply.started":"2022-02-22T13:25:40.013829Z","shell.execute_reply":"2022-02-22T13:25:40.022779Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"data = next(iter(test_loader))\nreal_A, real_B = data['datasetA'], data['datasetB']\nsample_images(real_A, real_B, model)","metadata":{"id":"uhiUpUf2dD5a","outputId":"5875d704-1d24-44fc-ce2b-36fe126176d7","execution":{"iopub.status.busy":"2022-02-22T13:25:45.520025Z","iopub.execute_input":"2022-02-22T13:25:45.520459Z","iopub.status.idle":"2022-02-22T13:25:47.048656Z","shell.execute_reply.started":"2022-02-22T13:25:45.520423Z","shell.execute_reply":"2022-02-22T13:25:47.047881Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## Generate Images","metadata":{}},{"cell_type":"code","source":"photo_dir = os.path.join(PARAMS['data_dir'], PARAMS['datasetB'])\nfiles = [os.path.join(photo_dir, name) for name in os.listdir(photo_dir)]\nlen(files)","metadata":{"execution":{"iopub.status.busy":"2022-02-22T13:26:41.922529Z","iopub.execute_input":"2022-02-22T13:26:41.92279Z","iopub.status.idle":"2022-02-22T13:26:41.945855Z","shell.execute_reply.started":"2022-02-22T13:26:41.92276Z","shell.execute_reply":"2022-02-22T13:26:41.945107Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"save_dir = '../images'\nif not os.path.exists(save_dir):\n    os.makedirs(save_dir)","metadata":{"execution":{"iopub.status.busy":"2022-02-22T13:26:44.424489Z","iopub.execute_input":"2022-02-22T13:26:44.424734Z","iopub.status.idle":"2022-02-22T13:26:44.431108Z","shell.execute_reply.started":"2022-02-22T13:26:44.424706Z","shell.execute_reply":"2022-02-22T13:26:44.428483Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"generate_transforms = tv.transforms.Compose([\n    tv.transforms.ToTensor(),\n    tv.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n])\n\nto_image = tv.transforms.ToPILImage()\nbatch_size = PARAMS['batch_size']\nmodel.generator_BA.eval()\nfor i in range(0, len(files), batch_size):\n    # read images\n    imgs = []\n    for j in range(i, min(len(files), i+batch_size)):\n        img = Image.open(files[j])\n        img = generate_transforms(img)\n        imgs.append(img)\n    imgs = torch.stack(imgs, 0).type(Tensor)\n    \n    # generate\n    fake_imgs = model.generator_BA(imgs.to(device)).detach().cpu()\n    \n    # save\n    for j in range(fake_imgs.size(0)):\n        img = fake_imgs[j].squeeze().permute(1, 2, 0)\n        img_arr = img.numpy()\n        img_arr = (img_arr - np.min(img_arr)) * 255 / (np.max(img_arr) - np.min(img_arr))\n        img_arr = img_arr.astype(np.uint8)\n        \n        img = to_image(img_arr)\n        _, name = os.path.split(files[i+j])\n        img.save(os.path.join(save_dir, name))","metadata":{"execution":{"iopub.status.busy":"2022-02-22T13:26:53.536472Z","iopub.execute_input":"2022-02-22T13:26:53.536726Z","iopub.status.idle":"2022-02-22T13:29:16.731995Z","shell.execute_reply.started":"2022-02-22T13:26:53.536691Z","shell.execute_reply":"2022-02-22T13:29:16.731251Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"shutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/images\")","metadata":{"execution":{"iopub.status.busy":"2022-02-22T13:29:50.494602Z","iopub.execute_input":"2022-02-22T13:29:50.494883Z","iopub.status.idle":"2022-02-22T13:29:54.440072Z","shell.execute_reply.started":"2022-02-22T13:29:50.494852Z","shell.execute_reply":"2022-02-22T13:29:54.439334Z"},"trusted":true},"execution_count":null,"outputs":[]}]}