{"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 implementation with pytorch_lightning","metadata":{}},{"cell_type":"markdown","source":"## 1. Build ImageFodler Dataset ","metadata":{"execution":{"iopub.status.busy":"2021-06-17T05:24:20.680051Z","iopub.execute_input":"2021-06-17T05:24:20.680463Z","iopub.status.idle":"2021-06-17T05:24:20.686988Z","shell.execute_reply.started":"2021-06-17T05:24:20.680382Z","shell.execute_reply":"2021-06-17T05:24:20.686085Z"}}},{"cell_type":"code","source":"import functools\nimport itertools\nfrom pathlib import Path\nfrom typing import List\nimport numpy as np\nfrom torch import nn\nimport torch.utils.data as data\nimport torch.nn.functional as F\nfrom PIL import Image\nimport os\nimport os.path\nimport shutil\n\nimport torchvision\n\nIMG_EXTENSIONS = [\n    '.jpg', '.JPG', '.jpeg', '.JPEG',\n    '.png', '.PNG', '.ppm', '.PPM', '.bmp', '.BMP',\n    '.tif', '.TIF', '.tiff', '.TIFF',\n]\n\n\ndef is_image_file(filename):\n    return any(filename.endswith(extension) for extension in IMG_EXTENSIONS)\n\n\ndef make_dataset(dir, max_dataset_size=float(\"inf\")):\n    images = []\n    assert os.path.isdir(dir) or os.path.islink(dir), '%s is not a valid directory' % dir\n\n    for root, _, fnames in sorted(os.walk(dir, followlinks=True)):\n        for fname in fnames:\n            if is_image_file(fname):\n                path = os.path.join(root, fname)\n                images.append(path)\n    return images[:min(max_dataset_size, len(images))]\n\n\ndef default_loader(path):\n    return Image.open(path).convert('RGB')\n\n\nclass ImageFolder(data.Dataset):\n\n    def __init__(self, root, transform=None, return_paths=False,\n                 loader=default_loader, max_dataset_size=float('inf')):\n        imgs = make_dataset(root, max_dataset_size=max_dataset_size)\n        if len(imgs) == 0:\n            raise(RuntimeError(\"Found 0 images in: \" + root + \"\\n\"\n                               \"Supported image extensions are: \" + \",\".join(IMG_EXTENSIONS)))\n\n        self.root = root\n        self.imgs = imgs\n        self.transform = transform\n        self.return_paths = return_paths\n        self.loader = loader\n\n    def __getitem__(self, index):\n        path = self.imgs[index]\n        img = self.loader(path)\n        if self.transform is not None:\n            img = self.transform(img)\n        if self.return_paths:\n            return img, path\n        else:\n            return img\n\n    def __len__(self):\n        return len(self.imgs)\n\nclass MergeDataset(data.Dataset):\n    def __init__(self, *datasets):\n        \"\"\"Merge multiple datasets to one dataset, and each time retrives a combinations of items in all sub datasets.\n        \"\"\"\n        self.datasets = datasets \n        self.sizes = [len(dataset) for dataset in datasets]\n        print('dataset size', self.sizes)\n\n    def __getitem__(self, indexs: List[int]):\n        return tuple(dataset[idx] for idx, dataset in zip(indexs, self.datasets))\n\n    def __len__(self):\n        return max(self.sizes)\n\nclass DistributedSamplerWrapper(data.DistributedSampler):\n    def __init__(\n            self, sampler,\n            num_replicas = None,\n            rank = None,\n            shuffle = True):\n        super(DistributedSamplerWrapper, self).__init__(\n            sampler.data_source, num_replicas, rank, shuffle)\n        self.sampler = sampler\n\n    def __iter__(self):\n        indices = list(self.sampler)\n        indices = indices[self.rank:self.total_size:self.num_replicas]\n        return iter(indices)\n\n    def __len__(self):\n        return len(self.sampler) // self.num_replicas\n    \nclass MultiRandomSampler(data.RandomSampler):\n    \"\"\" a Random Sampler for MergeDataset. NOTE will padding all dataset to same length\n    Each time it generates an index for each subdataset in MergeDataset.\n\n    Args:\n        data_source (MergeDataset): MergeDataset object\n        replacement (bool, optional): shuffle index use replacement. Defaults to True.\n        num_samples ([type], optional): Defaults to None.\n        generator ([type], optional): Defaults to None.\n    \"\"\"\n    def __init__(self,\n                 data_source: MergeDataset,\n                 replacement=True,\n                 num_samples=None,\n                 generator=None):\n        self.data_source: MergeDataset = data_source\n        self.replacement = replacement\n        self._num_samples = num_samples\n        self.generator = generator\n        self.maxn = len(self.data_source)\n\n    @property\n    def num_samples(self):\n        # dataset size might change at runtime\n        if self._num_samples is None:\n            self._num_samples = self.data_source.sizes\n        return self._num_samples\n\n    def __iter__(self):\n        rands = []\n        for size in self.num_samples:\n            if self.maxn == size:\n                rands.append(\n                    torch.randperm(size, generator=self.generator).tolist())\n            else:\n                rands.append(\n                    torch.randint(high=size,\n                                  size=(self.maxn, ),\n                                  dtype=torch.int64,\n                                  generator=self.generator).tolist())\n        return zip(*rands)\n\n    def __len__(self):\n        return len(self.data_source)\n","metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","execution":{"iopub.status.busy":"2021-06-22T10:58:02.439414Z","iopub.execute_input":"2021-06-22T10:58:02.439771Z","iopub.status.idle":"2021-06-22T10:58:02.46513Z","shell.execute_reply.started":"2021-06-22T10:58:02.439738Z","shell.execute_reply":"2021-06-22T10:58:02.464236Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 2. Build LinghtningDataModule","metadata":{}},{"cell_type":"code","source":"import pytorch_lightning as pl\nfrom torchvision import transforms\nimport torch\n\nclass Photo2MonetDataset(pl.LightningDataModule):\n    def __init__(self, batch_size=64, use_tpu=False, num_workers=8):\n        super().__init__()\n        self.batch_size = batch_size\n        self.use_tpu = use_tpu\n        self.num_workers = num_workers\n\n        transform = transforms.Compose([\n            transforms.ToTensor(),\n            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))\n        ])\n        train_photo = ImageFolder('/kaggle/input/gan-getting-started/photo_jpg', transform=transform)\n        train_paint = ImageFolder('/kaggle/input/gan-getting-started/monet_jpg', transform=transform)\n        self.train_data = MergeDataset(train_photo, train_paint)\n        self.test_data = ImageFolder('/kaggle/input/gan-getting-started/photo_jpg', transform=transform, return_paths=True)\n    \n    def train_dataloader(self):\n        sampler = MultiRandomSampler(self.train_data)\n        train_loader = data.DataLoader(self.train_data, batch_size=self.batch_size, sampler=sampler, num_workers=self.num_workers)\n        return train_loader\n    \n    def test_dataloader(self):\n        return data.DataLoader(self.test_data, batch_size=self.batch_size, num_workers=self.num_workers)\n","metadata":{"execution":{"iopub.status.busy":"2021-06-22T10:58:03.954015Z","iopub.execute_input":"2021-06-22T10:58:03.954347Z","iopub.status.idle":"2021-06-22T10:58:05.378411Z","shell.execute_reply.started":"2021-06-22T10:58:03.954304Z","shell.execute_reply":"2021-06-22T10:58:05.377474Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 3. Build Generator and Discriminator","metadata":{}},{"cell_type":"code","source":"################################################# Model #############################\n\nclass UnetGenerator(nn.Module):\n    \"\"\"Create a Unet-based generator\"\"\"\n\n    def __init__(self, input_nc, output_nc, num_downs, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False):\n        \"\"\"Construct a Unet generator\n        Parameters:\n            input_nc (int)  -- the number of channels in input images\n            output_nc (int) -- the number of channels in output images\n            num_downs (int) -- the number of downsamplings in UNet. For example, # if |num_downs| == 7,\n                                image of size 128x128 will become of size 1x1 # at the bottleneck\n            ngf (int)       -- the number of filters in the last conv layer\n            norm_layer      -- normalization layer\n\n        We construct the U-Net from the innermost layer to the outermost layer.\n        It is a recursive process.\n        \"\"\"\n        super(UnetGenerator, self).__init__()\n        # construct unet structure\n        unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=None, norm_layer=norm_layer, innermost=True)  # add the innermost layer\n        for i in range(num_downs - 5):          # add intermediate layers with ngf * 8 filters\n            unet_block = UnetSkipConnectionBlock(ngf * 8, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer, use_dropout=use_dropout)\n        # gradually reduce the number of filters from ngf * 8 to ngf\n        unet_block = UnetSkipConnectionBlock(ngf * 4, ngf * 8, input_nc=None, submodule=unet_block, norm_layer=norm_layer)\n        unet_block = UnetSkipConnectionBlock(ngf * 2, ngf * 4, input_nc=None, submodule=unet_block, norm_layer=norm_layer)\n        unet_block = UnetSkipConnectionBlock(ngf, ngf * 2, input_nc=None, submodule=unet_block, norm_layer=norm_layer)\n        self.model = UnetSkipConnectionBlock(output_nc, ngf, input_nc=input_nc, submodule=unet_block, outermost=True, norm_layer=norm_layer)  # add the outermost layer\n\n    def forward(self, input):\n        \"\"\"Standard forward\"\"\"\n        return self.model(input)\n\n\nclass UnetSkipConnectionBlock(nn.Module):\n    \"\"\"Defines the Unet submodule with skip connection.\n        X -------------------identity----------------------\n        |-- downsampling -- |submodule| -- upsampling --|\n    \"\"\"\n\n    def __init__(self, outer_nc, inner_nc, input_nc=None,\n                 submodule=None, outermost=False, innermost=False, norm_layer=nn.BatchNorm2d, use_dropout=False):\n        \"\"\"Construct a Unet submodule with skip connections.\n\n        Parameters:\n            outer_nc (int) -- the number of filters in the outer conv layer\n            inner_nc (int) -- the number of filters in the inner conv layer\n            input_nc (int) -- the number of channels in input images/features\n            submodule (UnetSkipConnectionBlock) -- previously defined submodules\n            outermost (bool)    -- if this module is the outermost module\n            innermost (bool)    -- if this module is the innermost module\n            norm_layer          -- normalization layer\n            use_dropout (bool)  -- if use dropout layers.\n        \"\"\"\n        super(UnetSkipConnectionBlock, self).__init__()\n        self.outermost = outermost\n        if type(norm_layer) == functools.partial:\n            use_bias = norm_layer.func == nn.InstanceNorm2d\n        else:\n            use_bias = norm_layer == nn.InstanceNorm2d\n        if input_nc is None:\n            input_nc = outer_nc\n        downconv = nn.Conv2d(input_nc, inner_nc, kernel_size=4,\n                             stride=2, padding=1, bias=use_bias)\n        downrelu = nn.LeakyReLU(0.2, True)\n        downnorm = norm_layer(inner_nc)\n        uprelu = nn.ReLU(True)\n        upnorm = norm_layer(outer_nc)\n\n        if outermost:\n            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,\n                                        kernel_size=4, stride=2,\n                                        padding=1)\n            down = [downconv]\n            up = [uprelu, upconv, nn.Tanh()]\n            model = down + [submodule] + up\n        elif innermost:\n            upconv = nn.ConvTranspose2d(inner_nc, outer_nc,\n                                        kernel_size=4, stride=2,\n                                        padding=1, bias=use_bias)\n            down = [downrelu, downconv]\n            up = [uprelu, upconv, upnorm]\n            model = down + up\n        else:\n            upconv = nn.ConvTranspose2d(inner_nc * 2, outer_nc,\n                                        kernel_size=4, stride=2,\n                                        padding=1, bias=use_bias)\n            down = [downrelu, downconv, downnorm]\n            up = [uprelu, upconv, upnorm]\n\n            if use_dropout:\n                model = down + [submodule] + up + [nn.Dropout(0.5)]\n            else:\n                model = down + [submodule] + up\n\n        self.model = nn.Sequential(*model)\n\n    def forward(self, x):\n        if self.outermost:\n            return self.model(x)\n        else:   # add skip connections\n            return torch.cat([x, self.model(x)], 1)\n\ndef get_filter(filt_size=3):\n    if(filt_size == 1):\n        a = np.array([1., ])\n    elif(filt_size == 2):\n        a = np.array([1., 1.])\n    elif(filt_size == 3):\n        a = np.array([1., 2., 1.])\n    elif(filt_size == 4):\n        a = np.array([1., 3., 3., 1.])\n    elif(filt_size == 5):\n        a = np.array([1., 4., 6., 4., 1.])\n    elif(filt_size == 6):\n        a = np.array([1., 5., 10., 10., 5., 1.])\n    elif(filt_size == 7):\n        a = np.array([1., 6., 15., 20., 15., 6., 1.])\n\n    filt = torch.Tensor(a[:, None] * a[None, :])\n    filt = filt / torch.sum(filt)\n\n    return filt\n\ndef get_pad_layer(pad_type):\n    if(pad_type in ['refl', 'reflect']):\n        PadLayer = nn.ReflectionPad2d\n    elif(pad_type in ['repl', 'replicate']):\n        PadLayer = nn.ReplicationPad2d\n    elif(pad_type == 'zero'):\n        PadLayer = nn.ZeroPad2d\n    else:\n        print('Pad type [%s] not recognized' % pad_type)\n    return PadLayer\n\nclass Downsample(nn.Module):\n    def __init__(self, channels, pad_type='reflect', filt_size=3, stride=2, pad_off=0):\n        super(Downsample, self).__init__()\n        self.filt_size = filt_size\n        self.pad_off = pad_off\n        self.pad_sizes = [int(1. * (filt_size - 1) / 2), int(np.ceil(1. * (filt_size - 1) / 2)), int(1. * (filt_size - 1) / 2), int(np.ceil(1. * (filt_size - 1) / 2))]\n        self.pad_sizes = [pad_size + pad_off for pad_size in self.pad_sizes]\n        self.stride = stride\n        self.off = int((self.stride - 1) / 2.)\n        self.channels = channels\n\n        filt = get_filter(filt_size=self.filt_size)\n        self.register_buffer('filt', filt[None, None, :, :].repeat((self.channels, 1, 1, 1)))\n\n        self.pad = get_pad_layer(pad_type)(self.pad_sizes)\n\n    def forward(self, inp):\n        if(self.filt_size == 1):\n            if(self.pad_off == 0):\n                return inp[:, :, ::self.stride, ::self.stride]\n            else:\n                return self.pad(inp)[:, :, ::self.stride, ::self.stride]\n        else:\n            return F.conv2d(self.pad(inp), self.filt, stride=self.stride, groups=inp.shape[1])\n\nclass Upsample(nn.Module):\n    def __init__(self, channels, pad_type='repl', filt_size=4, stride=2):\n        super(Upsample, self).__init__()\n        self.filt_size = filt_size\n        self.filt_odd = np.mod(filt_size, 2) == 1\n        self.pad_size = int((filt_size - 1) / 2)\n        self.stride = stride\n        self.off = int((self.stride - 1) / 2.)\n        self.channels = channels\n\n        filt = get_filter(filt_size=self.filt_size) * (stride**2)\n        self.register_buffer('filt', filt[None, None, :, :].repeat((self.channels, 1, 1, 1)))\n\n        self.pad = get_pad_layer(pad_type)([1, 1, 1, 1])\n\n    def forward(self, inp):\n        ret_val = F.conv_transpose2d(self.pad(inp), self.filt, stride=self.stride, padding=1 + self.pad_size, groups=inp.shape[1])[:, :, 1:, 1:]\n        if(self.filt_odd):\n            return ret_val\n        else:\n            return ret_val[:, :, :-1, :-1]\n\nclass ResnetGenerator(nn.Module):\n    \"\"\"Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.\n\n    We adapt Torch code and idea from Justin Johnson's neural style transfer project(https://github.com/jcjohnson/fast-neural-style)\n    \"\"\"\n\n    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect', no_antialias=False, no_antialias_up=False, opt=None):\n        \"\"\"Construct a Resnet-based generator\n\n        Parameters:\n            input_nc (int)      -- the number of channels in input images\n            output_nc (int)     -- the number of channels in output images\n            ngf (int)           -- the number of filters in the last conv layer\n            norm_layer          -- normalization layer\n            use_dropout (bool)  -- if use dropout layers\n            n_blocks (int)      -- the number of ResNet blocks\n            padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero\n        \"\"\"\n        assert(n_blocks >= 0)\n        super(ResnetGenerator, self).__init__()\n        self.opt = opt\n        if type(norm_layer) == functools.partial:\n            use_bias = norm_layer.func == nn.InstanceNorm2d\n        else:\n            use_bias = norm_layer == nn.InstanceNorm2d\n\n        model = [nn.ReflectionPad2d(3),\n                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),\n                 norm_layer(ngf),\n                 nn.ReLU(True)]\n\n        n_downsampling = 2\n        for i in range(n_downsampling):  # add downsampling layers\n            mult = 2 ** i\n            if(no_antialias):\n                model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),\n                          norm_layer(ngf * mult * 2),\n                          nn.ReLU(True)]\n            else:\n                model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=1, padding=1, bias=use_bias),\n                          norm_layer(ngf * mult * 2),\n                          nn.ReLU(True),\n                          Downsample(ngf * mult * 2)]\n\n        mult = 2 ** n_downsampling\n        for i in range(n_blocks):       # add ResNet blocks\n\n            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]\n\n        for i in range(n_downsampling):  # add upsampling layers\n            mult = 2 ** (n_downsampling - i)\n            if no_antialias_up:\n                model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),\n                                             kernel_size=3, stride=2,\n                                             padding=1, output_padding=1,\n                                             bias=use_bias),\n                          norm_layer(int(ngf * mult / 2)),\n                          nn.ReLU(True)]\n            else:\n                model += [Upsample(ngf * mult),\n                          nn.Conv2d(ngf * mult, int(ngf * mult / 2),\n                                    kernel_size=3, stride=1,\n                                    padding=1,  # output_padding=1,\n                                    bias=use_bias),\n                          norm_layer(int(ngf * mult / 2)),\n                          nn.ReLU(True)]\n        model += [nn.ReflectionPad2d(3)]\n        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]\n        model += [nn.Tanh()]\n\n        self.model = nn.Sequential(*model)\n\n    def forward(self, input, layers=[], encode_only=False):\n        if -1 in layers:\n            layers.append(len(self.model))\n        if len(layers) > 0:\n            feat = input\n            feats = []\n            for layer_id, layer in enumerate(self.model):\n                # print(layer_id, layer)\n                feat = layer(feat)\n                if layer_id in layers:\n                    # print(\"%d: adding the output of %s %d\" % (layer_id, layer.__class__.__name__, feat.size(1)))\n                    feats.append(feat)\n                else:\n                    # print(\"%d: skipping %s %d\" % (layer_id, layer.__class__.__name__, feat.size(1)))\n                    pass\n                if layer_id == layers[-1] and encode_only:\n                    # print('encoder only return features')\n                    return feats  # return intermediate features alone; stop in the last layers\n\n            return feat, feats  # return both output and intermediate features\n        else:\n            \"\"\"Standard forward\"\"\"\n            fake = self.model(input)\n            return fake\n\n\nclass ResnetDecoder(nn.Module):\n    \"\"\"Resnet-based decoder that consists of a few Resnet blocks + a few upsampling operations.\n    \"\"\"\n\n    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect', no_antialias=False):\n        \"\"\"Construct a Resnet-based decoder\n\n        Parameters:\n            input_nc (int)      -- the number of channels in input images\n            output_nc (int)     -- the number of channels in output images\n            ngf (int)           -- the number of filters in the last conv layer\n            norm_layer          -- normalization layer\n            use_dropout (bool)  -- if use dropout layers\n            n_blocks (int)      -- the number of ResNet blocks\n            padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero\n        \"\"\"\n        assert(n_blocks >= 0)\n        super(ResnetDecoder, self).__init__()\n        if type(norm_layer) == functools.partial:\n            use_bias = norm_layer.func == nn.InstanceNorm2d\n        else:\n            use_bias = norm_layer == nn.InstanceNorm2d\n        model = []\n        n_downsampling = 2\n        mult = 2 ** n_downsampling\n        for i in range(n_blocks):       # add ResNet blocks\n\n            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]\n\n        for i in range(n_downsampling):  # add upsampling layers\n            mult = 2 ** (n_downsampling - i)\n            if(no_antialias):\n                model += [nn.ConvTranspose2d(ngf * mult, int(ngf * mult / 2),\n                                             kernel_size=3, stride=2,\n                                             padding=1, output_padding=1,\n                                             bias=use_bias),\n                          norm_layer(int(ngf * mult / 2)),\n                          nn.ReLU(True)]\n            else:\n                model += [Upsample(ngf * mult),\n                          nn.Conv2d(ngf * mult, int(ngf * mult / 2),\n                                    kernel_size=3, stride=1,\n                                    padding=1,\n                                    bias=use_bias),\n                          norm_layer(int(ngf * mult / 2)),\n                          nn.ReLU(True)]\n        model += [nn.ReflectionPad2d(3)]\n        model += [nn.Conv2d(ngf, output_nc, kernel_size=7, padding=0)]\n        model += [nn.Tanh()]\n\n        self.model = nn.Sequential(*model)\n\n    def forward(self, input):\n        \"\"\"Standard forward\"\"\"\n        return self.model(input)\n\n\nclass ResnetEncoder(nn.Module):\n    \"\"\"Resnet-based encoder that consists of a few downsampling + several Resnet blocks\n    \"\"\"\n\n    def __init__(self, input_nc, output_nc, ngf=64, norm_layer=nn.BatchNorm2d, use_dropout=False, n_blocks=6, padding_type='reflect', no_antialias=False):\n        \"\"\"Construct a Resnet-based encoder\n\n        Parameters:\n            input_nc (int)      -- the number of channels in input images\n            output_nc (int)     -- the number of channels in output images\n            ngf (int)           -- the number of filters in the last conv layer\n            norm_layer          -- normalization layer\n            use_dropout (bool)  -- if use dropout layers\n            n_blocks (int)      -- the number of ResNet blocks\n            padding_type (str)  -- the name of padding layer in conv layers: reflect | replicate | zero\n        \"\"\"\n        assert(n_blocks >= 0)\n        super(ResnetEncoder, self).__init__()\n        if type(norm_layer) == functools.partial:\n            use_bias = norm_layer.func == nn.InstanceNorm2d\n        else:\n            use_bias = norm_layer == nn.InstanceNorm2d\n\n        model = [nn.ReflectionPad2d(3),\n                 nn.Conv2d(input_nc, ngf, kernel_size=7, padding=0, bias=use_bias),\n                 norm_layer(ngf),\n                 nn.ReLU(True)]\n\n        n_downsampling = 2\n        for i in range(n_downsampling):  # add downsampling layers\n            mult = 2 ** i\n            if(no_antialias):\n                model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=2, padding=1, bias=use_bias),\n                          norm_layer(ngf * mult * 2),\n                          nn.ReLU(True)]\n            else:\n                model += [nn.Conv2d(ngf * mult, ngf * mult * 2, kernel_size=3, stride=1, padding=1, bias=use_bias),\n                          norm_layer(ngf * mult * 2),\n                          nn.ReLU(True),\n                          Downsample(ngf * mult * 2)]\n\n        mult = 2 ** n_downsampling\n        for i in range(n_blocks):       # add ResNet blocks\n\n            model += [ResnetBlock(ngf * mult, padding_type=padding_type, norm_layer=norm_layer, use_dropout=use_dropout, use_bias=use_bias)]\n\n        self.model = nn.Sequential(*model)\n\n    def forward(self, input):\n        \"\"\"Standard forward\"\"\"\n        return self.model(input)\n\n\nclass ResnetBlock(nn.Module):\n    \"\"\"Define a Resnet block\"\"\"\n\n    def __init__(self, dim, padding_type, norm_layer, use_dropout, use_bias):\n        \"\"\"Initialize the Resnet block\n\n        A resnet block is a conv block with skip connections\n        We construct a conv block with build_conv_block function,\n        and implement skip connections in <forward> function.\n        Original Resnet paper: https://arxiv.org/pdf/1512.03385.pdf\n        \"\"\"\n        super(ResnetBlock, self).__init__()\n        self.conv_block = self.build_conv_block(dim, padding_type, norm_layer, use_dropout, use_bias)\n\n    def build_conv_block(self, dim, padding_type, norm_layer, use_dropout, use_bias):\n        \"\"\"Construct a convolutional block.\n\n        Parameters:\n            dim (int)           -- the number of channels in the conv layer.\n            padding_type (str)  -- the name of padding layer: reflect | replicate | zero\n            norm_layer          -- normalization layer\n            use_dropout (bool)  -- if use dropout layers.\n            use_bias (bool)     -- if the conv layer uses bias or not\n\n        Returns a conv block (with a conv layer, a normalization layer, and a non-linearity layer (ReLU))\n        \"\"\"\n        conv_block = []\n        p = 0\n        if padding_type == 'reflect':\n            conv_block += [nn.ReflectionPad2d(1)]\n        elif padding_type == 'replicate':\n            conv_block += [nn.ReplicationPad2d(1)]\n        elif padding_type == 'zero':\n            p = 1\n        else:\n            raise NotImplementedError('padding [%s] is not implemented' % padding_type)\n\n        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim), nn.ReLU(True)]\n        if use_dropout:\n            conv_block += [nn.Dropout(0.5)]\n\n        p = 0\n        if padding_type == 'reflect':\n            conv_block += [nn.ReflectionPad2d(1)]\n        elif padding_type == 'replicate':\n            conv_block += [nn.ReplicationPad2d(1)]\n        elif padding_type == 'zero':\n            p = 1\n        else:\n            raise NotImplementedError('padding [%s] is not implemented' % padding_type)\n        conv_block += [nn.Conv2d(dim, dim, kernel_size=3, padding=p, bias=use_bias), norm_layer(dim)]\n\n        return nn.Sequential(*conv_block)\n\n    def forward(self, x):\n        \"\"\"Forward function (with skip connections)\"\"\"\n        out = x + self.conv_block(x)  # add skip connections\n        return out\n\n\n\nclass NLayerDiscriminator(nn.Module):\n    \"\"\"Defines a PatchGAN discriminator\"\"\"\n\n    def __init__(self, input_nc, ndf=64, n_layers=3, norm_layer=nn.BatchNorm2d, no_antialias=False):\n        \"\"\"Construct a PatchGAN discriminator\n\n        Parameters:\n            input_nc (int)  -- the number of channels in input images\n            ndf (int)       -- the number of filters in the last conv layer\n            n_layers (int)  -- the number of conv layers in the discriminator\n            norm_layer      -- normalization layer\n        \"\"\"\n        super(NLayerDiscriminator, self).__init__()\n        if type(norm_layer) == functools.partial:  # no need to use bias as BatchNorm2d has affine parameters\n            use_bias = norm_layer.func == nn.InstanceNorm2d\n        else:\n            use_bias = norm_layer == nn.InstanceNorm2d\n\n        kw = 4\n        padw = 1\n        if(no_antialias):\n            sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=2, padding=padw), nn.LeakyReLU(0.2, True)]\n        else:\n            sequence = [nn.Conv2d(input_nc, ndf, kernel_size=kw, stride=1, padding=padw), nn.LeakyReLU(0.2, True), Downsample(ndf)]\n        nf_mult = 1\n        nf_mult_prev = 1\n        for n in range(1, n_layers):  # gradually increase the number of filters\n            nf_mult_prev = nf_mult\n            nf_mult = min(2 ** n, 8)\n            if(no_antialias):\n                sequence += [\n                    nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=2, padding=padw, bias=use_bias),\n                    norm_layer(ndf * nf_mult),\n                    nn.LeakyReLU(0.2, True)\n                ]\n            else:\n                sequence += [\n                    nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),\n                    norm_layer(ndf * nf_mult),\n                    nn.LeakyReLU(0.2, True),\n                    Downsample(ndf * nf_mult)]\n\n        nf_mult_prev = nf_mult\n        nf_mult = min(2 ** n_layers, 8)\n        sequence += [\n            nn.Conv2d(ndf * nf_mult_prev, ndf * nf_mult, kernel_size=kw, stride=1, padding=padw, bias=use_bias),\n            norm_layer(ndf * nf_mult),\n            nn.LeakyReLU(0.2, True)\n        ]\n\n        sequence += [nn.Conv2d(ndf * nf_mult, 1, kernel_size=kw, stride=1, padding=padw)]  # output 1 channel prediction map\n        self.model = nn.Sequential(*sequence)\n\n    def forward(self, input):\n        \"\"\"Standard forward.\"\"\"\n        return self.model(input)\n\n\n\n\nclass Identity(nn.Module):\n    def forward(self, x):\n        return x\n\ndef get_norm_layer(norm_type='instance'):\n    \"\"\"Return a normalization layer\n\n    Parameters:\n        norm_type (str) -- the name of the normalization layer: batch | instance | none\n\n    For BatchNorm, we use learnable affine parameters and track running statistics (mean/stddev).\n    For InstanceNorm, we do not use learnable affine parameters. We do not track running statistics.\n    \"\"\"\n    if norm_type == 'batch':\n        norm_layer = functools.partial(nn.BatchNorm2d, affine=True, track_running_stats=True)\n    elif norm_type == 'instance':\n        norm_layer = functools.partial(nn.InstanceNorm2d, affine=False, track_running_stats=False)\n    elif norm_type == 'none':\n        def norm_layer(x):\n            return Identity()\n    else:\n        raise NotImplementedError('normalization layer [%s] is not found' % norm_type)\n    return norm_layer\n\n\ndef define_G(input_nc, output_nc, ngf, netG, norm='batch', use_dropout=False, init_type='normal',\n             init_gain=0.02, no_antialias=False, no_antialias_up=False, gpu_ids=[], opt=None):\n    \"\"\"Create a generator\n\n    Parameters:\n        input_nc (int) -- the number of channels in input images\n        output_nc (int) -- the number of channels in output images\n        ngf (int) -- the number of filters in the last conv layer\n        netG (str) -- the architecture's name: resnet_9blocks | resnet_6blocks | unet_256 | unet_128\n        norm (str) -- the name of normalization layers used in the network: batch | instance | none\n        use_dropout (bool) -- if use dropout layers.\n        init_type (str)    -- the name of our initialization method.\n        init_gain (float)  -- scaling factor for normal, xavier and orthogonal.\n        gpu_ids (int list) -- which GPUs the network runs on: e.g., 0,1,2\n\n    Returns a generator\n\n    Our current implementation provides two types of generators:\n        U-Net: [unet_128] (for 128x128 input images) and [unet_256] (for 256x256 input images)\n        The original U-Net paper: https://arxiv.org/abs/1505.04597\n\n        Resnet-based generator: [resnet_6blocks] (with 6 Resnet blocks) and [resnet_9blocks] (with 9 Resnet blocks)\n        Resnet-based generator consists of several Resnet blocks between a few downsampling/upsampling operations.\n        We adapt Torch code from Justin Johnson's neural style transfer project (https://github.com/jcjohnson/fast-neural-style).\n\n\n    The generator has been initialized by <init_net>. It uses RELU for non-linearity.\n    \"\"\"\n    net = None\n    norm_layer = get_norm_layer(norm_type=norm)\n\n    if netG == 'resnet_9blocks':\n        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, no_antialias=no_antialias, no_antialias_up=no_antialias_up, n_blocks=9, opt=opt)\n    elif netG == 'resnet_6blocks':\n        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, no_antialias=no_antialias, no_antialias_up=no_antialias_up, n_blocks=6, opt=opt)\n    elif netG == 'resnet_4blocks':\n        net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout, no_antialias=no_antialias, no_antialias_up=no_antialias_up, n_blocks=4, opt=opt)\n    elif netG == 'unet_128':\n        net = UnetGenerator(input_nc, output_nc, 7, ngf, norm_layer=norm_layer, use_dropout=use_dropout)\n    elif netG == 'unet_256':\n        net = UnetGenerator(input_nc, output_nc, 8, ngf, norm_layer=norm_layer, use_dropout=use_dropout)\n    else:\n        raise NotImplementedError('Generator model name [%s] is not recognized' % netG)\n    return net","metadata":{"execution":{"iopub.status.busy":"2021-06-22T10:58:07.371918Z","iopub.execute_input":"2021-06-22T10:58:07.372255Z","iopub.status.idle":"2021-06-22T10:58:07.460061Z","shell.execute_reply.started":"2021-06-22T10:58:07.372219Z","shell.execute_reply":"2021-06-22T10:58:07.459088Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 4. Build Loss function","metadata":{}},{"cell_type":"code","source":"class GANLoss(nn.Module):\n    \"\"\"Define different GAN objectives.\n\n    The GANLoss class abstracts away the need to create the target label tensor\n    that has the same size as the input.\n    \"\"\"\n\n    def __init__(self, gan_mode, target_real_label=1.0, target_fake_label=0.0):\n        \"\"\" Initialize the GANLoss class.\n\n        Parameters:\n            gan_mode (str) - - the type of GAN objective. It currently supports vanilla, lsgan, and wgangp.\n            target_real_label (bool) - - label for a real image\n            target_fake_label (bool) - - label of a fake image\n\n        Note: Do not use sigmoid as the last layer of Discriminator.\n        LSGAN needs no sigmoid. vanilla GANs will handle it with BCEWithLogitsLoss.\n        \"\"\"\n        super(GANLoss, self).__init__()\n        self.register_buffer('real_label', torch.tensor(target_real_label))\n        self.register_buffer('fake_label', torch.tensor(target_fake_label))\n        self.gan_mode = gan_mode\n        if gan_mode == 'lsgan':\n            self.loss = nn.MSELoss()\n        elif gan_mode == 'vanilla':\n            self.loss = nn.BCEWithLogitsLoss()\n        elif gan_mode in ['wgangp', 'nonsaturating']:\n            self.loss = None\n        else:\n            raise NotImplementedError('gan mode %s not implemented' % gan_mode)\n\n    def get_target_tensor(self, prediction, target_is_real):\n        \"\"\"Create label tensors with the same size as the input.\n\n        Parameters:\n            prediction (tensor) - - tpyically the prediction from a discriminator\n            target_is_real (bool) - - if the ground truth label is for real images or fake images\n\n        Returns:\n            A label tensor filled with ground truth label, and with the size of the input\n        \"\"\"\n\n        if target_is_real:\n            target_tensor = self.real_label\n        else:\n            target_tensor = self.fake_label\n        return target_tensor.expand_as(prediction)\n\n    def __call__(self, prediction, target_is_real):\n        \"\"\"Calculate loss given Discriminator's output and grount truth labels.\n\n        Parameters:\n            prediction (tensor) - - tpyically the prediction output from a discriminator\n            target_is_real (bool) - - if the ground truth label is for real images or fake images\n\n        Returns:\n            the calculated loss.\n        \"\"\"\n        bs = prediction.size(0)\n        if self.gan_mode in ['lsgan', 'vanilla']:\n            target_tensor = self.get_target_tensor(prediction, target_is_real)\n            loss = self.loss(prediction, target_tensor)\n        elif self.gan_mode == 'wgangp':\n            if target_is_real:\n                loss = -prediction.mean()\n            else:\n                loss = prediction.mean()\n        elif self.gan_mode == 'nonsaturating':\n            if target_is_real:\n                loss = F.softplus(-prediction).view(bs, -1).mean(dim=1)\n            else:\n                loss = F.softplus(prediction).view(bs, -1).mean(dim=1)\n        return loss \n","metadata":{"execution":{"iopub.status.busy":"2021-06-22T10:58:07.69336Z","iopub.execute_input":"2021-06-22T10:58:07.693667Z","iopub.status.idle":"2021-06-22T10:58:07.704366Z","shell.execute_reply.started":"2021-06-22T10:58:07.693636Z","shell.execute_reply":"2021-06-22T10:58:07.703369Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 5. Build Model and training step","metadata":{}},{"cell_type":"code","source":"class CycleGANModel(pl.LightningModule):\n    def __init__(self, netG):\n        super().__init__()\n        self.lrG = 2e-4\n        self.lrD = 2e-4\n        # model\n        self.netG_A = define_G(3, 3, 64, netG)\n        self.netG_B = define_G(3, 3, 64, netG)\n        self.netD_A = NLayerDiscriminator(3)\n        self.netD_B = NLayerDiscriminator(3)\n        # loss\n        self.criterionGAN = GANLoss('lsgan').to(self.device)  # define GAN loss.\n        self.criterionCycle = torch.nn.L1Loss()\n        self.criterionIdt = torch.nn.L1Loss()\n        # loss weight\n        self.lambda_A = 10\n        self.lambda_B = 10\n        self.lambda_identity = 0.5\n\n\n    def configure_optimizers(self):\n        optim_g = torch.optim.Adam(itertools.chain(self.netG_A.parameters(), self.netG_B.parameters()), lr=self.lrG)\n        optim_d = torch.optim.Adam(itertools.chain(self.netD_A.parameters(), self.netD_B.parameters()), lr=self.lrD)\n        return optim_d, optim_g\n    \n    def loss_D(self, netD, real, fake):\n        # real\n        pred_real = netD(real)\n        loss_D_real = self.criterionGAN(pred_real, True)\n        # fake\n        pred_fake = netD(fake.detach())\n        loss_D_fake = self.criterionGAN(pred_fake, False)\n        loss_D = (loss_D_real + loss_D_fake) * 0.5\n        return loss_D\n    \n    def log_images(self, images_dict, num=4):\n        for k, images in images_dict.items():\n            image_show = torchvision.utils.make_grid((images[:num] + 1) / 2,\n                                                    nrow=4,\n                                                    normalize=False)  # to [0~1]\n            self.logger.experiment.add_image(k, image_show, self.global_step)\n\n    def training_step(self, batch, batch_idx, optimizer_idx):\n        real_A, real_B = batch\n        # training discriminator \n        if optimizer_idx == 0:\n            fake_B = self.netG_A(real_A)\n            fake_A = self.netG_B(real_B)\n\n            loss_D_A = self.loss_D(self.netD_A, real_A, fake_A)\n            loss_D_B = self.loss_D(self.netD_B, real_B, fake_B)\n            self.log('loss/D_A', loss_D_A)\n            self.log('loss/D_B', loss_D_B)\n            return loss_D_A+loss_D_B\n        # training generator\n        elif optimizer_idx == 1:\n            fake_B = self.netG_A(real_A)\n            fake_A = self.netG_B(real_B)\n            recon_B = self.netG_A(fake_A)\n            recon_A = self.netG_B(fake_B)\n            idt_B = self.netG_A(real_B)\n            idt_A = self.netG_B(real_A)\n\n            adv_loss_G_A = self.criterionGAN(self.netD_B(fake_B), True)\n            cycle_loss_A = self.criterionCycle(real_A, recon_A)\n\n            adv_loss_G_B = self.criterionGAN(self.netD_A(fake_A), True)\n            cycle_loss_B = self.criterionCycle(real_B, recon_B)\n            idt_loss_A = self.criterionIdt(idt_B, real_B)\n            idt_loss_B = self.criterionIdt(idt_A, real_A)\n\n            loss_g = adv_loss_G_A + adv_loss_G_B + cycle_loss_A*self.lambda_A + cycle_loss_B*self.lambda_B + idt_loss_A*self.lambda_A*self.lambda_identity + idt_loss_B*self.lambda_B+self.lambda_identity\n            self.log_dict({\n                'loss/G_A': adv_loss_G_A,\n                'loss/G_B': adv_loss_G_B,\n                'loss/cycle_A': cycle_loss_A,\n                'loss/cycle_B': cycle_loss_B,\n                'loss/idt_A': idt_loss_A,\n                'loss/idt_B': idt_loss_B\n            })\n            if self.global_step % 50 == 0:\n                self.log_images({\n                    'img/real_A': real_A,\n                    'img/real_B': real_B,\n                    'img/fake_A': fake_A,\n                    'img/fake_B': fake_B,\n                    'img/recon_A': recon_A,\n                    'img/recon_B': recon_B\n                })\n            return loss_g\n\n    def test_step(self, batch, batch_idx):\n        imgs, paths = batch\n        N = imgs.shape[0]\n        with torch.no_grad():\n            output_imgs = torch.clip((self.netG_A(imgs)+1)*127.5, 0, 255).permute(0, 2, 3, 1).cpu().numpy().astype(np.uint8)\n        save_dir = Path('/kaggle/images')\n        save_dir.mkdir(exist_ok=True)\n        for i in range(N):\n            save_path = save_dir / Path(paths[i]).name\n            Image.fromarray(output_imgs[i]).save(save_path)\n\n\n    def test_epoch_end(self, outputs):\n        shutil.make_archive(\"/kaggle/working/images\", 'zip', \"/kaggle/images\")","metadata":{"execution":{"iopub.status.busy":"2021-06-22T11:01:45.067015Z","iopub.execute_input":"2021-06-22T11:01:45.067366Z","iopub.status.idle":"2021-06-22T11:01:45.088581Z","shell.execute_reply.started":"2021-06-22T11:01:45.067321Z","shell.execute_reply":"2021-06-22T11:01:45.087705Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"markdown","source":"## 6. Start training","metadata":{}},{"cell_type":"code","source":"trainer = pl.Trainer(max_epochs=3, gpus=1)\nmodel = CycleGANModel(netG='resnet_4blocks')\ndata_module = Photo2MonetDataset(batch_size=1)\ntrainer.fit(model, datamodule=data_module)\ntrainer.test()","metadata":{"execution":{"iopub.status.busy":"2021-06-22T11:01:48.287114Z","iopub.execute_input":"2021-06-22T11:01:48.287438Z","iopub.status.idle":"2021-06-22T11:02:21.111881Z","shell.execute_reply.started":"2021-06-22T11:01:48.287405Z","shell.execute_reply":"2021-06-22T11:02:21.108931Z"},"trusted":true},"execution_count":null,"outputs":[]},{"cell_type":"code","source":"%load_ext tensorboard\n%tensorboard --logdir lightning_logs","metadata":{"execution":{"iopub.status.busy":"2021-06-17T02:21:11.606744Z","iopub.execute_input":"2021-06-17T02:21:11.607095Z","iopub.status.idle":"2021-06-17T02:21:15.696029Z","shell.execute_reply.started":"2021-06-17T02:21:11.607055Z","shell.execute_reply":"2021-06-17T02:21:15.695158Z"},"trusted":true},"execution_count":null,"outputs":[]}]}