{"cells":[{"metadata":{},"cell_type":"markdown","source":"# CUT v2"},{"metadata":{},"cell_type":"markdown","source":"Copy of Otto's notebook"},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"%config Completer.use_jedi = False # Enables code auto-completion\n\nfrom kaggle_datasets import KaggleDatasets\nimport os\nimport matplotlib.pyplot as plt\nimport math\nimport random\nimport cv2","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.nn import init\nfrom packaging import version\n\nimport time\n\nfrom collections import OrderedDict\nimport functools\nfrom torch.optim import lr_scheduler","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"import numpy as np","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Set the paths for both the monet paintings and the photo"},{"metadata":{"trusted":true},"cell_type":"code","source":"BASE_PATH = \"../input/gan-getting-started/\"\nMONET_PATH = os.path.join(BASE_PATH, \"monet_jpg\")\nPHOTO_PATH = os.path.join(BASE_PATH, \"photo_jpg\")","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"After loading the paintings and the photo's, we want to quickly get an idea what they look like. The method below shows n (random) images with a fixed size for a given path.\n\n*Note: cv2 is a library for solving computer vision problems (reading/ showing images)*"},{"metadata":{"trusted":true},"cell_type":"code","source":"def batch_visualization(path, n_images, is_random=True, figsize=(16, 16)):\n    \"\"\"\n    Taken from: \n      https://www.kaggle.com/ihelon/monet-visualization-and-augmentation\n    \"\"\"\n    \n    plt.figure(figsize=figsize)\n    \n    w = int(n_images ** .5)\n    h = math.ceil(n_images / w)\n    \n    all_names = os.listdir(path)\n    \n    image_names = all_names[:n_images]\n    if is_random:\n        image_names = random.sample(all_names, n_images)\n    \n    # the color is set below\n    for ind, image_name in enumerate(image_names):\n        img = cv2.imread(os.path.join(path, image_name))\n        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) \n        plt.subplot(h, w, ind + 1)\n        plt.imshow(img)\n        plt.axis(\"off\")\n    \n    plt.show()","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"4 Random Monet paintings:"},{"metadata":{"trusted":true},"cell_type":"code","source":"batch_visualization(MONET_PATH, 4, is_random=True, figsize=(5, 5))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"4 Random photo's:"},{"metadata":{"trusted":true},"cell_type":"code","source":"batch_visualization(PHOTO_PATH, 4, is_random=True, figsize=(5, 5))","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"# CUT"},{"metadata":{},"cell_type":"markdown","source":"### Definitions\nAffine: A type of transformation of a geometric object, which does preserve straight and paralel lines, but can alter distances between them. \n"},{"metadata":{},"cell_type":"markdown","source":"*Note: the 'nn' referred to comes from torch.nn, and stands for Neural Network. The module is used for training neural networks*"},{"metadata":{"trusted":true},"cell_type":"code","source":"def get_norm_layer(norm_type='instance'):\n    \"\"\"Return a normalization layer\n    Parameters:\n        norm_type (str) -- the name of the normalization layer: batch | instance | none\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","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"The class below is used to 'downsample' aka reformat the image to smaller configurations\n\nStride: the 'jump' needed to go from one element to the other within different dimensions (within the same or other tuple for example)"},{"metadata":{"trusted":true},"cell_type":"code","source":"class 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])","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"Conv explanation: https://machinelearningmastery.com/convolutional-layers-for-deep-learning-neural-networks/\n\nConvolutional Layer within a (CNN), a convolution is a linear operation between an input and a filter/kernel layer. These are multiplied and the result is the scalar product.\n"},{"metadata":{"trusted":true},"cell_type":"code","source":"class 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        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)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"conv block wit skip connection explanation: https://theaisummer.com/skip-connections/\n\nPadding: The adding of values to an image to increase/ preserve the size of an image without losing/ creating information. This can be done by only adding 0's. Another type is by reflecting the values from an original image as done below: \n![reflection_padding](https://www.machinecurve.com/wp-content/uploads/2020/02/reflection_pad.jpg).\nFinally, there is also replication padding, which uses symmetry: \n![replication_padding](https://www.machinecurve.com/wp-content/uploads/2020/02/replication_pad.png)\n\nidentity function: a function wich is able to produce an output which is similar to the input. I preserves the idenity of the input.\n\nloss function: distance between two tensors (multidimensional matrix), when optimizing a problem (from input to prediction), we want to minimize this distance between the two.\n\nbackpropagation: understanding how changing the weights of a network affects the loss function, by calculating partial derivatives with respect to the model parameters.\n\nGradient descent: Is simply a chain-rule function, which tries to find the lowest value for the loss function. This lowest value is then where the derivative is equal to 0, and both sides of this position descent. Since this loss function is a function of functions, we have to calculate the partial derivatives. These themselves have a gradient, but these functions are then again functions of functions, and therefore you continously go back through the network. Note, that then the gradients of the partial derivatives keep getting smaller the more you go back (butterfly effect). \n\nSkip connections: these 'skip' a or multiple layers and feed their output more direct to layers further in the network instead of simply the next layer. The idenity function is used to add skip blocks together and then the identity function is used to preserve the gradient ","attachments":{}},{"metadata":{"trusted":true},"cell_type":"code","source":"class 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        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        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        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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def init_weights(net, init_type='normal', init_gain=0.02, debug=False):\n    \"\"\"Initialize network weights.\n    Parameters:\n        net (network)   -- network to be initialized\n        init_type (str) -- the name of an initialization method: normal | xavier | kaiming | orthogonal\n        init_gain (float)    -- scaling factor for normal, xavier and orthogonal.\n    We use 'normal' in the original pix2pix and CycleGAN paper. But xavier and kaiming might\n    work better for some applications. Feel free to try yourself.\n    \"\"\"\n    def init_func(m):  # define the initialization function\n        classname = m.__class__.__name__\n        if hasattr(m, 'weight') and (classname.find('Conv') != -1 or classname.find('Linear') != -1):\n            if debug:\n                print(classname)\n            if init_type == 'normal':\n                init.normal_(m.weight.data, 0.0, init_gain)\n            elif init_type == 'xavier':\n                init.xavier_normal_(m.weight.data, gain=init_gain)\n            elif init_type == 'kaiming':\n                init.kaiming_normal_(m.weight.data, a=0, mode='fan_in')\n            elif init_type == 'orthogonal':\n                init.orthogonal_(m.weight.data, gain=init_gain)\n            else:\n                raise NotImplementedError('initialization method [%s] is not implemented' % init_type)\n            if hasattr(m, 'bias') and m.bias is not None:\n                init.constant_(m.bias.data, 0.0)\n        elif classname.find('BatchNorm2d') != -1:  # BatchNorm Layer's weight is not a matrix; only normal distribution applies.\n            init.normal_(m.weight.data, 1.0, init_gain)\n            init.constant_(m.bias.data, 0.0)\n\n    net.apply(init_func)  # apply the initialization function <init_func>\n        \ndef init_net(net, init_type='normal', init_gain=0.02, gpu_ids=[], debug=False, initialize_weights=True):\n    \"\"\"Initialize a network: 1. register CPU/GPU device (with multi-GPU support); 2. initialize the network weights\n    Parameters:\n        net (network)      -- the network to be initialized\n        init_type (str)    -- the name of an initialization method: normal | xavier | kaiming | orthogonal\n        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    Return an initialized network.\n    \"\"\"\n    if len(gpu_ids) > 0:\n        assert(torch.cuda.is_available())\n        net.to(gpu_ids[0])\n        # if not amp:\n        # net = torch.nn.DataParallel(net, gpu_ids)  # multi-GPUs for non-AMP training\n    if initialize_weights:\n        init_weights(net, init_type, init_gain=init_gain, debug=debug)\n    return net\n\nclass ResnetGenerator(nn.Module):\n    \"\"\"Resnet-based generator that consists of Resnet blocks between a few downsampling/upsampling operations.\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        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(2), # ORIGINALLY 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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Normalize(nn.Module):\n\n    def __init__(self, power=2):\n        super(Normalize, self).__init__()\n        self.power = power\n\n    def forward(self, x):\n        norm = x.pow(self.power).sum(1, keepdim=True).pow(1. / self.power)\n        out = x.div(norm + 1e-7)\n        return out\n\nclass PatchSampleF(nn.Module):\n    def __init__(self, use_mlp=False, init_type='normal', init_gain=0.02, nc=256, gpu_ids=[]):\n        # potential issues: currently, we use the same patch_ids for multiple images in the batch\n        super(PatchSampleF, self).__init__()\n        self.l2norm = Normalize(2)\n        self.use_mlp = use_mlp\n        self.nc = nc  # hard-coded\n        self.mlp_init = False\n        self.init_type = init_type\n        self.init_gain = init_gain\n        self.gpu_ids = gpu_ids\n\n    def create_mlp(self, feats):\n        for mlp_id, feat in enumerate(feats):\n            input_nc = feat.shape[1]\n            mlp = nn.Sequential(*[nn.Linear(input_nc, self.nc), nn.ReLU(), nn.Linear(self.nc, self.nc)])\n            if len(self.gpu_ids) > 0:\n                mlp.cuda()\n            setattr(self, 'mlp_%d' % mlp_id, mlp)\n        init_net(self, self.init_type, self.init_gain, self.gpu_ids)\n        self.mlp_init = True\n\n    def forward(self, feats, num_patches=64, patch_ids=None):\n        return_ids = []\n        return_feats = []\n        if self.use_mlp and not self.mlp_init:\n            self.create_mlp(feats)\n        for feat_id, feat in enumerate(feats):\n            B, H, W = feat.shape[0], feat.shape[2], feat.shape[3]\n            feat_reshape = feat.permute(0, 2, 3, 1).flatten(1, 2)\n            if num_patches > 0:\n                if patch_ids is not None:\n                    patch_id = patch_ids[feat_id]\n                else:\n                    patch_id = torch.randperm(feat_reshape.shape[1], device=feats[0].device)\n                    patch_id = patch_id[:int(min(num_patches, patch_id.shape[0]))]  # .to(patch_ids.device)\n                x_sample = feat_reshape[:, patch_id, :].flatten(0, 1)  # reshape(-1, x.shape[1])\n            else:\n                x_sample = feat_reshape\n                patch_id = []\n            if self.use_mlp:\n                mlp = getattr(self, 'mlp_%d' % feat_id)\n                x_sample = mlp(x_sample)\n            return_ids.append(patch_id)\n            x_sample = self.l2norm(x_sample)\n\n            if num_patches == 0:\n                x_sample = x_sample.permute(0, 2, 1).reshape([B, x_sample.shape[-1], H, W])\n            return_feats.append(x_sample)\n        return return_feats, return_ids","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class GANLoss(nn.Module):\n    \"\"\"Define different GAN objectives.\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        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        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        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        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        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        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","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class PatchNCELoss(nn.Module):\n    def __init__(self, opt):\n        super().__init__()\n        self.opt = opt\n        self.cross_entropy_loss = torch.nn.CrossEntropyLoss(reduction='none')\n        self.mask_dtype = torch.uint8 if version.parse(torch.__version__) < version.parse('1.2.0') else torch.bool\n\n    def forward(self, feat_q, feat_k):\n        batchSize = feat_q.shape[0]\n        dim = feat_q.shape[1]\n        feat_k = feat_k.detach()\n\n        # pos logit\n        l_pos = torch.bmm(feat_q.view(batchSize, 1, -1), feat_k.view(batchSize, -1, 1))\n        l_pos = l_pos.view(batchSize, 1)\n\n        # neg logit\n\n        # Should the negatives from the other samples of a minibatch be utilized?\n        # In CUT and FastCUT, we found that it's best to only include negatives\n        # from the same image. Therefore, we set\n        # --nce_includes_all_negatives_from_minibatch as False\n        # However, for single-image translation, the minibatch consists of\n        # crops from the \"same\" high-resolution image.\n        # Therefore, we will include the negatives from the entire minibatch.\n        #if self.opt.nce_includes_all_negatives_from_minibatch: #***\n            # reshape features as if they are all negatives of minibatch of size 1.\n        #    batch_dim_for_bmm = 1\n        #else: #***\n        batch_dim_for_bmm = model.batch_size #*** was self.opt.batch_size #***\n\n        # reshape features to batch size\n        feat_q = feat_q.view(batch_dim_for_bmm, -1, dim)\n        feat_k = feat_k.view(batch_dim_for_bmm, -1, dim)\n        npatches = feat_q.size(1)\n        l_neg_curbatch = torch.bmm(feat_q, feat_k.transpose(2, 1))\n\n        # diagonal entries are similarity between same features, and hence meaningless.\n        # just fill the diagonal with very small number, which is exp(-10) and almost zero\n        diagonal = torch.eye(npatches, device=feat_q.device, dtype=self.mask_dtype)[None, :, :]\n        l_neg_curbatch.masked_fill_(diagonal, -10.0)\n        l_neg = l_neg_curbatch.view(-1, npatches)\n\n        out = torch.cat((l_pos, l_neg), dim=1) / model.nce_T # was self.opt.nce_T ***\n\n        loss = self.cross_entropy_loss(out, torch.zeros(out.size(0), dtype=torch.long,\n                                                        device=feat_q.device))\n\n        return loss","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def 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    norm_layer = get_norm_layer(norm_type=norm)\n    net = ResnetGenerator(input_nc, output_nc, ngf, norm_layer=norm_layer, use_dropout=use_dropout,\n                          no_antialias=no_antialias, no_antialias_up=no_antialias_up, n_blocks=9, opt=opt)\n    return init_net(net, init_type, init_gain, gpu_ids, initialize_weights=('stylegan2' not in netG))\n\ndef define_F(input_nc, netF, norm='batch', use_dropout=False, init_type='normal', init_gain=0.02, no_antialias=False, gpu_ids=[], netF_nc=256):\n    net = PatchSampleF(use_mlp=True, init_type=init_type, init_gain=init_gain, gpu_ids=gpu_ids, nc=netF_nc)\n    return init_net(net, init_type, init_gain, gpu_ids)\n\ndef define_D(input_nc, ndf, netD, n_layers_D=3, norm='batch', init_type='normal', init_gain=0.02, no_antialias=False, gpu_ids=[], opt=None):\n    norm_layer = get_norm_layer(norm_type=norm)\n    net = NLayerDiscriminator(input_nc, ndf, n_layers=3, norm_layer=norm_layer, no_antialias=no_antialias,)\n    return init_net(net, init_type, init_gain, gpu_ids, initialize_weights=('stylegan2' not in netD))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class CUT():\n    \n    def __init__(self):\n        \n        # https://github.com/taesungp/contrastive-unpaired-translation/blob/master/options/train_options.py\n        self.n_epochs = 200 # The number of eopchs with the inital learning rate\n        self.n_epochs_decay = 200 # The number of epochs to linearly decay learning rate to zero\n        self.beta1 = 0.5 # Momentum term of adam\n        self.beta2 = 0.999 # Momentum term of adam\n        self.lr = 0.0002 # Initial learning rate of adam\n        self.gan_mode = 'lsgan' # The type of GAN objective\n        self.pool_size = 50 # The size of image buffer that stores previously generated images\n        self.lr_policy = 'linear' # Learning rate policy\n        self.lr_decay_iters = 50 # Multiply by gamma every lr_decay_iters iterations\n        self.isTrain = True # Train or test\n        self.epoch_count = 1 # The starting epoch count\n        \n        \n        # https://github.com/taesungp/contrastive-unpaired-translation/blob/master/options/base_options.py\n        self.gpu_ids = [0] # Determines which GPU to use\n        self.input_nc = 3 # Number of input image channels: 3 for RGB and 1 for grayscale\n        self.output_nc = 3 # Number of output image channels: 3 for RGB and 1 for grayscale\n        self.ngf = 64 # Number of gen filters in the last convolutional layer\n        self.ndf = 64 # Number of discrim filters in the first convolutional layer\n        self.opt_netD = 'basic' # Specify discriminator architecture; the basic model is a 70x70 PatchGAN.\n        self.opt_netG = 'resnet_9blocks' # Specify the generator architecture\n        self.n_layers_D = 3 # Only used if netD=='n_layers'\n        self.normG = 'instance' # Specify the type of normalization for G\n        self.normD = 'instance' # Specify the type of normalization for D\n        self.init_type = 'xavier' # Network initialization\n        self.init_gain = 0.02 # Scaling factor for normal, xavier and orthogonal initialization\n        self.no_dropout = True # No dropout for the generator\n        self.direction = 'AtoB'\n        # ...\n        self.batch_size = 1 # Input batch size\n        #self.load_size = 256 # Scale the images to this size\n        #self.crop_size = 256 # Then crop to this size\n        \n        # https://github.com/taesungp/contrastive-unpaired-translation/blob/87ab89cdca651f87742844016b0cfa49fa7bd3ee/models/base_model.py#L8\n        self.device = torch.device('cuda:{}'.format(self.gpu_ids[0])) if self.gpu_ids else torch.device('cpu')  # get device name: CPU or GPU\n        self.optimizers = []\n        \n        # https://github.com/taesungp/contrastive-unpaired-translation/blob/87ab89cdca651f87742844016b0cfa49fa7bd3ee/models/cut_model.py#L9\n        self.lambda_GAN = 1.0 # Weight for GAN loss: GAN(G(X))\n        self.lambda_NCE = 1.0 # Weight for NCE loss: NCE(G(X), X)\n        self.nce_idt = True # Use NCE loss for identity mapping: NCE(G(Y), Y); True for CUT, False for SinCUT\n        self.nce_layers = [0, 4, 8, 12, 16]  # Compute NCE loss on these layers\n        #self.nce_includes_all_negatives_from_minibatch = False # (used for single image translation)\n        self.opt_netF = 'mlp_sample' # How to downsample the feature map\n        self.netF_nc = 256 # 256 # ***\n        self.nce_T = 0.07 # Temperature for NCE loss\n        self.num_patches = 256 # 256 # *** # Number of patches per layer\n        self.flip_equivariance = False # Used by FastCUT, but not CUT\n        \n        self.loss_names = ['G_GAN', 'D_real', 'D_fake', 'G', 'NCE', 'NCE_Y']\n        self.visual_names = ['real_A', 'fake_B', 'real_B', 'idt_B']\n        \n        if self.isTrain:\n            self.model_names = ['G', 'F', 'D']\n        else:  # during test time, only load G\n            self.model_names = ['G']\n            \n        # Not sure about Park et al's settings for:\n        self.no_antialias = True # TRUE OR FALSE?\n        self.no_antialias_up = True # TRUE OR FALSE?\n        opt = None\n        \n        # define networks (both generator and discriminator)\n        self.netG = define_G(self.input_nc, self.output_nc, self.ngf, self.opt_netG, self.normG,\n                                      not self.no_dropout, self.init_type, self.init_gain,\n                                      self.no_antialias, self.no_antialias_up, self.gpu_ids, opt)\n        self.netF = define_F(self.input_nc, self.opt_netF, self.normG, not self.no_dropout,\n                                      self.init_type, self.init_gain, self.no_antialias, self.gpu_ids, self.netF_nc)\n        \n        if self.isTrain:\n            self.netD = define_D(self.output_nc, self.ndf, self.opt_netD, self.n_layers_D, self.normD,\n                                          self.init_type, self.init_gain, self.no_antialias, self.gpu_ids, opt)\n\n            # define loss functions\n            self.criterionGAN = GANLoss(self.gan_mode).to(self.device)\n            self.criterionNCE = []\n\n            for nce_layer in self.nce_layers:\n                self.criterionNCE.append(PatchNCELoss(opt).to(self.device))\n\n            self.criterionIdt = torch.nn.L1Loss().to(self.device)\n            self.optimizer_G = torch.optim.Adam(self.netG.parameters(), lr=self.lr, betas=(self.beta1, self.beta2))\n            self.optimizer_D = torch.optim.Adam(self.netD.parameters(), lr=self.lr, betas=(self.beta1, self.beta2))\n            self.optimizers.append(self.optimizer_G)\n            self.optimizers.append(self.optimizer_D)\n            \n            self.schedulers = [self.get_scheduler(optimizer, opt) for optimizer in self.optimizers]\n    \n    \n    def get_current_losses(self):\n        \"\"\"Return training losses / errors. train.py will print out these errors on console, and save them to a file\"\"\"\n        errors_ret = OrderedDict()\n        for name in self.loss_names:\n            if isinstance(name, str):\n                errors_ret[name] = float(getattr(self, 'loss_' + name))  # float(...) works for both scalar tensor and float number\n        return errors_ret\n    \n    def update_learning_rate(self):\n        \"\"\"Update learning rates for all the networks; called at the end of every epoch\"\"\"\n        for scheduler in self.schedulers:\n            if self.lr_policy == 'plateau':\n                scheduler.step(self.metric)\n            else:\n                scheduler.step()\n\n        lr = self.optimizers[0].param_groups[0]['lr']\n        print('learning rate = %.7f' % lr)\n\n    def data_dependent_initialize(self, data):\n        \"\"\"\n        The feature network netF is defined in terms of the shape of the intermediate, extracted\n        features of the encoder portion of netG. Because of this, the weights of netF are\n        initialized at the first feedforward pass with some input images.\n        Please also see PatchSampleF.create_mlp(), which is called at the first forward() call.\n        \"\"\"\n        \n        #self.set_input(data)  # basically: set model.real_A to images from one domain, and model.real_B to images from the other domain\n        self.real_A = real_A\n        self.real_B = real_B\n        \n        bs_per_gpu = self.real_A.size(0) // max(len(self.gpu_ids), 1)\n        self.real_A = self.real_A[:bs_per_gpu]\n        self.real_B = self.real_B[:bs_per_gpu]\n        self.forward()                     # compute fake images: G(A)\n        if self.isTrain:\n            self.compute_D_loss().backward()                  # calculate gradients for D\n            self.compute_G_loss().backward()                   # calculate graidents for G\n            if self.lambda_NCE > 0.0:\n                self.optimizer_F = torch.optim.Adam(self.netF.parameters(), lr=self.lr, betas=(self.beta1, self.beta2))\n                self.optimizers.append(self.optimizer_F)\n                \n    def optimize_parameters(self):\n        # forward\n        self.forward()\n\n        # update D\n        self.set_requires_grad(self.netD, True)\n        self.optimizer_D.zero_grad()\n        self.loss_D = self.compute_D_loss()\n        self.loss_D.backward() # They have a custom backward function? Is this necessary? I think the standard pytorch function is used here ***\n        self.optimizer_D.step()\n\n        # update G\n        self.set_requires_grad(self.netD, False)\n        self.optimizer_G.zero_grad()\n        if self.netF == 'mlp_sample':\n            self.optimizer_F.zero_grad()\n        self.loss_G = self.compute_G_loss()\n        self.loss_G.backward()\n        self.optimizer_G.step()\n        if self.netF == 'mlp_sample':\n            self.optimizer_F.step()\n    \n    def set_requires_grad(self, nets, requires_grad=False):\n        \"\"\"Set requies_grad=Fasle for all the networks to avoid unnecessary computations\n        Parameters:\n            nets (network list)   -- a list of networks\n            requires_grad (bool)  -- whether the networks require gradients or not\n        \"\"\"\n        if not isinstance(nets, list):\n            nets = [nets]\n        for net in nets:\n            if net is not None:\n                for param in net.parameters():\n                    param.requires_grad = requires_grad\n    \n    def get_scheduler(self, optimizer, opt):\n        \"\"\"Return a learning rate scheduler\n        Parameters:\n            optimizer          -- the optimizer of the network\n            opt (option class) -- stores all the experiment flags; needs to be a subclass of BaseOptions．　\n                                  opt.lr_policy is the name of learning rate policy: linear | step | plateau | cosine\n        For 'linear', we keep the same learning rate for the first <opt.n_epochs> epochs\n        and linearly decay the rate to zero over the next <opt.n_epochs_decay> epochs.\n        For other schedulers (step, plateau, and cosine), we use the default PyTorch schedulers.\n        See https://pytorch.org/docs/stable/optim.html for more details.\n        \"\"\"\n        if self.lr_policy == 'linear':\n            def lambda_rule(epoch):\n                lr_l = 1.0 - max(0, epoch + self.epoch_count - self.n_epochs) / float(self.n_epochs_decay + 1)\n                return lr_l\n            scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda_rule)\n        elif self.lr_policy == 'step':\n            scheduler = lr_scheduler.StepLR(optimizer, step_size=self.lr_decay_iters, gamma=0.1)\n        elif self.lr_policy == 'plateau':\n            scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)\n        elif self.lr_policy == 'cosine':\n            scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=self.n_epochs, eta_min=0)\n        else:\n            return NotImplementedError('learning rate policy [%s] is not implemented', self.lr_policy)\n        return scheduler\n\n    def set_input(self, input):\n        \"\"\"Unpack input data from the dataloader and perform necessary pre-processing steps.\n        Parameters:\n            input (dict): include the data itself and its metadata information.\n        The option 'direction' can be used to swap domain A and domain B.\n        \"\"\"\n        AtoB = self.direction == 'AtoB'\n        self.real_A = input['A' if AtoB else 'B'].to(self.device)\n        self.real_B = input['B' if AtoB else 'A'].to(self.device)\n        self.image_paths = input['A_paths' if AtoB else 'B_paths']\n\n    def parallelize(self):\n        for name in self.model_names:\n            if isinstance(name, str):\n                net = getattr(self, 'net' + name)\n                setattr(self, 'net' + name, torch.nn.DataParallel(net, self.gpu_ids))\n    \n\n    def forward(self):\n        \"\"\"Run forward pass; called by both functions <optimize_parameters> and <test>.\"\"\"\n        self.real = torch.cat((self.real_A, self.real_B), dim=0) if self.nce_idt and self.isTrain else self.real_A\n        if self.flip_equivariance:\n            self.flipped_for_equivariance = self.isTrain and (np.random.random() < 0.5)\n            if self.flipped_for_equivariance:\n                self.real = torch.flip(self.real, [3])\n\n        self.fake = self.netG(self.real)\n        self.fake_B = self.fake[:self.real_A.size(0)]\n        \n        #print(self.fake.shape)\n        #print(self.fake_B.shape)\n        \n        aaa = self.fake.permute(0, 2, 3, 1)\n        bbb = self.fake_B.permute(0, 2, 3, 1)\n        \n        print()\n        plt.imshow(aaa[0].detach().cpu().squeeze(0).numpy(), interpolation='none')\n        plt.show()\n        plt.imshow(bbb[0].detach().cpu().squeeze(0).numpy(), interpolation='none')\n        plt.show()\n        \n        if self.nce_idt:\n            self.idt_B = self.fake[self.real_A.size(0):]\n    \n    def normal_CUT_D_loss(self):\n        \"\"\"Calculate GAN loss for the discriminator\"\"\"\n        fake = self.fake_B.detach()\n        # Fake; stop backprop to the generator by detaching fake_B\n        pred_fake = self.netD(fake)\n        self.loss_D_fake = self.criterionGAN(pred_fake, False).mean()\n        # Real\n        self.pred_real = self.netD(self.real_B)\n        loss_D_real = self.criterionGAN(self.pred_real, True)\n        self.loss_D_real = loss_D_real.mean()\n\n        # combine loss and calculate gradients\n        self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5\n    \n    def normal_CUT_G_loss(self):\n        \"\"\"Calculate GAN and NCE loss for the generator\"\"\"\n        fake = self.fake_B\n        # First, G(A) should fake the discriminator\n        if self.lambda_GAN > 0.0:\n            pred_fake = self.netD(fake)\n            self.loss_G_GAN = self.criterionGAN(pred_fake, True).mean() * self.lambda_GAN\n        else:\n            self.loss_G_GAN = 0.0\n\n        if self.lambda_NCE > 0.0:\n            self.loss_NCE = self.calculate_NCE_loss(self.real_A, self.fake_B)\n        else:\n            self.loss_NCE, self.loss_NCE_bd = 0.0, 0.0\n\n        if self.nce_idt and self.lambda_NCE > 0.0:\n            self.loss_NCE_Y = self.calculate_NCE_loss(self.real_B, self.idt_B)\n            loss_NCE_both = (self.loss_NCE + self.loss_NCE_Y) * 0.5\n        else:\n            loss_NCE_both = self.loss_NCE\n\n        self.loss_G = self.loss_G_GAN + loss_NCE_both\n        return self.loss_G\n\n    def compute_D_loss(self):\n        \"\"\"Calculate GAN loss for the discriminator\"\"\"\n        fake = self.fake_B.detach()\n        # Fake; stop backprop to the generator by detaching fake_B\n        pred_fake = self.netD(fake)\n        self.loss_D_fake = self.criterionGAN(pred_fake, False).mean()\n        # Real\n        self.pred_real = self.netD(self.real_B)\n        loss_D_real = self.criterionGAN(self.pred_real, True)\n        self.loss_D_real = loss_D_real.mean()\n\n        # combine loss and calculate gradients\n        self.loss_D = (self.loss_D_fake + self.loss_D_real) * 0.5\n        return self.loss_D\n\n    def compute_G_loss(self):\n        \"\"\"Calculate GAN and NCE loss for the generator\"\"\"\n        fake = self.fake_B\n        # First, G(A) should fake the discriminator\n        if self.lambda_GAN > 0.0:\n            pred_fake = self.netD(fake)\n            self.loss_G_GAN = self.criterionGAN(pred_fake, True).mean() * self.lambda_GAN\n        else:\n            self.loss_G_GAN = 0.0\n\n        if self.lambda_NCE > 0.0:\n            self.loss_NCE = self.calculate_NCE_loss(self.real_A, self.fake_B)\n        else:\n            self.loss_NCE, self.loss_NCE_bd = 0.0, 0.0\n\n        if self.nce_idt and self.lambda_NCE > 0.0:\n            self.loss_NCE_Y = self.calculate_NCE_loss(self.real_B, self.idt_B)\n            loss_NCE_both = (self.loss_NCE + self.loss_NCE_Y) * 0.5\n        else:\n            loss_NCE_both = self.loss_NCE\n\n        self.loss_G = self.loss_G_GAN + loss_NCE_both\n        return self.loss_G\n    \n    def calculate_NCE_loss(self, src, tgt):\n        n_layers = len(self.nce_layers)\n        feat_q = self.netG(tgt, self.nce_layers, encode_only=True)\n\n        if self.flip_equivariance and self.flipped_for_equivariance:\n            feat_q = [torch.flip(fq, [3]) for fq in feat_q]\n\n        feat_k = self.netG(src, self.nce_layers, encode_only=True)\n        feat_k_pool, sample_ids = self.netF(feat_k, self.num_patches, None)\n        feat_q_pool, _ = self.netF(feat_q, self.num_patches, sample_ids)\n\n        total_nce_loss = 0.0\n        for f_q, f_k, crit, nce_layer in zip(feat_q_pool, feat_k_pool, self.criterionNCE, self.nce_layers):\n            loss = crit(f_q, f_k) * self.lambda_NCE\n            total_nce_loss += loss.mean()\n\n        return total_nce_loss / n_layers","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def print_current_losses(epoch, iters, losses, t_comp, t_data):\n    \"\"\"print current losses on console; also save the losses to the disk\n    Parameters:\n        epoch (int) -- current epoch\n        iters (int) -- current training iteration during this epoch (reset to 0 at the end of every epoch)\n        losses (OrderedDict) -- training losses stored in the format of (name, float) pairs\n        t_comp (float) -- computational time per data point (normalized by batch_size)\n        t_data (float) -- data loading time per data point (normalized by batch_size)\n    \"\"\"\n    message = '(epoch: %d, iters: %d, time: %.3f, data: %.3f) ' % (epoch, iters, t_comp, t_data)\n    for k, v in losses.items():\n        message += '%s: %.3f ' % (k, v)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"print_freq = 10 # Frequency of showing training results on console\n\nmodel = CUT()\n\ntotal_iters = 0                # the total number of training iterations\noptimize_time = 0.1","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#all_names = os.listdir(PHOTO_PATH)\n#img_name_0 = all_names[0]\n\n#img = cv2.imread(os.path.join(PHOTO_PATH, img_name_0))\n#img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)\n\n#img_tensor = torch.tensor(img, dtype=torch.float64, device=model.device)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#print(img_tensor[None, :, :].shape)\n#real_A = torch.empty([300, 256, 256, 3])\n#print(real_A.shape)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"n_images = 10 # I\n\nreal_A = torch.empty([n_images, 256, 256, 3], dtype=torch.float32, device=model.device)\nreal_B = torch.empty([n_images, 256, 256, 3], dtype=torch.float32, device=model.device)\n\nalternate_image_size = 256\nreal_A = torch.empty([n_images, alternate_image_size, alternate_image_size, 3], dtype=torch.float32, device=model.device) # ***\nreal_B = torch.empty([n_images, alternate_image_size, alternate_image_size, 3], dtype=torch.float32, device=model.device) # ***\n\nall_photo_names = os.listdir(PHOTO_PATH)\nall_monet_names = os.listdir(MONET_PATH)\nfor i in range(n_images):\n    imgA = cv2.imread(os.path.join(PHOTO_PATH, all_photo_names[i]))\n    imgA = cv2.cvtColor(imgA, cv2.COLOR_BGR2RGB)\n    imgA = imgA[:alternate_image_size, :alternate_image_size, :] # ***\n    imgA = np.ascontiguousarray(imgA, dtype=np.float32) / 255\n    real_A[i] = torch.from_numpy(imgA).to(model.device)\n    \n    imgB = cv2.imread(os.path.join(MONET_PATH, all_monet_names[i]))\n    imgB = cv2.cvtColor(imgB, cv2.COLOR_BGR2RGB)\n    imgB = imgB[:alternate_image_size, :alternate_image_size, :] # ***\n    imgB = np.ascontiguousarray(imgB, dtype=np.float32) / 255\n    real_B[i] = torch.from_numpy(imgB).to(model.device)\n    \n    #\"\"\"\n    print()\n    print(i)\n    plt.imshow(real_A[i].detach().cpu().squeeze(0).numpy(), interpolation='none')\n    plt.show()\n    plt.imshow(real_B[i].detach().cpu().squeeze(0).numpy(), interpolation='none')\n    plt.show()\n    #\"\"\"\n\nreal_A = real_A.permute(0, 3, 1, 2)\nreal_B = real_B.permute(0, 3, 1, 2)\nprint(real_A.shape)\nprint(real_B.shape)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"dataset = {'A':real_A, 'B':real_B}\ndataset_size = len(dataset['A'])   # get the number of images in the dataset.\nprint('The number of training images = %d' % dataset_size)\n\ntimes = []\nfor epoch in range(model.epoch_count, model.n_epochs + model.n_epochs_decay + 1):    # outer loop for different epochs\n    epoch_start_time = time.time()  # timer for entire epoch\n    iter_data_time = time.time()    # timer for data loading per iteration\n    epoch_iter = 0                  # the number of training iterations in current epoch, reset to 0 every epoch\n\n    #dataset.set_epoch(epoch)\n    for i, data in enumerate(dataset):  # inner loop within one epoch\n        \n        #print(f\"\\nIteration {i} starting.\")\n        \n        iter_start_time = time.time()  # timer for computation per iteration\n        if total_iters % print_freq == 0:\n            t_data = iter_start_time - iter_data_time\n\n        # batch_size = data[\"A\"].size(0)\n        batch_size = model.batch_size ###################\n        \n        total_iters += batch_size\n        epoch_iter += batch_size\n        if len(model.gpu_ids) > 0:\n            torch.cuda.synchronize()\n        optimize_start_time = time.time()\n        \n        if epoch == model.epoch_count and i == 0:\n            model.data_dependent_initialize(data)\n            #model.setup(opt)               # regular setup: load and print networks; create schedulers\n            model.parallelize()\n        \n        #model.set_input(data)  # basically: set model.real_A to images from one domain, and model.real_B to images from the other domain\n        model.real_A = real_A\n        model.real_B = real_B\n        \n        model.optimize_parameters()   # calculate loss functions, get gradients, update network weights\n        if len(model.gpu_ids) > 0:\n            torch.cuda.synchronize()\n        optimize_time = (time.time() - optimize_start_time) / batch_size * 0.005 + 0.995 * optimize_time\n\n        if total_iters % print_freq == 0:    # print training losses and save logging information to the disk\n            losses = model.get_current_losses()\n            print_current_losses(epoch, epoch_iter, losses, optimize_time, t_data)\n\n        iter_data_time = time.time()\n        \n        #print(f\"Iteration {i} ending.\")\n\n    print('End of epoch %d / %d \\t Time Taken: %d sec' % (epoch, model.n_epochs + model.n_epochs_decay, time.time() - epoch_start_time))\n    model.update_learning_rate() # update learning rates at the end of every epoch.","execution_count":null,"outputs":[]}],"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":4,"nbformat_minor":4}