{"cells":[{"metadata":{"trusted":true},"cell_type":"code","source":"import shutil\nimport os\nimport random\nfrom time import time\nimport matplotlib.pyplot as plt\nimport numpy as np\nimport torch\nfrom torch.nn.utils import spectral_norm\nfrom torch.utils.data import Dataset\nimport torchvision\nimport xml.etree.ElementTree as ET\nfrom tqdm import tqdm_notebook as tqdm\nfrom scipy.stats import truncnorm\n","execution_count":null,"outputs":[]},{"metadata":{"_uuid":"8f2839f25d086af736a60e9eeb907d3b93b6e0e5","_cell_guid":"b1076dfc-b9ad-4769-8c92-a6c4dae69d19","trusted":true},"cell_type":"code","source":"# Define RGB images\nCHANNEL = {'RGB': 3, 'BW': 1}\nIMG_CHANNEL = CHANNEL['RGB']\n\n# Defining path to dog images \nDATA_LOCATION = '../input/all-dogs/all-dogs/'\n\n# Define batch size\nBATCH_SIZE = 32\n\n# Each image in the generator starts from LATENT_DIM points generated from a gaussian distribution (variance = 1 sigma = 0)\nLATENT_DIM = 128\n\n# The amount of parameters in the networks scales with those number\nCONV_DEPTHS_G = 32\nCONV_DEPTHS_D = 48\n\n\n# Learning rate functions\n# The learning scheduale \nLEARNING_RATE_G = 0.0003 # was 0.0003 \nLEARNING_RATE_D = 0.0001 # was 0.0001\nBETA_1 = 0.5\nT0_interval = 200\n# T0_interval = 400\n\n\n# Amount of epochs, EPOCHS/T0_interval has to be integer\n\nEPOCHS = 600\n# EPOCHS = 400 \n\n\n# If archtiecture changed to CGAN then it might be needed later\nNUM_CLASSES = 10\n\n# Define limit for training time\nTIME_FOR_TRAIN = 25000 # 8.3 hours (in seconds)\n# TIME_FOR_TRAIN = 7 * 60 * 60 # hours\n# TIME_FOR_TRAIN = 5 * 60","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"def seed_everything(seed=42):\n    random.seed(seed)\n    os.environ['PYTHONHASHSEED'] = str(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.backends.cudnn.deterministic = True\n    \nseed_everything()\ndevice = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n\n\n","execution_count":null,"outputs":[]},{"metadata":{"_cell_guid":"79c7e3d0-c299-4dcb-8224-4455121ee9b0","_uuid":"d629ff2d2480ee46fbb7e2d37f6b5fab8052498a","trusted":true},"cell_type":"code","source":"class DataGenerator(Dataset):\n    def __init__(self, directory, transform=None, n_samples=np.inf):\n        self.directory = directory\n        self.transform = transform\n        self.n_samples = n_samples\n        self.samples = self._load_subfolders_images(directory)\n        if len(self.samples) == 0:\n            raise RuntimeError(\"Found 0 files in subfolders of: {}\".format(directory))\n\n    def _load_subfolders_images(self, root):\n        IMG_EXTENSIONS = (\n            '.jpg', '.jpeg', '.png', '.ppm', '.bmp', '.pgm', '.tif', '.tiff', '.webp')\n\n        def is_valid_file(x):\n            return torchvision.datasets.folder.has_file_allowed_extension(x, IMG_EXTENSIONS)\n\n        \n        \n        required_transforms = torchvision.transforms.Compose([\n            torchvision.transforms.Resize(64),\n            torchvision.transforms.CenterCrop(64),\n        ])\n        imgs = []\n        paths = []\n        for root, _, fnames in sorted(os.walk(root)):\n            for fname in sorted(fnames)[:min(self.n_samples, 999999999999999)]:\n                path = os.path.join(root, fname)\n                paths.append(path)\n\n        for path in paths:\n            if is_valid_file(path):\n                # Load image\n                img = torchvision.datasets.folder.default_loader(path)\n\n                # Get bounding boxes\n                annotation_basename = os.path.splitext(os.path.basename(path))[0]\n                annotation_dirname = next(\n                    dirname for dirname in os.listdir('../input/annotation/Annotation/') if\n                    dirname.startswith(annotation_basename.split('_')[0]))\n                annotation_filename = os.path.join('../input/annotation/Annotation/',\n                                                   annotation_dirname, annotation_basename)\n                tree = ET.parse(annotation_filename)\n                root = tree.getroot()\n                objects = root.findall('object')\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\n                    w = np.min((xmax - xmin, ymax - ymin))\n                    bbox = (xmin, ymin, xmin + w, ymin + w)\n                    object_img = required_transforms(img.crop(bbox))\n                    # object_img = object_img.resize((64,64), Image.ANTIALIAS)\n                    imgs.append(object_img)\n        return imgs\n\n    def __getitem__(self, index):\n        sample = self.samples[index]\n\n        if self.transform is not None:\n            sample = self.transform(sample)\n\n        return np.asarray(sample)\n\n    def __len__(self):\n        return len(self.samples)\n    \n\ntransform = torchvision.transforms.Compose([torchvision.transforms.RandomHorizontalFlip(p=0.1),\n                                            torchvision.transforms.ToTensor(),\n                                            torchvision.transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])    \ndata_gen = DataGenerator(DATA_LOCATION, transform=transform, n_samples=25000)\ntrain_loader = torch.utils.data.DataLoader(data_gen, shuffle=True, batch_size=BATCH_SIZE, num_workers=4)","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 5.1 Generator"},{"metadata":{"trusted":true},"cell_type":"code","source":"class PixelwiseNorm(torch.nn.Module):\n    def __init__(self):\n        super(PixelwiseNorm, self).__init__()\n\n    def forward(self, x, alpha=1e-8):\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        y = x.pow(2.).mean(dim=1, keepdim=True).add(alpha).sqrt()  # [N1HW]\n        y = x / y  # normalize the input x volume\n        return y\n    \n    \nclass Generator(torch.nn.Module):\n    def __init__(self, latent_dim, nfeats, nchannels):  #was nz\n        super(Generator, self).__init__()\n\n        # input is Z, going into a convolution\n        self.conv1 = spectral_norm(torch.nn.ConvTranspose2d(latent_dim, nfeats * 8, 4, 1, 0, bias=False)) #was nz\n\n        self.conv2 = spectral_norm(torch.nn.ConvTranspose2d(nfeats * 8, nfeats * 8, 4, 2, 1, bias=False))\n\n        self.conv3 = spectral_norm(torch.nn.ConvTranspose2d(nfeats * 8, nfeats * 4, 4, 2, 1, bias=False))\n\n        self.conv4 = spectral_norm(torch.nn.ConvTranspose2d(nfeats * 4, nfeats * 2, 4, 2, 1, bias=False))\n\n        self.conv5 = spectral_norm(torch.nn.ConvTranspose2d(nfeats * 2, nfeats, 4, 2, 1, bias=False))\n\n        self.conv6 = spectral_norm(torch.nn.ConvTranspose2d(nfeats, nchannels, 3, 1, 1, bias=False))\n        self.pixnorm = PixelwiseNorm()\n\n    def forward(self, x):\n\n        x = torch.nn.functional.leaky_relu(self.conv1(x))\n        x = self.pixnorm(x) # wasnt here\n        x = torch.nn.functional.leaky_relu(self.conv2(x))\n        x = torch.nn.Dropout(0.05)(x)\n        x = self.pixnorm(x)\n        x = torch.nn.functional.leaky_relu(self.conv3(x))\n        x = self.pixnorm(x)        \n        x = torch.nn.functional.leaky_relu(self.conv4(x))\n        x = torch.nn.Dropout(0.05)(x)\n        x = self.pixnorm(x)\n        x = torch.nn.functional.leaky_relu(self.conv5(x))\n        x = self.pixnorm(x)\n        x = torch.tanh(self.conv6(x))\n\n        return x\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 5.2 Discriminator"},{"metadata":{"trusted":true},"cell_type":"code","source":"class MinibatchStdDev(torch.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\n    def forward(self, x, alpha=1e-8):\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 = torch.sqrt(y.pow(2.).mean(dim=0, keepdim=False) + 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 = torch.cat([x, y], 1)\n        # return the computed values:\n        return y","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"class Discriminator(torch.nn.Module):\n    def __init__(self, nchannels, nfeats):\n        super(Discriminator, self).__init__()\n\n        # input is (nchannels) x 64 x 64\n        self.conv1 = torch.nn.Conv2d(nchannels, nfeats, 4, 2, 1, bias=False)\n        # state size. (nfeats) x 32 x 32\n\n        self.conv2 = spectral_norm(torch.nn.Conv2d(nfeats, nfeats * 2, 4, 2, 1, bias=False))\n        self.bn2 = torch.nn.BatchNorm2d(nfeats * 2)\n        # state size. (nfeats*2) x 16 x 16\n\n        self.conv3 = spectral_norm(torch.nn.Conv2d(nfeats * 2, nfeats * 4, 4, 2, 1, bias=False))\n        self.bn3 = torch.nn.BatchNorm2d(nfeats * 4)\n        # state size. (nfeats*4) x 8 x 8\n\n        self.conv4 = spectral_norm(torch.nn.Conv2d(nfeats * 4, nfeats * 8, 4, 2, 1, bias=False))\n        self.bn4 = torch.nn.MaxPool2d(2)\n        # state size. (nfeats*8) x 4 x 4\n        self.batch_discriminator = MinibatchStdDev()\n\n        self.conv5 = spectral_norm(torch.nn.Conv2d(nfeats * 8 + 1, 1, 2, 1, 0, bias=False))\n        # state size. 1 x 1 x 1\n\n    def forward(self, x):\n        x = torch.nn.functional.leaky_relu(self.conv1(x), 0.2)\n        x = torch.nn.functional.leaky_relu(self.bn2(self.conv2(x)), 0.2)\n        x = torch.nn.functional.leaky_relu(self.bn3(self.conv3(x)), 0.2)\n        x = torch.nn.functional.leaky_relu(self.bn4(self.conv4(x)), 0.2)\n        x = self.batch_discriminator(x)\n        x = torch.sigmoid(self.conv5(x))\n        # x= self.conv5(x)\n        return x.view(-1, 1)\n","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 6. Functions "},{"metadata":{},"cell_type":"markdown","source":"### 6.1 Image show functions"},{"metadata":{"trusted":true},"cell_type":"code","source":"def show_generated_img_all():\n    gen_z = torch.randn(32, LATENT_DIM, 1, 1, device=device)\n    gen_images = netG(gen_z).to(\"cpu\").clone().detach()\n    gen_images = gen_images.numpy().transpose(0, 2, 3, 1)\n    gen_images = (gen_images + 1.0) / 2.0\n    fig = plt.figure(figsize=(25, 16))\n    for ii, img in enumerate(gen_images):\n        ax = fig.add_subplot(4, 8, ii + 1, xticks=[], yticks=[])\n        plt.imshow(img)\n    # plt.savefig(filename)\n\n\ndef show_generated_img():\n    row_num = 1\n    col_num = 10 \n    gen_z = torch.randn(row_num * col_num , LATENT_DIM, 1, 1, device=device)\n    gen_images = netG(gen_z).to(\"cpu\").clone().detach()\n    gen_images = gen_images.numpy().transpose(0, 2, 3, 1)\n    gen_images = (gen_images + 1.0) / 2.0\n    fig = plt.figure(figsize=(20, 4))\n    for ii, img in enumerate(gen_images):\n        ax = fig.add_subplot(row_num, col_num, ii + 1, xticks=[], yticks=[])\n        plt.imshow(img)\n    plt.show()\n    # plt.savefig(filename)\n    \n    ","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"### 6.2 Truncate function"},{"metadata":{"trusted":true},"cell_type":"code","source":"def truncated_normal(size, threshold=1):\n    values = truncnorm.rvs(-threshold, threshold, size=size)\n    return values","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 7. Train Module"},{"metadata":{"trusted":true},"cell_type":"code","source":"def train_module(epochs):\n    step = 0\n    start = time()\n    for epoch in range(epochs):\n        for ii, (real_images) in enumerate(train_loader):\n            end = time()\n            if (end - start) > TIME_FOR_TRAIN:\n                break\n            ############################\n            # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))\n            ###########################\n            # train with real\n            netD.zero_grad()\n            real_images = real_images.to(device)\n            batch_size = real_images.size(0)\n            labels = torch.full((batch_size, 1), real_label, device=device) + np.random.uniform(-0.1, 0.1)\n\n            output = netD(real_images)\n            errD_real = criterion(output, labels)\n            errD_real.backward()\n            D_x = output.mean().item()\n\n            # train with fake\n            noise = torch.randn(batch_size, LATENT_DIM, 1, 1, device=device)\n            fake = netG(noise)\n            labels.fill_(fake_label) + np.random.uniform(0, 0.2)\n            output = netD(fake.detach())\n            errD_fake = criterion(output, labels)\n            errD_fake.backward()\n            D_G_z1 = output.mean().item()\n            errD = errD_real + errD_fake\n            optimizerD.step()\n\n            ############################\n            # (2) Update G network: maximize log(D(G(z)))\n            ###########################\n            netG.zero_grad()\n            labels.fill_(real_label)  # fake labels are real for generator cost\n            output = netD(fake)\n            errG = criterion(output, labels)\n            errG.backward()\n            D_G_z2 = output.mean().item()\n            optimizerG.step()\n\n            if step % 500 == 0:\n                print('[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f'\n                      % (epoch + 1, EPOCHS, ii, len(train_loader),\n                         errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))\n\n                valid_image = netG(fixed_noise)\n            step += 1\n            lr_schedulerG.step(epoch)\n            lr_schedulerD.step(epoch)\n            \n        if epoch % 5 == 0:\n            show_generated_img()\n            print(end - start)\n        if (end - start) > TIME_FOR_TRAIN:\n            break","execution_count":null,"outputs":[]},{"metadata":{},"cell_type":"markdown","source":"## 8. Define and train the GAN"},{"metadata":{"trusted":true},"cell_type":"code","source":"netG = Generator(LATENT_DIM, CONV_DEPTHS_G, IMG_CHANNEL).to(device)\nnetD = Discriminator(3, CONV_DEPTHS_D).to(device)\ncriterion = torch.nn.BCELoss()\noptimizerD = torch.optim.Adam(netD.parameters(), lr=LEARNING_RATE_G, betas=(BETA_1, 0.999))\noptimizerG = torch.optim.Adam(netG.parameters(), lr=LEARNING_RATE_D, betas=(BETA_1, 0.999))\nlr_schedulerG = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizerG,\n                                                                     T_0=EPOCHS // T0_interval, eta_min=0.00005)\nlr_schedulerD = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizerD,\n                                                                     T_0=EPOCHS // T0_interval, eta_min=0.00005)\n\nfixed_noise = torch.randn(25, LATENT_DIM, 1, 1, device=device)\n\nreal_label = 0.8\nfake_label = 0.15\nbatch_size = train_loader.batch_size\n# criterion = nn.MSELoss()\ntrain_module(EPOCHS)","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"netG.eval()\n\nif not os.path.exists('../output_images'):\n    os.mkdir('../output_images')\nim_batch_size = 50\nn_images = 10000\nfor i_batch in range(0, n_images, im_batch_size):\n    z = truncated_normal((im_batch_size, LATENT_DIM, 1, 1), threshold=1)\n    gen_z = torch.from_numpy(z).float().to(device)\n    # gen_z = torch.randn(im_batch_size, 100, 1, 1, device=device)\n    gen_images = netG(gen_z)\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        torchvision.utils.save_image((gen_images[i_image, :, :, :] + 1.0) / 2.0,\n                                     os.path.join('../output_images', f'image_{i_batch + i_image:05d}.png'))\n\nshutil.make_archive('images', 'zip', '../output_images')","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"# show_generated_img_all()","execution_count":null,"outputs":[]},{"metadata":{"trusted":true},"cell_type":"code","source":"","execution_count":null,"outputs":[]}],"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.6.4","mimetype":"text/x-python","codemirror_mode":{"name":"ipython","version":3},"pygments_lexer":"ipython3","nbconvert_exporter":"python","file_extension":".py"}},"nbformat":4,"nbformat_minor":1}