{"cells":[{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# This Python 3 environment comes with many helpful analytics libraries installed\n# It is defined by the kaggle/python docker image: https://github.com/kaggle/docker-python\n# For example, here's several helpful packages to load in \n\nimport numpy as np # linear algebra\nimport pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)\n\n# Input data files are available in the \"../input/\" directory.\n# For example, running this (by clicking run or pressing Shift+Enter) will list all files under the input directory\n\nimport os\nfor dirname, _, filenames in os.walk('/kaggle/input'):\n    for filename in filenames:\n        #print(os.path.join(dirname, filename))\n        pass\n\n# Any results you write to the current directory are saved as output.","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","trusted":true},"cell_type":"code","source":"# coding: utf-8\n\n# As far as we know, I/O operations are most expensive in training Neural Networks.<br> In this kernel I wrote pytorch Dataset loader that loads and preprocess all images once and stores them into RAM.\n\n# ### Benchmark \n# As a baseline I took <a href=\"https://www.kaggle.com/speedwagon/ralsgan-dogs\">this</a> kernel \n#  - dataloader with 1 worker takes 2 min for 1 epoch\n#  - dataloader with 4 workers takes 75s for 1 epoch\n#  - RAMdataloader (this kernel) with 4 workers takes 30s for 1 epoch\n#  \n# With this dataloader you can make much more experiments and epochs!\n\n# In[1]:\n\n\nimport numpy as np\nimport pandas as pd\nimport os\nimport sys\nimport matplotlib.pyplot as plt\n\nimport torch\nfrom torch import nn, optim\nimport torch.nn.functional as F\nfrom torchvision import datasets, transforms\nfrom torchvision.utils import save_image\nfrom torch.utils.data import Dataset, DataLoader\n\nfrom PIL import Image\n\nimport xml.etree.ElementTree as ET # for parsing XML\n\nimport pprint\nimport argparse\nimport cv2\nimport copy\nimport random\n\n###\n# utility\nimport time\ndef tic():\n    start_time = time.time()\n    return start_time\n\ndef toc(start_time = 0):\n    elapsed_time = time.time() - start_time\n    return elapsed_time\n\n      \nclass EasyDict(dict):\n    def __init__(self, d=None, **kwargs):\n        if d is None:\n            d = {}\n        if kwargs:\n            d.update(**kwargs)\n        for k, v in d.items():\n            setattr(self, k, v)\n        # Class attributes\n        for k in self.__class__.__dict__.keys():\n            if not (k.startswith('__') and k.endswith('__')) and not k in ('update', 'pop'):\n                setattr(self, k, getattr(self, k))\n\n    def __setattr__(self, name, value):\n        if isinstance(value, (list, tuple)):\n            value = [self.__class__(x)\n                     if isinstance(x, dict) else x for x in value]\n        elif isinstance(value, dict) and not isinstance(value, self.__class__):\n            value = self.__class__(value)\n        super(EasyDict, self).__setattr__(name, value)\n        super(EasyDict, self).__setitem__(name, value)\n\n    __setitem__ = __setattr__\n\n    def update(self, e=None, **f):\n        d = e or dict()\n        d.update(f)\n        for k in d:\n            setattr(self, k, d[k])\n\n    def pop(self, k, d=None):\n        delattr(self, k)\n        return super(EasyDict, self).pop(k, d)\n\ndef _get_default_config():\n  c = EasyDict()\n\n  # dataset\n  c.data = EasyDict()\n  c.data.n_images = -1\n  c.data.sbox = True\n  c.data.real_label = [0.6, 0.8]\n  c.data.fake_label = [0.0, 0.2]\n\n  # model\n  c.modelG = EasyDict()\n  c.modelG.name = 'generator_nroyUS'\n  c.modelG.params = EasyDict()\n  c.modelG.params.winit = 'default'\n  c.modelG.params.nz = 128\n  c.modelG.params.gchs = [1024, 512, 256, 128, 64]\n  c.modelG.params.upsample = 'nearest'\n  c.modelG.params.kernel = 3\n  \n  c.modelD = EasyDict()\n  c.modelD.name = 'discriminator_PFN'\n  c.modelD.params = EasyDict()\n  c.modelD.params.winit = 'default'\n  c.modelD.params.dchs = [32, 64, 128, 256]\n  c.modelD.params.bns = [False, False, True, True, True]\n  c.modelD.params.acti = True\n  c.modelD.params.lrelu_slope = 0.2\n  c.modelD.params.bn = False\n  c.modelD.params.mbstd = True\n  c.modelD.params.nfilt = 32\n  c.modelD.params.kernel = 5\n\n  # train\n  c.train = EasyDict()\n  c.train.batch_size = 32\n  c.train.epoch = 2000\n  c.train.log_step = 100\n  c.train.params = EasyDict()\n  c.train.params.use_averaging_model = True\n  c.train.params.ema_alpha = 0.9999\n  c.train.params.ma_start = 100\n  c.train.params.ema = True # False in submitted version\n\n  # evaluation\n  c.eval = EasyDict()\n  c.eval.batch_size = 50\n\n  # optimizer\n  c.optimizerG = EasyDict()\n  c.optimizerG.name = 'adam'\n  c.optimizerG.params = EasyDict()\n  c.optimizerG.params.lr = 0.0005\n  c.optimizerG.params.beta1 = 0.5\n  c.optimizerG.params.beta2 = 0.999\n\n  c.optimizerD = EasyDict()\n  c.optimizerD.name = 'adam'\n  c.optimizerD.params = EasyDict()\n  c.optimizerD.params.lr = 0.0005\n  c.optimizerD.params.beta1 = 0.5\n  c.optimizerD.params.beta2 = 0.999\n\n  # loss\n  c.loss = EasyDict()\n  c.loss.name = 'sGAN'\n\n\n  # scheduler\n  c.scheduler = EasyDict()\n  c.scheduler.name = 'none'\n  c.scheduler.params = EasyDict()\n\n  return c\n\nconfig = _get_default_config()\npprint.PrettyPrinter(indent=2).pprint(config)\n\nfrom tqdm import tqdm_notebook as tqdm\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# ## Parameters of GAN\n\nif True:\n    root_dir = '../input/generative-dog-images'\nelse:\n    root_dir = '../input'\n\neval_do = False\nuse_scheduler = False\n\n# In[2]:\n\nif config.data.sbox:\n    exrw = 1\n    exrh = 1\nelse:\n    exrw = 0\n    exrh = 0\nn_images = config.data.n_images    \n\n\ng_model = config.modelG.name\nd_model = config.modelD.name\n\ng_params = config.modelG.params\nd_params = config.modelD.params\n\nnz = g_params['nz']\n\nprint(g_params)\nprint(d_params)\n\nepochs = config.train.epoch\nbatch_size = config.train.batch_size\nlrG = config.optimizerG.params.lr\nG_beta1 = config.optimizerG.params.beta1\nG_beta2 = config.optimizerD.params.beta2\n\nlrD = config.optimizerD.params.lr\nD_beta1 = config.optimizerD.params.beta1\nD_beta2 = config.optimizerD.params.beta2\n\nuse_averaging_model = config.train.params.use_averaging_model\n\n# labels\n#label_type = 1\n#if label_type==0:\n#    # default\n#    real_label = +1\n#    fake_label = -1\n#elif label_type==1:\n#    # smooth labeling\n#    real_label = +0.5\n#    fake_label = -0.5\n#else:\n#    # smooth labeling with random values\n#    real_label = [0.7, 1.2]\n#    fake_label = [-1.2, -0.7]\n\nreal_label = config.data.real_label\nfake_label = config.data.fake_label\n\n#\nnum_classes = 120\nsz = 64\n\n#\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# ## Pytorch Dataset and DataLoader\n\n# In[3]:\n\ndef get_list_objects_recursively(tdir):\n    obj_list = []\n    for root, dirnames, filenames in os.walk(tdir):\n        for filename in filenames:\n            obj_list.append(root + '/' + filename)\n    obj_list.sort()\n    return obj_list\n\n\ndef get_bboxes(anno_list, img_name):\n    basename = os.path.splitext(img_name)[0]\n    anno_fn = None\n    for fn in anno_list:\n        fn_base = os.path.basename(fn)\n        if basename==fn_base:\n            #print(img_name, fn)\n            if anno_fn is None:\n                anno_fn = fn\n            else:\n                print('filename error')\n                sys.exit()\n                \n    dog_label = anno_fn.split('/')[-2]\n    tree = ET.parse(anno_fn)\n    root = tree.getroot()\n    objects = root.findall('object')\n    bboxes = []\n    dog_labels = []\n    for o in objects:\n        bndbox = o.find('bndbox')\n        xmin = int(bndbox.find('xmin').text)\n        ymin = int(bndbox.find('ymin').text)\n        xmax = int(bndbox.find('xmax').text)\n        ymax = int(bndbox.find('ymax').text)\n        bbox = (xmin, ymin, xmax, ymax)\n        bboxes.append(bbox)\n        dog_labels.append(dog_label)\n    return bboxes, dog_labels\n    \nclass DogDataset(Dataset):\n    def __init__(self, img_dir, anno_dir, transform1=None, transform2=None, exrw=0, exrh=0, n_images=-1):\n    \n        self.img_dir = img_dir\n        self.img_names = os.listdir(img_dir)\n        self.anno_dir = anno_dir\n        self.anno_list = get_list_objects_recursively(anno_dir)\n        self.transform1 = transform1\n        self.transform2 = transform2\n        \n        self.dog_label_dict = {}\n        \n        self.imgs = []\n        self.dog_classes = []\n        \n        \n        if n_images>0:\n            self.n_images = n_images\n            seed = 190804\n            random.seed(seed)\n            random.shuffle(self.anno_list)\n        else:\n            self.n_images = len(self.anno_list)\n        \n        cnt = 0\n        for img_name in tqdm(self.img_names):\n            cnt += 1\n            if cnt > self.n_images:\n                break\n            img_org = Image.open(os.path.join(img_dir, img_name))\n            width, height = img_org.size\n            bboxes, dog_labels = get_bboxes(self.anno_list, img_name)\n            for bbox, dog_label in zip(bboxes, dog_labels):\n                xmin, ymin, xmax, ymax = bbox\n                if False:\n                    # exact square\n                    if width < height:\n                        d = (height - width)//2\n                        xmin -= d\n                        xmax += d\n                    else:\n                        d = (width - height)//2\n                        ymin -= d\n                        ymax += d\n                else:\n                    # rectangular\n                    if width < height:\n                        dmax = height - width\n                        dxl = xmin - int(max(xmin*(1-exrw), 0))\n                        dxr = int(min(xmax*(1+exrw), width)) - xmax\n                        dx = min(dxl, dxr, dmax)\n                        xmin -= dx\n                        xmax += dx\n                    else:\n                        dmax = width - height\n                        dyt = ymin - int(max(ymin*(1-exrh), 0)) \n                        dyb = int(min(ymax*(1+exrh), height)) - ymax\n                        dy = min(dyt, dyb, dmax)\n                        ymin -= dy\n                        ymax += dy\n    \n                   \n                bbox = (xmin, ymin, xmax, ymax)\n                img = img_org.crop(bbox)\n            \n                if self.transform1 is not None:\n                    img = self.transform1(img)\n                \n                if dog_label not in self.dog_label_dict:\n                    self.dog_label_dict[dog_label] = len(self.dog_label_dict)\n                dog_class = self.dog_label_dict[dog_label]\n\n                self.imgs.append(img)\n                self.dog_classes.append(dog_class)\n                #print(img_name, dog_label, dog_class)\n\n    def __getitem__(self, index):\n        img = self.imgs[index]\n        dog_class = self.dog_classes[index]\n        \n        if self.transform2 is not None:\n            img = self.transform2(img)\n        \n        return img, dog_class\n\n    def __len__(self):\n        return len(self.imgs)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# First preprocessing of data\ntransform1 = transforms.Compose([transforms.Resize(sz),\n                                transforms.CenterCrop(sz)])\n\n# Data augmentation and converting to tensors\nrandom_transforms = [transforms.RandomRotation(degrees=5)]\ntransform2 = transforms.Compose([transforms.RandomHorizontalFlip(p=0.5),\n                                 transforms.RandomApply(random_transforms, p=0.3), \n                                 transforms.ToTensor(),\n                                 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])\n\ntrain_dataset = DogDataset(img_dir=root_dir + '/all-dogs/all-dogs/',\n                           anno_dir=root_dir + '/annotation/Annotation',\n                           transform1=transform1,\n                           transform2=transform2,\n                           exrw=exrw,\n                           exrh=exrh,\n                           n_images=n_images)\n\ntrain_loader = DataLoader(dataset=train_dataset,\n                          batch_size=batch_size,\n                          shuffle=True,\n                          num_workers=4)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"x, y = next(iter(train_loader))\n\n\n# In[6]:\n\nfig = plt.figure(figsize=(25, 16))\nfor ii, img in enumerate(x):\n    if ii>31:\n        break\n    ax = fig.add_subplot(4, 8, ii + 1, xticks=[], yticks=[])\n    \n    img = img.numpy().transpose(1, 2, 0)\n    if img.shape[0]<64:\n        img = cv2.resize(img, (64, 64))\n    plt.imshow((img+1)/2)\n    plt.pause(0.1)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# taken from https://www.kaggle.com/speedwagon/ralsgan-dogs\ndef to_img(x):\n    x = 0.5 * (x + 1)\n    x = x.clamp(0, 1)\n    x = x.view(x.size(0), 3, 64, 64)\n    return x\n    \n    \n    \n### This is to show one sample image for iteration of chosing\ndef show_generated_img(epoch=0):\n    if not eval_do:\n        return\n    if epoch % 10 == 0 or epoch==epochs-1:\n        noise = fixed_noise\n        labels = fixed_labels\n        gen_tensors = netG(noise, labels).to(\"cpu\").clone().detach().squeeze(0)\n        gen_tensors = gen_tensors/2+0.5\n        gen_tensors = gen_tensors.clamp(0, 1)\n        gen_images = gen_tensors.numpy().transpose(0, 2, 3, 1)\n        plt.subplot(1,2,1)\n        plt.imshow(gen_images[0])\n        plt.subplot(1,2,2)\n        plt.imshow(gen_images[1])\n        plt.pause(0.1)\n        print('gen_image: (%f, %f)' % (gen_images.min(), gen_images.max()))\n\n# In[10]:\n\ndef generate_samples(outimg_dir, netG, n_images=10000):\n    if not os.path.exists(outimg_dir):\n        os.mkdir(outimg_dir)\n    im_batch_size = 50\n    #n_images=10000\n    for i_batch in range(0, n_images, im_batch_size):\n        gen_z = torch.randn(im_batch_size, nz, 1, 1, device=device)\n        gen_labels = torch.randint(num_classes, (im_batch_size,), device=device)\n        gen_images = (netG(gen_z, gen_labels) + 1)/2\n        images = gen_images.to(\"cpu\").clone().detach()\n        images = images.numpy().transpose(0, 2, 3, 1)\n        for i_image in range(gen_images.size(0)):\n            save_image(gen_images[i_image, :, :, :], os.path.join(outimg_dir, f'image_{i_batch+i_image:05d}.png'))\n\n\nfrom scipy.stats import truncnorm\ndef truncated_noise_sample(batch_size=1, dim_z=128, truncation=1., seed=None):\n    \"\"\" Create a truncated noise vector.\n        Params:\n            batch_size: batch size.\n            dim_z: dimension of z\n            truncation: truncation value to use\n            seed: seed for the random generator\n        Output:\n            array of shape (batch_size, dim_z)\n    \"\"\"\n    state = None if seed is None else np.random.RandomState(seed)\n    values = truncnorm.rvs(-2, 2, size=(batch_size, dim_z), random_state=state).astype(np.float32)\n    return truncation * values\n\n\ndef generate_samples_truncated(outimg_dir, netG, th=1, n_images=10000):\n    if not os.path.exists(outimg_dir):\n        os.mkdir(outimg_dir)\n    im_batch_size = 50\n    #n_images=10000\n    for i_batch in range(0, n_images, im_batch_size):\n        if th>0:\n            #print('truncated')\n            tns = truncated_noise_sample(im_batch_size, nz, th)\n            tns = np.expand_dims(tns, -1)\n            tns = np.expand_dims(tns, -1)\n            gen_z = torch.from_numpy(tns).to(device)\n        else:\n            gen_z = torch.randn(im_batch_size, nz, 1, 1, device=device)\n        gen_labels = torch.randint(num_classes, (im_batch_size,), device=device)\n        gen_images = (netG(gen_z, gen_labels) + 1)/2\n        images = gen_images.to(\"cpu\").clone().detach()\n        images = images.numpy().transpose(0, 2, 3, 1)\n        for i_image in range(gen_images.size(0)):\n            save_image(gen_images[i_image, :, :, :], os.path.join(outimg_dir, f'image_{i_batch+i_image:05d}.png'))","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# ## Initialize models and optimizers\nimport torch as th\nfrom torch.nn.utils import spectral_norm\n\ndef init_weights_normal(m):\n    w_mean = 0.0\n    w_std = 0.02\n    if type(m) == nn.ConvTranspose2d:\n        print('gaussian initialization')\n        torch.nn.init.normal_(m.weight, w_mean, w_std)\n    elif type(m) == nn.Conv2d:\n        print('gaussian initialization')\n        torch.nn.init.normal_(m.weight, w_mean, w_std)\n\n\ndef init_weights_xavier(m):\n    if type(m) == nn.Linear or type(m) == nn.Conv2d:\n        print('xavier_uniform initialization')\n        torch.nn.init.xavier_uniform_(m.weight)\n        if m.bias is not None:\n            m.bias.data.fill_(0.)        \n\n\nclass PixelwiseNorm(nn.Module):\n    def __init__(self):\n        super(PixelwiseNorm, self).__init__()\n        self.alpha = 1e-8\n\n    def forward(self, x):\n        \"\"\"\n        forward pass of the module\n        :param x: input activations volume\n        :param alpha: small number for numerical stability\n        :return: y => pixel normalized activations\n        \"\"\"\n        if False:\n            y = x.pow(2.).mean(dim=1, keepdim=True).add(self.alpha).sqrt()  # [N1HW]\n            y = x / y  # normalize the input x volume\n        else:\n            # https://github.com/deepsound-project/pggan-pytorch/blob/master/network.py\n            mean = torch.mean(x * x, 1, keepdim=True)\n            dom = torch.rsqrt(mean + self.alpha)\n            y = x * dom\n        return y\n        \n\nclass MinibatchStdDev(nn.Module):\n    \"\"\"\n    Minibatch standard deviation layer for the discriminator\n    \"\"\"\n\n    def __init__(self):\n        \"\"\"\n        derived class constructor\n        \"\"\"\n        super(MinibatchStdDev, self).__init__()\n        self.alpha = 1e-8\n\n    def forward(self, x):\n        \"\"\"\n        forward pass of the layer\n        :param x: input activation volume\n        :param alpha: small number for numerical stability\n        :return: y => x appended with standard deviation constant map\n        \"\"\"\n        batch_size, _, height, width = x.shape\n        # [B x C x H x W] Subtract mean over batch.\n        y = x - x.mean(dim=0, keepdim=True)\n        # [1 x C x H x W]  Calc standard deviation over batch\n        y = th.sqrt(y.pow(2.).mean(dim=0, keepdim=False) + self.alpha)\n\n        # [1]  Take average over feature_maps and pixels.\n        y = y.mean().view(1, 1, 1, 1)\n\n        # [B x 1 x H x W]  Replicate over group and pixels.\n        y = y.repeat(batch_size,1, height, width)\n\n        # [B x C x H x W]  Append as new feature_map.\n        y = th.cat([x, y], 1)\n        # return the computed values:\n        return y\n\n\nclass UpsamplingNN(nn.Module):\n    def __init__(self, n_input, n_output, k_size=3, stride=1, padding=1, bias=False):\n        super().__init__()\n        layers = [\n                nn.Upsample(scale_factor=2, mode='nearest'),\n                nn.Conv2d(n_input, n_output, kernel_size=k_size, stride=stride, padding=padding, bias=bias),\n            ]\n        self.convlayer = nn.Sequential(*layers)\n        self.weight = self.convlayer[1].weight\n        \n    def forward(self, x):\n        y = self.convlayer(x)\n        return y\n        \nclass Generator(nn.Module):\n    def __init__(self, nz, nfeats, nchannels, upsample='nearest', kernel=3):\n        super().__init__()\n        if upsample=='pixelshuffle_conv':\n            upsample = ConvPS\n        elif upsample=='conv_pixelshuffle':\n            upsample = ConvPS_convfirst\n        elif upsample=='nearest':\n            upsample = UpsamplingNN\n        elif upsample=='bilinear':\n            upsample = UpsamplingBL\n        else:\n            print('unknown upsample: ', upsample)\n            sys.exit()\n\n        print(self.__class__)\n        print('upsample: ', upsample)\n        print('nfeats: ', nfeats)\n        print('kernel: ', kernel)\n        pad = kernel//2\n        # input is Z, going into a convolution\n        self.conv1 = spectral_norm(nn.ConvTranspose2d(nz, nfeats * 8, 4, 1, 0, bias=False))\n        #self.bn1 = nn.BatchNorm2d(nfeats * 8)\n        # state size. (nfeats*8) x 4 x 4\n        \n        self.conv2 = spectral_norm(upsample(nfeats * 8, nfeats * 8, kernel, 1, pad, bias=False))\n        #self.bn2 = nn.BatchNorm2d(nfeats * 8)\n        # state size. (nfeats*8) x 8 x 8\n        \n        self.conv3 = spectral_norm(upsample(nfeats * 8, nfeats * 4, kernel, 1, pad, bias=False))\n        #self.bn3 = nn.BatchNorm2d(nfeats * 4)\n        # state size. (nfeats*4) x 16 x 16\n        \n        self.conv4 = spectral_norm(upsample(nfeats * 4, nfeats * 2, kernel, 1, pad, bias=False))\n        #self.bn4 = nn.BatchNorm2d(nfeats * 2)\n        # state size. (nfeats * 2) x 32 x 32\n        \n        self.conv5 = spectral_norm(upsample(nfeats * 2, nfeats, kernel, 1, pad, bias=False))\n        #self.bn5 = nn.BatchNorm2d(nfeats)\n        # state size. (nfeats) x 64 x 64\n        \n        self.conv6 = spectral_norm(nn.Conv2d(nfeats, nchannels, 5, 1, 2, bias=False))\n        # state size. (nchannels) x 64 x 64\n        self.pixnorm = PixelwiseNorm()\n    def forward(self, x, labels=None):\n        x = F.leaky_relu(self.conv1(x))\n        x = F.leaky_relu(self.conv2(x))\n        x = self.pixnorm(x)\n        x = F.leaky_relu(self.conv3(x))\n        x = self.pixnorm(x)\n        x = F.leaky_relu(self.conv4(x))\n        x = self.pixnorm(x)\n        x = F.leaky_relu(self.conv5(x))\n        x = self.pixnorm(x)\n        x = torch.tanh(self.conv6(x))\n        return x\n\ndef get_generator(nz, channels=3, nfilt=32, winit=None, upsample='nearest', kernel=3, **_):\n    model = Generator(nz, nfeats=nfilt, nchannels=channels, upsample=upsample, kernel=kernel)\n    if winit=='normal':\n        model.apply(init_weights_normal)\n    elif winit=='xavier':\n        model.apply(init_weights_xavier)\n    return model","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# https://github.com/pfnet-research/chainer-gan-lib/blob/master/common/net.py\ndef snconv(chin, chout, kernel, stride, pad, bias=False, bn=False):\n    layers = []\n    layers.append(spectral_norm(nn.Conv2d(chin, chout, kernel, stride, pad, bias=bias)))\n    if bn:\n        layers.append(nn.BatchNorm2d(chout))\n    return nn.Sequential(*layers)\n\n\nclass Discriminator(nn.Module):\n    def __init__(self, nchannels, nfeats, acti=False, lrelu_slope=0.2, bn=False, mbstd=False, kernel=3):\n        super().__init__()\n        bias = True\n        pad = kernel//2\n\n        # input is (nchannels) x 64 x 64\n        self.snconv1_1 = snconv(nchannels, nfeats//2, kernel, 1, pad, bias=bias, bn=bn)\n        self.snconv1_2 = snconv(nfeats//2, nfeats*1, 4, 2, 1, bias=bias, bn=bn)\n        self.snconv2_1 = snconv(nfeats*1,  nfeats*1, kernel, 1, pad, bias=bias, bn=bn)\n        self.snconv2_2 = snconv(nfeats*1,  nfeats*2, 4, 2, 1, bias=bias, bn=bn)\n        self.snconv3_1 = snconv(nfeats*2,  nfeats*2, kernel, 1, pad, bias=bias, bn=bn)\n        self.snconv3_2 = snconv(nfeats*2,  nfeats*4, 4, 2, 1, bias=bias, bn=bn)\n        self.snconv4_1 = snconv(nfeats*4,  nfeats*4, kernel, 1, pad, bias=bias, bn=bn)\n        self.snconv4_2 = snconv(nfeats*4,  nfeats*8, 4, 2, 1, bias=bias, bn=bn)        \n        # state size. (nfeats*8) x 4 x 4\n        self.batch_discriminator = MinibatchStdDev()\n        self.pixnorm = PixelwiseNorm()\n        # state size. (nfeats*8) x 4 x 4\n        self.snconv5 = snconv(nfeats*8+mbstd, 1, 4, 1, 0, bias=True, bn=False)\n        self.conv5 = nn.Conv2d(nfeats*8+mbstd, 1, 4, 1, 0, bias=True)\n        # state size. 1 x 1 x 1\n        \n        self.acti = acti\n        self.lrelu_slope = lrelu_slope\n        self.bn = bn\n        self.mbstd = mbstd\n        print('D: last sigmoid: ', self.acti)\n        print('   lrelu_slope: ', self.lrelu_slope)\n        print('   batch_norm: ', self.bn)\n        print('   kernel_size: ', kernel)\n        \n    def forward(self, x, labels=None):\n        x = F.leaky_relu(self.snconv1_1(x), self.lrelu_slope)\n        x = F.leaky_relu(self.snconv1_2(x), self.lrelu_slope)\n        x = F.leaky_relu(self.snconv2_1(x), self.lrelu_slope)\n        x = F.leaky_relu(self.snconv2_2(x), self.lrelu_slope)\n        x = F.leaky_relu(self.snconv3_1(x), self.lrelu_slope)\n        x = F.leaky_relu(self.snconv3_2(x), self.lrelu_slope)\n        x = F.leaky_relu(self.snconv4_1(x), self.lrelu_slope)\n        x = F.leaky_relu(self.snconv4_2(x), self.lrelu_slope)\n        if self.mbstd:\n            x = self.batch_discriminator(x)\n        #x = torch.sigmoid(self.conv5(x))\n        x= self.snconv5(x)\n        if self.acti:\n            x = torch.sigmoid(x)\n        return x.view(-1, 1)\n\n\n\ndef get_discriminator(channels=3, nfilt=32, winit=None, acti=True, lrelu_slope=0.2, bn=False, mbstd=False, kernel=3, **_):\n    model = Discriminator(nchannels=channels, nfeats=nfilt, acti=acti, lrelu_slope=lrelu_slope, bn=bn, mbstd=mbstd, kernel=kernel)\n    if winit=='normal':\n        model.apply(init_weights_normal)\n    elif winit=='xavier':\n        model.apply(init_weights_xavier)\n    return model\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"netG = get_generator(**g_params).to(device)\nnetD = get_discriminator(**d_params).to(device)\n\noptimizerD = optim.Adam(netD.parameters(), lr=lrD, betas=(D_beta1, D_beta2))\noptimizerG = optim.Adam(netG.parameters(), lr=lrG, betas=(G_beta1, G_beta2))\n\nlr_schedulerD = torch.optim.lr_scheduler.StepLR(optimizerD, step_size=100, gamma=0.5)\nlr_schedulerG = torch.optim.lr_scheduler.StepLR(optimizerG, step_size=100, gamma=0.5)\n\nfixed_noise = torch.randn(25, nz, 1, 1, device=device)\n#fixed_labels = torch.randint(num_classes, (25,), device=device)\nfixed_labels = torch.randint(1, (25,), device=device)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def ema_model(target_model, source_model, alpha):\n    target_params = target_model.state_dict()    \n    source_params = source_model.state_dict()\n    for param_name, param in source_params.items():\n        #print(param_name)\n        #target_params[param_name] *= (1-alpha)\n        #target_params[param_name] += alpha * param\n        target_params[param_name] *= alpha\n        target_params[param_name] += (1-alpha) * param\n\n\ndef ma_model(target_model, source_model, n):\n    target_params = target_model.state_dict()    \n    source_params = source_model.state_dict()\n    for param_name, param in source_params.items():\n        #print(param_name)\n        target_params[param_name] *= (n / (n+1))\n        target_params[param_name] += (1/(n+1)) * param\n        \ndef copy_model_params(target_model, source_model):\n    target_model.load_state_dict(source_model.state_dict())\n\nif use_averaging_model:\n    ma_count = 0\n    netG_ma = copy.deepcopy(netG)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# loss function\n# 1=sGAN\ndef get_lossD_sGAN():\n    criterion = nn.BCELoss()\n    def lossD_real(outputR, real_labels):\n        return criterion(outputR, real_labels)\n    def lossD_fake(outputF, fake_labels):        \n        return criterion(outputF, fake_labels)\n    return lossD_real, lossD_fake\n\ndef get_lossG_sGAN():\n    criterion = nn.BCELoss()\n    def lossG(outputR, outputF, real_labels, fake_labels):\n        return criterion(outputF, real_labels)\n    return lossG\n    \n\n# 7=RaLSGAN\ndef get_lossD_RaLSGAN():\n    def lossD_rerative(outputR, outputF, real_labels, fake_labels):\n        errD = (torch.mean((outputR - torch.mean(outputF) - real_labels) ** 2) +\n                torch.mean((outputF - torch.mean(outputR) - fake_labels) ** 2))/2\n        return errD\n    return lossD_rerative\n\ndef get_lossG_RaLSGAN():\n    def lossG_rerative(outputR, outputF, real_labels, fake_labels):\n        errG = (torch.mean((outputR - torch.mean(outputF) + real_labels) ** 2) +\n                torch.mean((outputF - torch.mean(outputR) + fake_labels) ** 2))/2\n        return errG\n    return lossG_rerative\n\n\n\nloss_type = config.loss.name\nif loss_type=='sGAN':\n    print('loss type: sGAN')\n    criterionD_real, criterionD_fake = get_lossD_sGAN()\n    criterionG = get_lossG_sGAN()\n    relative_loss = False\nelse:\n    print('loss type: RaLSGAN')\n    criterionD = get_lossD_RaLSGAN()\n    criterionG = get_lossG_RaLSGAN()\n    relative_loss = True","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"#epochs = 150","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# ## Training loop\noutput_dir = '../tmp'\noutimg_dir = '../tmp/output_images'\nif not os.path.exists(output_dir):\n    os.mkdir(output_dir)\nif not os.path.exists(outimg_dir):\n    os.mkdir(outimg_dir)\n\nstart_time = tic()    \n    \nlossD = []\nlossG = []\nfor epoch in range(epochs):\n    \n    elasped_time = toc(start_time)\n    print('elasped_time[h]: ', elasped_time/3600)\n    if elasped_time/3600 > 8.5:\n        break\n    \n    lossD_epoch = 0\n    lossG_epoch = 0\n    for ii, (real_images, bleed_labels) in tqdm(enumerate(train_loader), total=len(train_loader)):\n        ###########################\n        # input\n        ###########################\n        real_images = real_images.to(device)\n        bleed_labels = bleed_labels.to(device)\n        bleed_labels[:] = 0\n        #print(bleed_labels.shape)\n        batch_size = real_images.size(0)\n        #labels = torch.full((batch_size, 1), real_label, device=device)\n\n        ###########################\n        # labels\n        ###########################\n        if isinstance(real_label, list):\n            real_labels = torch.empty((batch_size, 1), device=device).uniform_(real_label[0], real_label[1])\n            fake_labels = torch.empty((batch_size, 1), device=device).uniform_(fake_label[0], fake_label[1])\n            #for p in np.random.choice(batch_size, size=np.random.randint((batch_size//8)), replace=False):\n            #    real_labels[p], fake_labels[p] = fake_labels[p].clone(), real_labels[p].clone() # swap labels\n        else:\n            # simple\n            real_labels = torch.full((batch_size, 1), real_label, device=device)\n            fake_labels = torch.full((batch_size, 1), fake_label, device=device)\n    \n    \n        ###########################\n        # (1) Update D network\n        ###########################\n        netD.zero_grad()\n        outputR = netD(real_images, bleed_labels)\n        \n        if relative_loss:\n            noise = torch.randn(batch_size, nz, 1, 1, device=device)\n            fake = netG(noise, bleed_labels)\n            outputF = netD(fake.detach(), bleed_labels)\n            errD = criterionD(outputR, outputF, real_labels, fake_labels)\n            errD.backward(retain_graph=True)\n        else:\n            errD_real = criterionD_real(outputR, real_labels)\n            errD_real.backward()\n            \n            noise = torch.randn(batch_size, nz, 1, 1, device=device)\n            fake = netG(noise, bleed_labels)\n            outputF = netD(fake.detach(), bleed_labels)\n            errD_fake = criterionD_fake(outputF, fake_labels)\n            errD_fake.backward()\n            errD = errD_real + errD_fake\n            \n        optimizerD.step()\n        ###########################\n        # (2) Update G network\n        ###########################\n        netG.zero_grad()\n        outputF = netD(fake, bleed_labels)\n        errG = criterionG(outputR, outputF, real_labels, fake_labels)\n                \n        errG.backward()\n        optimizerG.step()\n        \n        # log\n        lossD_epoch += errD.item()\n        lossG_epoch += errG.item()\n        \n        if (ii+1) % (len(train_loader)//2) == 0:\n            print('[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f'\n                  % (epoch + 1, epochs, ii+1, len(train_loader),\n                     errD.item(), errG.item()))\n                     \n        # averaging\n        if use_averaging_model:\n            if epoch >= config.train.params.ma_start:\n                if ma_count==0:\n                    #ma_model(netG_ma, netG_ema, ma_count)\n                    copy_model_params(netG_ma, netG)\n                else:\n                    if config.train.params.ema:\n                        ema_model(netG_ma, netG, config.train.params.ema_alpha)\n                    else:\n                        ma_model(netG_ma, netG, ma_count)\n                ma_count += 1\n                \n                     \n\n    # epoch end\n    if use_scheduler:\n        lr_schedulerD.step()\n        lr_schedulerG.step()\n\n    lossD.append(lossD_epoch)\n    lossG.append(lossG_epoch)\n    show_generated_img(epoch)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# for debug","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# plot\nplt.figure()\nplt.plot(lossD, label='lossD')\nplt.plot(lossG, label='lossG')\nplt.xlabel('epoch')\nplt.ylabel('loss')\nplt.legend()\n#plt.show()\nplt.pause(0.01)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# ## Let's visualise generated results\n\n# In[11]:\n\nnetG.eval()\ngen_z = torch.randn(32, nz, 1, 1, device=device)\ngen_labels = torch.randint(num_classes, (32,), device=device)\ngen_images = netG(gen_z, gen_labels).to(\"cpu\").clone().detach()\ngen_images = gen_images.numpy().transpose(0, 2, 3, 1)\n\n\nfig = plt.figure(figsize=(25, 16))\nfor ii, img in enumerate(gen_images):\n    if ii>31:\n        break\n    ax = fig.add_subplot(4, 8, ii + 1, xticks=[], yticks=[])\n    if img.shape[0]<64:\n        img = cv2.resize(img, (64, 64))\n    plt.imshow((img+1)/2)\nplt.pause(0.1)\n","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# ## Make predictions and submit\n# final_evaluation\n\nth = 1.0\nif use_averaging_model:\n    # ma\n    generate_samples_truncated(outimg_dir, netG_ma)\nelse:\n    generate_samples_truncated(outimg_dir, netG)\n    \n\nimport shutil\nshutil.make_archive('images', 'zip', outimg_dir)\n\nprint('last epoch: ', epoch)\nelapsed_time = toc(start_time)\nprint('elapsed time[h] :', elapsed_time/3600)","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":1}